main.py

import os
import io
import json
import re
import tempfile
import asyncio
from typing import Optional
import logging
from contextlib import asynccontextmanager
from fastapi import FastAPI, Request, status, Depends, Header, HTTPException
from fastapi.concurrency import run_in_threadpool
from pydantic import BaseModel
from dotenv import load_dotenv
from openai import OpenAI
from elevenlabs.client import ElevenLabs
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_postgres.vectorstores import PGVector
from sqlalchemy import create_engine

# --- GRADIO ---
import gradio as gr

# --- SETUP ---
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
logging.getLogger('tensorflow').setLevel(logging.ERROR)
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

load_dotenv()
NEON_DATABASE_URL = os.getenv("NEON_DATABASE_URL")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
ELEVENLABS_API_KEY = os.getenv("ELEVENLABS_API_KEY")
SHARED_SECRET = os.getenv("SHARED_SECRET")

# --- CONFIG ---
COLLECTION_NAME = "real_estate_embeddings"
EMBEDDING_MODEL = "hkunlp/instructor-large"
ELEVENLABS_VOICE_NAME = "Leo"
PLANNER_MODEL = "gpt-4o-mini"
ANSWERER_MODEL = "gpt-4o"
TABLE_DESCRIPTIONS = """
- "ongoing_projects_source": Details about projects currently under construction.
- "upcoming_projects_source": Information on future planned projects.
- "completed_projects_source": Facts about projects that are already finished.
- "historical_sales_source": Specific sales records, including price, date, and property ID.
- "past_customers_source": Information about previous customers.
- "feedback_source": Customer feedback and ratings for projects.
"""

# --- CLIENTS ---
embeddings = None
vector_store = None
client_openai = OpenAI(api_key=OPENAI_API_KEY)
client_elevenlabs = ElevenLabs(api_key=ELEVENLABS_API_KEY)


# --- LIFESPAN ---
@asynccontextmanager
async def lifespan(app: FastAPI):
    global embeddings, vector_store
    logging.info(f"Loading embedding model: {EMBEDDING_MODEL}")
    embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)

    logging.info(f"Connecting to vector store: {COLLECTION_NAME}")
    engine = create_engine(NEON_DATABASE_URL, pool_pre_ping=True)
    vector_store = PGVector(
        connection=engine,
        collection_name=COLLECTION_NAME,
        embeddings=embeddings,
    )
    logging.info("Vector store ready.")
    yield
    logging.info("Shutting down.")


app = FastAPI(lifespan=lifespan)


# --- PROMPTS ---
QUERY_FORMULATION_PROMPT = """
You are a query analysis agent. Transform the user's query into a precise search query and determine the correct table to filter by.
**Available Tables:**
{table_descriptions}
**User's Query:** "{user_query}"
**Task:**
1. Rephrase into a clear, keyword-focused English search query.
2. If status keywords (ongoing, completed, upcoming, etc.) are present, pick the matching table.
3. If no status keyword, set filter_table to null.
4. Return JSON: {{"search_query": "...", "filter_table": "table_name or null"}}
""".format(table_descriptions=TABLE_DESCRIPTIONS)

ANSWER_SYSTEM_PROMPT = """
You are an expert AI assistant for a premier real estate developer.
## YOUR PERSONA
- You are professional, helpful, and highly knowledgeable. Your tone should be polite and articulate.
## CORE BUSINESS KNOWLEDGE
- **Operational Cities:** We are currently operational in Pune, Mumbai, Bengaluru, Delhi, Chennai, Hyderabad, Goa, Gurgaon, Kolkata.
- **Property Types:** We offer luxury apartments, villas, and commercial properties.
- **Budget Range:** Our residential properties typically range from 45 lakhs to 5 crores.
## CORE RULES
1.  **Language Adaptation:** If the user's original query was in Hinglish, respond in Hinglish. If in English, respond in English.
2.  **Fact-Based Answers:** Use the provided CONTEXT to answer the user's question. If the context is empty, use your Core Business Knowledge.
3.  **Stay on Topic:** Only answer questions related to real estate.
"""


# --- AUDIO & LLM HELPERS ---
def transcribe_audio(audio_path: str, audio_bytes: bytes) -> str:
    for attempt in range(3):
        try:
            audio_file = io.BytesIO(audio_bytes)
            filename = os.path.basename(audio_path)  # e.g., "audio.wav"

            logging.info(f"Transcribing audio: {filename} ({len(audio_bytes)} bytes)")

            transcript = client_openai.audio.transcriptions.create(
                model="whisper-1",
                file=(filename, audio_file)  # ← Critical: gives format hint
            )
            text = transcript.text.strip()

            # Hinglish transliteration
            if re.search(r'[\u0900-\u097F]', text):
                response = client_openai.chat.completions.create(
                    model="gpt-4o-mini",
                    messages=[{"role": "user", "content": f"Transliterate to Roman (Hinglish): {text}"}],
                    temperature=0.0
                )
                text = response.choices[0].message.content.strip()

            logging.info(f"Transcribed: {text}")
            return text

        except Exception as e:
            logging.error(f"Transcription error (attempt {attempt+1}): {e}", exc_info=True) # Added exc_info
            if attempt == 2:
                return ""
    return ""


def generate_elevenlabs_sync(text: str, voice: str) -> bytes:
    for attempt in range(3):
        try:
            return client_elevenlabs.generate(
                text=text,
                voice=voice,
                model="eleven_multilingual_v2",
                output_format="mp3_44100_128"
            )
        except Exception as e:
            logging.error(f"ElevenLabs error (attempt {attempt+1}): {e}", exc_info=True) # Added exc_info
            if attempt == 2:
                return b''
    return b''

