backend/app/api/api_v1/endpoints/chat.py

import json
from typing import List, Optional
from fastapi import APIRouter, Depends, HTTPException, Request
from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session
from ... import deps
from ....core.config import settings
from ....models.user import User
from ....models.chat import ChatSession, ChatMessage
from ....schemas.chat import ChatQuery
from ....services.rag_service import rag_service
from ....db.init_db import q_client, COLLECTION_NAME, get_vector_store
from ....db.session import SessionLocal
from qdrant_client.http import models as rest
from ....services.tts import get_tts_wav, stream_tts_wav_chunks
import re
from pydantic import BaseModel

def clean_context_text(text: str) -> str:
    if not text:
        return ""
    # Normalize unicode spaces and control characters
    text = text.replace('\xa0', ' ')
    text = text.replace('\u200b', '')
    lines = []
    for line in text.splitlines():
        cleaned = line.strip()
        if not cleaned:
            continue
        # Replace 3 or more repeating divider symbols inside lines with a single one (e.g. ------ to -)
        cleaned = re.sub(r'([=\-_*#\.\~\+\|\\/\u2014])\1{2,}', r'\1', cleaned)
        cleaned = cleaned.strip()
        # Skip lines that collapse to a single divider symbol
        if cleaned in ['=', '-', '_', '*', '#', '.', '~', '+', '|', '\\', '/', '\u2014']:
            continue
        # Remove spaces to check if it's a spacer pattern like ". . . ." or "- - - -"
        no_spaces = re.sub(r'\s+', '', cleaned)
        if re.match(r'^[=\-_*#\.\~\+\|\\/\u2014]{2,}$', no_spaces):
            continue
        # Skip page reference headers and footers
        if re.match(r'^(page\s*\d+|\d+\s*of\s*\d+)$', cleaned, re.IGNORECASE):
            continue
        # Collapse tabs and consecutive spaces
        cleaned = re.sub(r'[ \t]+', ' ', cleaned)
        if cleaned:
            lines.append(cleaned)
    return "\n".join(lines)


router = APIRouter()

class SpeakRequest(BaseModel):
    text: str

async def build_chat_title(query: str) -> str:
    cleaned = re.sub(r"\s+", " ", query or "").strip()
    if not cleaned: return "New Chat"
    try:
        from langchain_core.messages import HumanMessage
        prompt = f"Short title (2-5 words) for: {cleaned}"
        messages = [HumanMessage(content=prompt)]
        response = await rag_service.llm.ainvoke(messages)
        return response.content.strip().strip('"').strip("'")[:60]
    except:
        return cleaned[:30] + "..."

@router.get("/sessions")
def list_sessions(db: Session = Depends(deps.get_db), current_user: User = Depends(deps.get_current_active_user)):
    user_id_str = str(current_user.id)
    sessions = db.query(ChatSession).filter(ChatSession.user_id == user_id_str).order_by(ChatSession.created_at.desc()).all()
    return {"sessions": [{"id": s.id, "title": s.title, "date": s.created_at} for s in sessions]}

@router.get("/history/{session_id}")
def get_history(session_id: str, db: Session = Depends(deps.get_db), current_user: User = Depends(deps.get_current_active_user)):
    user_id_str = str(current_user.id)
    messages = db.query(ChatMessage).filter(
        ChatMessage.user_id == user_id_str,
        ChatMessage.session_id == session_id
    ).order_by(ChatMessage.timestamp.asc()).all()
    
    # Also fetch linked documents
    from ....models.document import Document
    docs = db.query(Document).filter(
        Document.user_id == user_id_str,
        Document.session_id == session_id
    ).all()
    
    return {
        "history": [{"role": m.role, "text": m.content, "sources": json.loads(m.sources) if m.sources else []} for m in messages],
        "documents": [{"filename": d.filename, "chunks": d.chunk_count} for d in docs]
    }

class ChatQuery(BaseModel):
    query: str
    session_id: str
    filename: Optional[str] = None
    filenames: Optional[List[str]] = None

