backend/api/query.py

# backend/api/query.py — FIXED: Multi-source, context, and history

from fastapi import APIRouter, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from backend.rag.retriever import retrieve
from backend.rag.generator import generate_answer, generate_answer_stream, _build_citations
from backend.rag.query_classifier import classify_query, QueryAnalysis, extract_source_filter
from backend.rag.multi_retriever import (
    MultiSourceResult, multi_retrieve, retrieve_multi_selected, retrieve_single_source
)
from backend.rag.multi_generator import generate_multi_answer
from backend.rag.image_rag import enrich_query_with_image_context
from backend.database.connection import get_connection
import uuid
import json

router = APIRouter()


class ChatMessageModel(BaseModel):
    role   : str
    content: str


class QueryRequest(BaseModel):
    question        : str
    source_ids      : list[str] | None = None
    history         : list[ChatMessageModel] | None = None
    mode            : str | None = None
    conversation_id : str | None = None
    image_id        : str | None = None
    include_images  : bool = False
    llm_provider    : str | None = "groq"
    is_legal_mode   : bool = False
    legal_filter    : str | None = None # "statute", "judgment", or None
    agentic_mode    : bool = False  # ← Deep Research: 3-stage Planner→Searcher→Validator


def _history_to_dicts(history: list[ChatMessageModel] | None) -> list[dict] | None:
    if not history:
        return None
    return [{"role": m.role, "content": m.content} for m in history]


def _contextualize_query(question: str, history: list[dict] | None) -> str:
    """
    Rewrite the user's question by injecting the last assistant response
    as context. This resolves pronouns ("they", "it", "those", "that")
    so the classifier and reranker see a self-contained question.

    Example:
        history[-1] = {role: assistant, content: "Cipher techniques include...
                       digital signatures are used for..."}
        question = "how are they different from message digests?"
        → contextualized = "[Context: Cipher techniques include...
                             digital signatures are used for...]
                             how are they different from message digests?"

    The LLM receives the original `question` for display purposes.
    The `contextualized` version is only used for retrieval + classification.
    """
    if not history or not question.strip():
        return question

    # Only do this if query contains common pronouns / relative references
    import re
    pronoun_pattern = re.compile(
        r'\b(they|them|their|it|its|this|that|those|these|he|she|his|her|the same|the above)\b',
        re.IGNORECASE
    )
    if not pronoun_pattern.search(question):
        return question

    # Find the last assistant message
    last_assistant = None
    for msg in reversed(history):
        if msg.get("role") == "assistant" and msg.get("content"):
            last_assistant = msg["content"]
            break

    if not last_assistant:
        return question

    # Take first 400 chars of last answer as context prefix (keeps prompt short)
    context_snippet = last_assistant[:400].strip()
    if len(last_assistant) > 400:
        context_snippet += "..."

    contextualized = f"[Previous context: {context_snippet}]\n{question}"
    print(f"[Query] Contextualized query for pronoun resolution ({len(question)} → {len(contextualized)} chars)")
    return contextualized


def _safe_classify(question: str) -> QueryAnalysis:
    """Always returns a valid QueryAnalysis, never raises."""
    try:
        return classify_query(question)
    except Exception as e:
        print(f"[Query] Classifier failed, using default: {e}")
        return QueryAnalysis(
            intent="single_source", source_types=["any"], topics=[],
            ipc_sections=[], time_filter=None, language_hint="en",
            requires_compare=False, requires_summary=False, source_names=[]
        )


def _do_retrieve(question: str, req_source_ids: list[str] | None, analysis: QueryAnalysis) -> MultiSourceResult:
    """
    THE RETRIEVAL ROUTER.
    Priority:
      1. User explicitly selected sources → use retrieve_multi_selected (CORE FEATURE)
         - 1 source → single_source path
         - 2+ sources → multi_selected path (synthesis intent forced)
      2. Otherwise → let classifier intent decide
    """
    if req_source_ids and len(req_source_ids) > 0:
        if len(req_source_ids) == 1:
            # Single explicit source
            return retrieve_single_source(question, source_ids=req_source_ids)
        else:
            # MULTI-SOURCE CONSOLIDATION — the core product feature
            # Force synthesis intent regardless of what the classifier said
            return retrieve_multi_selected(question, source_ids=req_source_ids)
    else:
        # No manual selection — let classifier decide
        return multi_retrieve(question, analysis)