# --- UPDATED formulate_search_plan with logging ---
async def formulate_search_plan(user_query: str) -> dict:
    logging.info(f"Formulating search plan for query: {user_query}") # Log incoming query
    for attempt in range(3):
        try:
            response = await run_in_threadpool(
                client_openai.chat.completions.create,
                model=PLANNER_MODEL,
                messages=[{"role": "user", "content": QUERY_FORMULATION_PROMPT.format(user_query=user_query)}],
                response_format={"type": "json_object"},
                temperature=0.0
            )
            # Log the raw response BEFORE trying to parse
            raw_response_content = response.choices[0].message.content
            logging.info(f"Raw Planner LLM response content: {raw_response_content}")

            # Try parsing
            plan = json.loads(raw_response_content)
            logging.info(f"Successfully parsed search plan: {plan}")
            return plan
        except Exception as e:
            # Log the specific error during parsing or API call, with traceback
            logging.error(f"Planner error (attempt {attempt+1}): {e}", exc_info=True)
            if attempt == 2:
                logging.warning("Planner failed after 3 attempts. Using fallback.")
                return {"search_query": user_query, "filter_table": None}
    # Fallback if loop finishes unexpectedly
    logging.error("Planner loop finished unexpectedly. Using fallback.")
    return {"search_query": user_query, "filter_table": None}
# --- END UPDATED FUNCTION ---

async def get_agent_response(user_text: str) -> str:
    for attempt in range(3):
        try:
            plan = await formulate_search_plan(user_text)
            search_query = plan.get("search_query", user_text)
            filter_table = plan.get("filter_table")
            search_filter = {"source_table": filter_table} if filter_table else {}

            docs = await run_in_threadpool(
                vector_store.similarity_search,
                search_query, k=3, filter=search_filter
            )
            if not docs:
                docs = await run_in_threadpool(vector_store.similarity_search, search_query, k=3)

            context = "\n\n".join([d.page_content for d in docs])

            response = await run_in_threadpool(
                client_openai.chat.completions.create,
                model=ANSWERER_MODEL,
                messages=[
                    {"role": "system", "content": ANSWER_SYSTEM_PROMPT},
                    {"role": "system", "content": f"CONTEXT:\n{context}"},
                    {"role": "user", "content": f"Question: {user_text}"}
                ]
            )
            return response.choices[0].message.content.strip()
        except Exception as e:
            logging.error(f"RAG error (attempt {attempt+1}): {e}", exc_info=True) # Added exc_info
            if attempt == 2:
                return "Sorry, I couldn't respond. Please try again."
    return "Sorry, I couldn't respond."


# --- AUTH ENDPOINT ---
class TextQuery(BaseModel):
    query: str

async def verify_token(x_auth_token: str = Header(...)):
    if not SHARED_SECRET or x_auth_token != SHARED_SECRET:
        logging.warning("Auth failed for /test-text-query")
        raise HTTPException(status_code=401, detail="Invalid token")
    logging.info("Auth passed")

@app.post("/test-text-query", dependencies=[Depends(verify_token)])
async def test_text_query_endpoint(query: TextQuery):
    logging.info(f"Text query: {query.query}")
    response = await get_agent_response(query.query)
    return {"response": response}


# --- GRADIO AUDIO PROCESSING ---
async def process_audio(audio_path):
    if not audio_path or not os.path.exists(audio_path):
        return None, "No valid audio file received."

    try:
        # Read raw bytes
        with open(audio_path, "rb") as f:
            audio_bytes = f.read()

        if len(audio_bytes) == 0:
            return None, "Empty audio file."

        # 1. Transcribe — pass path + bytes
        user_text = await run_in_threadpool(transcribe_audio, audio_path, audio_bytes)
        if not user_text:
            return None, "Couldn't understand audio. Try again."

        logging.info(f"User: {user_text}")

        # 2. AI Response
        agent_response = await get_agent_response(user_text)
        if not agent_response:
            return None, "No response generated."

        logging.info(f"AI: {agent_response[:100]}...")

        # 3. Generate Speech
        ai_audio_bytes = await run_in_threadpool(
            generate_elevenlabs_sync, agent_response, ELEVENLABS_VOICE_NAME
        )
        if not ai_audio_bytes:
            return None, "Failed to generate voice."

        # Save to temp file
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as f:
            f.write(ai_audio_bytes)
            out_path = f.name

        return out_path, f"**You:** {user_text}\n\n**AI:** {agent_response}"

    except Exception as e:
        logging.error(f"Audio processing error: {e}", exc_info=True) # Added exc_info
        return None, f"Error: {str(e)}"


# --- GRADIO UI ---
with gr.Blocks(title="Real Estate AI") as demo:
    gr.Markdown("# Real Estate Voice Assistant")
    gr.Markdown("Ask about projects in Pune, Mumbai, Bengaluru, etc.")

    with gr.Row():
        inp = gr.Audio(sources=["microphone"], type="filepath", label="Speak")
        out_audio = gr.Audio(label="AI Response", type="filepath")
    
    out_text = gr.Textbox(label="Conversation", lines=8)

    inp.change(process_audio, inp, [out_audio, out_text])

    # Removed examples to avoid FileNotFoundError with text inputs
    # gr.Examples(examples=[], inputs=inp)


# --- MOUNT GRADIO ---
app = gr.mount_gradio_app(app, demo, path="/")