@router.post("/")
async def query_chat(chat_data: ChatQuery, db: Session = Depends(deps.get_db), current_user: User = Depends(deps.get_current_active_user)):
    try:
        user_id_str = str(current_user.id)
        vector_store = get_vector_store(rag_service.embeddings)
        user_id_f = rest.FieldCondition(key="metadata.user_id", match=rest.MatchValue(value=user_id_str))
        
        # 1. Broad Session Search (Selection-Aware)
        must_conditions = [user_id_f, rest.FieldCondition(key="metadata.session_id", match=rest.MatchValue(value=chat_data.session_id))]
        if chat_data.filenames and len(chat_data.filenames) > 0:
            must_conditions.append(rest.FieldCondition(key="metadata.filename", match=rest.MatchAny(any=chat_data.filenames)))
        elif chat_data.filename:
            must_conditions.append(rest.FieldCondition(key="metadata.filename", match=rest.MatchValue(value=chat_data.filename)))

        search_results = vector_store.search(
            query=chat_data.query, 
            search_type="mmr", 
            k=settings.SEARCH_K, 
            fetch_k=settings.FETCH_K,
            filter=rest.Filter(must=must_conditions)
        )
        
        if not search_results:
            async def empty_gen():
                yield "I couldn't find any relevant information across your documents to answer this question."
            return StreamingResponse(empty_gen(), media_type="text/plain")

        # 2. Handle Session & Logging
        session = db.query(ChatSession).filter(ChatSession.id == chat_data.session_id).first()
        if not session:
            title = await build_chat_title(chat_data.query)
            session = ChatSession(id=chat_data.session_id, user_id=user_id_str, title=title)
            db.add(session)
            db.commit()

        db.add(ChatMessage(user_id=user_id_str, session_id=chat_data.session_id, role="user", content=chat_data.query))
        db.commit()

        # 3. Intelligent Grouping
        # Rerank first to ensure we are only using top relevant bits across all files
        candidates = [doc.page_content for doc in search_results]
        scores = rag_service.rerank_results(chat_data.query, candidates)
        scored_hits = sorted(zip(search_results, scores), key=lambda x: x[1], reverse=True)[:settings.RERANK_TOP_K]

        # Group the top hits by filename
        grouped_hits = {}
        all_sources_data = [] # For DB storage
        consolidated_citations = {} # For final display

        for hit, score in scored_hits:
            fname = hit.metadata.get('filename', 'Unknown Document')
            page = hit.metadata.get('page')
            
            all_sources_data.append({"file": fname, "page": page})
            if fname not in consolidated_citations: consolidated_citations[fname] = set()
            if page: consolidated_citations[fname].add(page)
            
            if fname not in grouped_hits: grouped_hits[fname] = []
            context_text = hit.metadata.get('parent_text', hit.page_content)
            clean_text = clean_context_text(context_text)
            grouped_hits[fname].append(f"[Page: {page}]\n{clean_text}")

        unique_files_found = list(grouped_hits.keys())
        
        # Ensure Page 1 context is included for each file (Metadata/Cover Page Injection)
        for fname in unique_files_found:
            has_page_1 = any("[Page: 1]\n" in item for item in grouped_hits[fname])
            if not has_page_1:
                try:
                    page_1_scroll, _ = q_client.scroll(
                        collection_name=COLLECTION_NAME,
                        scroll_filter=rest.Filter(
                            must=[
                                rest.FieldCondition(key="metadata.user_id", match=rest.MatchValue(value=user_id_str)),
                                rest.FieldCondition(key="metadata.filename", match=rest.MatchValue(value=fname)),
                                rest.FieldCondition(key="metadata.page", match=rest.MatchValue(value=1))
                            ]
                        ),
                        limit=10,
                        with_payload=True,
                        with_vectors=False
                    )
                    if page_1_scroll:
                        for point in reversed(page_1_scroll):
                            p_meta = point.payload.get("metadata", {})
                            p_text = p_meta.get("parent_text", point.payload.get("text"))
                            clean_p_text = clean_context_text(p_text)
                            page_1_item = f"[Page: 1]\n{clean_p_text}"
                            if page_1_item not in grouped_hits[fname]:
                                grouped_hits[fname].insert(0, page_1_item)
                except Exception as pg_err:
                    print(f"Error fetching page 1 metadata for {fname}: {pg_err}")

        is_sequential = len(unique_files_found) > 1

        async def response_generator():
            full_answer = ""
            
            for idx, fname in enumerate(unique_files_found):
                # A. Prepare section header
                header = f"### [DOCUMENT: {fname}]\n\n" if is_sequential else ""
                full_answer += header
                if header: yield header
                
                # B. Stream answer for THIS document's context
                doc_context = grouped_hits[fname]
                async for chunk in rag_service.generate_answer_stream(
                    chat_data.query, 
                    doc_context, 
                    brief=is_sequential, 
                    trace_metadata={"user_id": user_id_str, "file": fname}
                ):
                    full_answer += chunk
                    yield chunk
                
                # C. Separator
                if is_sequential and idx < len(unique_files_found) - 1:
                    sep = "\n\n---\n\n"
                    full_answer += sep
                    yield sep

            # Check if the generated answer is a refusal/no-information response
            is_refusal = False
            # Strip section headers/separators to analyze the raw LLM text content
            raw_llm_text = re.sub(r'###\s+\[DOCUMENT:.*?\]', '', full_answer)
            raw_llm_text = re.sub(r'---', '', raw_llm_text)
            text_lower = raw_llm_text.lower().strip()

            refusal_keywords = [
                "no information", "not mentioned", "not found", "not provide", 
                "not in the provided context", "not in context", "does not contain", 
                "doesn't contain", "does not mention", "doesn't mention", 
                "unable to answer", "cannot answer", "no reference", 
                "i don't know", "i do not know", "i couldn't find", "no mention of",
                "does not provide"
            ]
            if len(text_lower) < 500:
                for keyword in refusal_keywords:
                    if keyword in text_lower:
                        is_refusal = True
                        break

            # 4. Deterministic Python Citations (Only append if the query was successfully answered)
            citation_lines = []
            if not is_refusal:
                for f, pages in consolidated_citations.items():
                    sorted_pages = sorted(list(pages))
                    pages_str = ", ".join(map(str, sorted_pages))
                    citation_lines.append(f"[Source: {f}, Pages: {pages_str}]")
            
            python_citation_str = "\n\n***\n" + "\n".join(citation_lines) if citation_lines else ""
            
            if python_citation_str:
                full_answer += python_citation_str
                yield python_citation_str

            # 5. Final Save
            with SessionLocal() as final_db:
                final_db.add(ChatMessage(user_id=user_id_str, session_id=chat_data.session_id, role="assistant", content=full_answer, sources=json.dumps(all_sources_data)))
                final_db.commit()

        return StreamingResponse(response_generator(), media_type="text/plain")
    except Exception as e:
        import traceback
        traceback.print_exc()
        raise HTTPException(
            status_code=500,
            detail=f"Backend Error: {str(e)}"
        )