def _ensure_conversation(cursor, conv_id: str | None, question: str, conv_type: str = "general") -> str:
    if conv_id:
        return conv_id
    new_id = str(uuid.uuid4())
    title = question[:60] + ("..." if len(question) > 60 else "")
    cursor.execute(
        "INSERT INTO conversations (id, title, conv_type) VALUES (%s, %s, %s)", 
        (new_id, title, conv_type)
    )
    return new_id


def _save_to_db(chat_id: str, conv_id: str, question: str, answer: str, source_ids_used: list[str]) -> None:
    try:
        with get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute(
                "INSERT INTO chat_history (id, question, answer, sources_used, conversation_id) VALUES (%s, %s, %s, %s, %s)",
                (chat_id, question, answer, json.dumps(source_ids_used), conv_id)
            )
            cursor.execute("UPDATE conversations SET updated_at = NOW() WHERE id = %s", (conv_id,))
            conn.commit()
            print(f"[Query] Saved chat {chat_id[:8]} to conv {conv_id[:8]}")
    except Exception as e:
        print(f"[Query] DB save warning: {e}")


def _pre_create_conv(conv_id: str | None, question: str, conv_type: str = "general") -> str:
    """Pre-create a conversation row before streaming starts so meta event can carry real ID."""
    if conv_id:
        return conv_id
    try:
        with get_connection() as conn:
            cursor = conn.cursor()
            new_id = _ensure_conversation(cursor, None, question, conv_type)
            conn.commit()
            return new_id
    except Exception as e:
        print(f"[Query] Pre-create conv warning: {e}")
        return str(uuid.uuid4())


def _format_chunks_out(chunks: list) -> list[dict]:
    formatted = []
    for i, c in enumerate(chunks):
        # Build timestamped URL for YouTube
        final_url = c.url_ref
        if c.source_type == "youtube" and c.url_ref and c.timestamp_s is not None:
            sep = "&" if "?" in c.url_ref else "?"
            final_url = f"{c.url_ref}{sep}t={c.timestamp_s}s"
        
        # Human-readable timestamp
        time_str = None
        if c.timestamp_s is not None:
            time_str = f"{c.timestamp_s // 60}:{c.timestamp_s % 60:02d}"

        formatted.append({
            "id": c.chunk_id,
            "sourceId": c.source_id,
            "sourceName": c.source_title,
            "sourceType": c.source_type,
            "text": c.chunk_text,
            "similarityScore": round(c.score, 4),
            "language": c.language or "en",
            "metadata": {
                "page": c.page_number,
                "timestamp": time_str,
                "url": final_url
            }
        })
    return formatted



def _format_citations_out(citations: list) -> list[dict]:
    return [
        {
            "sourceId": c.source_id,
            "sourceType": c.source_type,
            "sourceTitle": c.source_title,
            "reference": c.reference,
            "snippet": c.snippet,
            "score": round(c.score, 4),
        }
        for c in citations
    ]


# ── /query — standard chitchat / conversational fallback helper ────────────────

def build_chat_prompt(question: str, history: list[dict] | None = None) -> list[dict]:
    """
    Build the messages prompt for a general conversational turn.
    Avoids retrieval entirely.
    """
    messages = [
        {
            "role": "system",
            "content": (
                "You are InteleX, a premium Senior Staff AI Research Assistant. "
                "The user is engaging in general chitchat or casual conversation, or asking about your capabilities. "
                "Respond in a warm, professional, and friendly manner. "
                "Briefly mention that you are equipped to perform Multi-Source Agentic RAG "
                "across PDFs, images, websites, and YouTube video transcripts, and you can "
                "do comparative and synthetic analyses. Keep your answer engaging, helpful, and concise."
            )
        }
    ]
    if history:
        for m in history:
            messages.append({"role": m["role"], "content": m["content"]})
    messages.append({"role": "user", "content": question})
    return messages


# ── /query — standard non-streaming ──────────────────────────────────────────

@router.post("/query")
def query(req: QueryRequest):
    if not req.question.strip():
        raise HTTPException(status_code=400, detail="Question cannot be empty.")

    history = _history_to_dicts(req.history)
    enriched_question, image_context_block = enrich_query_with_image_context(
        req.question, image_id=req.image_id, include_recent=req.include_images
    )

    # Contextualize: expand pronouns using the last assistant turn
    retrieval_question = _contextualize_query(enriched_question, history)
    analysis = _safe_classify(retrieval_question)

    # Graceful fallback: check if there are no sources at all in the DB
    has_sources = True
    try:
        with get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT COUNT(*) FROM sources")
            source_count = cursor.fetchone()[0]
            if source_count == 0:
                has_sources = False
    except Exception as e:
        print(f"[Query] Error checking sources count: {e}")

    # ROUTING DECISION: Conversational CHAT Mode
    if analysis.route == "chat" or not has_sources:
        chat_id = str(uuid.uuid4())
        conv_id = req.conversation_id or ""
        if not conv_id:
            try:
                with get_connection() as conn:
                    cursor = conn.cursor()
                    conv_id = _ensure_conversation(cursor, None, req.question, "general")
                    conn.commit()
            except Exception as e:
                print(f"[Query] Conv create warning: {e}")
                conv_id = str(uuid.uuid4())

        if not has_sources and analysis.route != "chat":
            answer = (
                "It looks like no knowledge sources have been added to my database yet! "
                "Please upload a PDF document, add a website URL, or ingest a YouTube video in the sidebar "
                "or tabs first, so that I can analyze and answer questions based on your specific documents."
            )
        else:
            messages = build_chat_prompt(req.question, history)
            try:
                from backend.rag.generator import _get_groq_client, GROQ_MODEL, GROQ_TIMEOUT
                client = _get_groq_client()
                resp = client.chat.completions.create(
                    model=GROQ_MODEL,
                    messages=messages,
                    timeout=GROQ_TIMEOUT,
                )
                answer = resp.choices[0].message.content.strip()
            except Exception as e:
                answer = f"Hello! I am InteleX, your Staff AI Research Assistant. I'm ready to assist, but I encountered an error generating a response: {e}"

        _save_to_db(chat_id, conv_id, req.question, answer, [])

        return {
            "chatId": chat_id,
            "conversationId": conv_id,
            "answer": answer,
            "citations": [],
            "retrievedChunks": [],
            "query_intent": "chat",
            "imageContextUsed": False,
        }

    # ROUTING DECISION: RAG Mode (Retrieve & Generate)
    try:
        multi_result = _do_retrieve(retrieval_question, req.source_ids, analysis)
        chunks = multi_result.all_chunks
        
        # ── Mandatory Image Consideration ─────────────────────────────────────
        if req.image_id:
            from backend.rag.retriever import fetch_image_chunk
            img_chunk = fetch_image_chunk(req.image_id)
            if img_chunk:
                chunks = [img_chunk] + chunks
                multi_result.all_chunks = chunks
                if img_chunk.source_title not in multi_result.source_groups:
                    multi_result.source_groups[img_chunk.source_title] = [img_chunk]
                    multi_result.source_count = len(multi_result.source_groups)
                print(f"[Query] Injected image chunk and group for {req.image_id}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Retrieval error: {str(e)}")

    if not chunks:
        return {
            "chatId": str(uuid.uuid4()), "conversationId": req.conversation_id,
            "answer": "No relevant information found in the selected sources. Please check that documents have been uploaded and try a different question.",
            "citations": [], "retrievedChunks": [], "query_intent": analysis.intent, "imageContextUsed": False
        }

    augmented_history = list(history) if history else []
    if image_context_block:
        augmented_history = [{"role": "system", "content": image_context_block}] + augmented_history

    try:
        is_legal = req.is_legal_mode or (req.llm_provider == "huggingface")
        result = generate_answer(
            req.question, 
            multi_result, 
            history=augmented_history, 
            image_context=image_context_block, 
            provider_name=req.llm_provider,
            is_legal=is_legal
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Generation error: {str(e)}")

    chat_id = str(uuid.uuid4())
    source_ids_used = list({c.source_id for c in chunks})

    conv_id = req.conversation_id or ""
    if not conv_id:
        try:
            with get_connection() as conn:
                cursor = conn.cursor()
                conv_type = "legal" if is_legal else "general"
                conv_id = _ensure_conversation(cursor, None, req.question, conv_type)
                conn.commit()
        except Exception as e:
            print(f"[Query] Conv create warning: {e}")

    _save_to_db(chat_id, conv_id, req.question, result.answer, source_ids_used)

    return {
        "chatId": chat_id, "conversationId": conv_id,
        "answer": result.answer,
        "citations": _format_citations_out(result.citations),
        "retrievedChunks": _format_chunks_out(result.chunks),
        "query_intent": analysis.intent,
        "imageContextUsed": bool(image_context_block),
    }


# ── /query-stream — SSE streaming (PRIMARY PATH) ─────────────────────────────

@router.post("/query-stream")
def query_stream(req: QueryRequest):
    if not req.question.strip():
        raise HTTPException(status_code=400, detail="Question cannot be empty.")

    history = _history_to_dicts(req.history)
    enriched_question, image_context_block = enrich_query_with_image_context(
        req.question, image_id=req.image_id, include_recent=req.include_images
    )

    # Contextualize: expand pronouns using the last assistant turn
    retrieval_question = _contextualize_query(enriched_question, history)

    # 1. Classify (on the contextualized question for better intent detection)
    analysis = _safe_classify(retrieval_question)

    # Graceful fallback: check if there are no sources at all in the DB
    has_sources = True
    try:
        with get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT COUNT(*) FROM sources")
            source_count = cursor.fetchone()[0]
            if source_count == 0:
                has_sources = False
    except Exception as e:
        print(f"[QueryStream] Error checking sources count: {e}")

    # ROUTING DECISION: Conversational CHAT Mode
    if analysis.route == "chat" or not has_sources:
        chat_id = str(uuid.uuid4())
        is_legal = req.is_legal_mode or (req.llm_provider == "huggingface")
        conv_type = "legal" if is_legal else "general"
        conv_id = _pre_create_conv(req.conversation_id, req.question, conv_type)

        def chat_event_stream():
            yield f"data: {json.dumps({'type': 'meta', 'chatId': chat_id, 'conversationId': conv_id, 'citations': [], 'retrievedChunks': [], 'sourceCount': 0, 'activeProvider': req.llm_provider or 'groq'})}\n\n"
            
            if not has_sources and analysis.route != "chat":
                fallback_msg = (
                    "It looks like no knowledge sources have been added to my database yet! "
                    "Please upload a PDF document, add a website URL, or ingest a YouTube video in the sidebar "
                    "or tabs first, so that I can analyze and answer questions based on your specific documents."
                )
                yield f"data: {json.dumps({'type': 'token', 'content': fallback_msg})}\n\n"
                yield f"data: {json.dumps({'type': 'done'})}\n\n"
                _save_to_db(chat_id, conv_id, req.question, fallback_msg, [])
                return

            messages = build_chat_prompt(req.question, history)
            collected = []
            try:
                from backend.rag.generator import _get_groq_client, GROQ_MODEL, GROQ_TIMEOUT
                client = _get_groq_client()
                stream = client.chat.completions.create(
                    model=GROQ_MODEL,
                    messages=messages,
                    stream=True,
                    timeout=GROQ_TIMEOUT,
                )
                for chunk_response in stream:
                    token = chunk_response.choices[0].delta.content
                    if token is not None:
                        collected.append(token)
                        yield f"data: {json.dumps({'type': 'token', 'content': token})}\n\n"
            except Exception as e:
                yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
                return

            yield f"data: {json.dumps({'type': 'done'})}\n\n"
            
            full_answer = "".join(collected).strip()
            _save_to_db(chat_id, conv_id, req.question, full_answer, [])

        return StreamingResponse(
            chat_event_stream(),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}
        )

    # ROUTING DECISION: RAG Mode (Retrieve & Generate)
    try:
        multi_result = _do_retrieve(retrieval_question, req.source_ids, analysis)
        chunks = multi_result.all_chunks
        
        # ── Mandatory Image Consideration ─────────────────────────────────────
        if req.image_id:
            from backend.rag.retriever import fetch_image_chunk
            img_chunk = fetch_image_chunk(req.image_id)
            if img_chunk:
                chunks = [img_chunk] + chunks
                if img_chunk.source_title not in multi_result.source_groups:
                    multi_result.source_groups[img_chunk.source_title] = [img_chunk]
                    multi_result.source_count = len(multi_result.source_groups)
                print(f"[QueryStream] Injected image chunk and group for {req.image_id}")
        
        print(f"[QueryStream] Retrieved {len(chunks)} chunks from {multi_result.source_count} sources")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Retrieval error: {str(e)}")

    chat_id = str(uuid.uuid4())
    is_legal = req.is_legal_mode or (req.llm_provider == "huggingface")
    conv_type = "legal" if is_legal else "general"
    conv_id = _pre_create_conv(req.conversation_id, req.question, conv_type)

    citations_out = _format_citations_out(_build_citations(chunks))
    chunks_out = _format_chunks_out(chunks)

    augmented_history = list(history) if history else []
    if image_context_block:
        augmented_history = [{"role": "system", "content": image_context_block}] + augmented_history

    def event_stream():
        nonlocal multi_result, chunks

        if req.agentic_mode:
            try:
                from backend.rag.agent_workflow import run_agentic_workflow

                def _retriever_fn(q: str, sids):
                    from backend.api.query import _safe_classify, _do_retrieve
                    analysis = _safe_classify(q)
                    return _do_retrieve(q, sids or req.source_ids, analysis)

                yield f"data: {json.dumps({'type': 'agent_status', 'stage': 0, 'message': '🚀 Deep Research Mode activated — starting multi-stage analysis...'})}\n\n"

                is_legal_flag = req.is_legal_mode or (req.llm_provider == "huggingface")
                multi_result, status_log = run_agentic_workflow(
                    question=req.question,
                    retriever_fn=_retriever_fn,
                    source_ids=req.source_ids,
                    is_legal=is_legal_flag,
                )
                chunks = multi_result.all_chunks

                for i, msg in enumerate(status_log):
                    yield f"data: {json.dumps({'type': 'agent_status', 'stage': i + 1, 'message': msg})}\n\n"

            except Exception as e:
                print(f"[QueryStream] Agentic workflow error: {e} — falling back to standard retrieval")
                yield f"data: {json.dumps({'type': 'agent_status', 'stage': 0, 'message': f'⚠️ Deep research unavailable ({str(e)[:60]}), using standard retrieval'})}\n\n"

        citations_final = _format_citations_out(_build_citations(chunks))
        chunks_final = _format_chunks_out(chunks)
        yield f"data: {json.dumps({'type': 'meta', 'chatId': chat_id, 'conversationId': conv_id, 'citations': citations_final, 'retrievedChunks': chunks_final, 'sourceCount': multi_result.source_count, 'activeProvider': req.llm_provider or 'groq'})}\n\n"

        if not chunks:
            yield f"data: {json.dumps({'type': 'token', 'content': 'No relevant information found in the selected sources. Please check that documents have been uploaded and try a different question.'})}\n\n"
            yield f"data: {json.dumps({'type': 'done'})}\n\n"
            return

        collected = []
        try:
            is_legal = req.is_legal_mode or (req.llm_provider == "huggingface")
            for token in generate_answer_stream(
                req.question, 
                multi_result, 
                history=augmented_history, 
                image_context=image_context_block, 
                provider_name=req.llm_provider,
                is_legal=is_legal
            ):
                collected.append(token)
                yield f"data: {json.dumps({'type': 'token', 'content': token})}\n\n"
        except Exception as e:
            yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
            return

        yield f"data: {json.dumps({'type': 'done'})}\n\n"

        full_answer = "".join(collected).strip()
        source_ids_used = list({c.source_id for c in chunks})
        _save_to_db(chat_id, conv_id, req.question, full_answer, source_ids_used)

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}
    )