@router.delete("/session/{session_id}")
def delete_session(session_id: str, db: Session = Depends(deps.get_db), current_user: User = Depends(deps.get_current_active_user)):
    user_id_str = str(current_user.id)
    
    # 1. Cleanup Documents and Embeddings associated with this session
    from ....models.document import Document
    db.query(Document).filter(Document.session_id == session_id, Document.user_id == user_id_str).delete()
    
    q_client.delete(
        collection_name=COLLECTION_NAME,
        points_selector=rest.Filter(
            must=[
                rest.FieldCondition(key="metadata.user_id", match=rest.MatchValue(value=user_id_str)),
                rest.FieldCondition(key="metadata.session_id", match=rest.MatchValue(value=session_id))
            ]
        )
    )

    # 2. Cleanup Messages and Session
    db.query(ChatMessage).filter(ChatMessage.session_id == session_id, ChatMessage.user_id == user_id_str).delete()
    db.query(ChatSession).filter(ChatSession.id == session_id, ChatSession.user_id == user_id_str).delete()
    
    db.commit()
    return {"message": "Session and associated documents deleted"}

@router.post("/speak")
async def speak(request: Request, speak_data: SpeakRequest):
    # Sanitize text for TTS
    clean_text = speak_data.text
    clean_text = re.sub(r'#+\s+', '', clean_text)
    clean_text = re.sub(r'\*+', '', clean_text)
    clean_text = re.sub(r'_{3,}', '', clean_text)
    clean_text = re.sub(r'-{3,}', '', clean_text)
    clean_text = re.sub(r'\[Source:.*?\]', '', clean_text)
    
    # Internal stop signal for THIS specific request
    import threading
    import asyncio
    disconnect_event = threading.Event()

    async def watch_disconnect():
        try:
            while not disconnect_event.is_set():
                if await request.is_disconnected():
                    disconnect_event.set()
                    break
                await asyncio.sleep(0.1)
        except asyncio.CancelledError:
            pass

    watch_task = asyncio.create_task(watch_disconnect())

    # Generator wrapper to monitor disconnection
    async def disconnect_monitor_gen():
        generator = stream_tts_wav_chunks(clean_text, disconnect_event)
        try:
            for chunk in generator:
                if disconnect_event.is_set():
                    break
                yield chunk
        except Exception as e:
            disconnect_event.set()
            raise e
        finally:
            disconnect_event.set()
            watch_task.cancel()

    return StreamingResponse(
        disconnect_monitor_gen(),
        media_type="application/x-ndjson"
    )