# ── /query/debug — Diagnostic ─────────────────────────────────────────────────

@router.get("/query/debug")
def query_debug(question: str = "test query"):
    """GET /query/debug?question=... — trace FAISS + MySQL pipeline."""
    from backend.ingestion.embedder import embed_query
    from backend.vectorstore import search_vectors, get_stats
    import os

    results = {"question": question, "faiss_stats": get_stats(), "steps": []}

    try:
        results["steps"].append("1. Embedding query...")
        vec = embed_query(question)

        results["steps"].append("2. Searching FAISS...")
        raw_hits = search_vectors(vec, top_k=5)
        results["faiss_hits"] = raw_hits

        if not raw_hits:
            results["steps"].append("WARNING: FAISS returned 0 hits.")
            return results

        results["steps"].append(f"3. Querying MySQL for {len(raw_hits)} IDs...")
        chunk_ids = [h["chunk_id"] for h in raw_hits]
        placeholders = ", ".join(["%s"] * len(chunk_ids))
        with get_connection() as conn:
            cursor = conn.cursor(dictionary=True)
            cursor.execute(f"SELECT id, source_id, chunk_text FROM chunks WHERE id IN ({placeholders})", chunk_ids)
            db_rows = cursor.fetchall()
            for r in db_rows:
                r["snippet"] = (r.get("chunk_text") or "")[:100] + "..."
                r.pop("chunk_text", None)
            results["db_rows"] = db_rows
            results["steps"].append(f"Found {len(db_rows)} matching rows in DB.")

    except Exception as e:
        results["error"] = str(e)

    return results