backend/app.py

"""
themis — FastAPI backend (Phase 2: hybrid statutes + judgments)

Wraps the unified legal RAG router (hybrid_rag/unified_legal_Rag.py) and streams
every step over Server-Sent Events so the UI can render the reasoning live:

    planning  -> route decision (statute / judgment / hybrid)
    retrieval -> statute semantic search + SCI judgment dense+BM25+RRF
    rerank    -> unified cross-encoder rerank across BOTH sources
    answer    -> streamed DeepSeek tokens (intent-specific prompt)
    verify    -> Tier-1 grounding for cited sections AND case citations
    done      -> final answer + citation hyperlinks (sections + cases)

Heavy modules (two Chroma indices, BM25 over 28k judgment chunks, embedding +
cross-encoder models) load lazily on the first /ask, so the server boots
instantly and the SSE stream emits a "warming up" step before the load.
"""

import asyncio
import json
import os
import re
import sys
import time
import logging
from urllib.parse import quote_plus

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel

log = logging.getLogger("themis.backend")
logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(levelname)-8s  %(message)s", datefmt="%H:%M:%S")

BACKEND_DIR = os.path.dirname(os.path.abspath(__file__))
PROJECT_ROOT = os.path.dirname(BACKEND_DIR)

try:
    from dotenv import load_dotenv
    load_dotenv(os.path.join(PROJECT_ROOT, ".env"))
except Exception:
    pass

# The unified router lives in hybrid_rag/ and wires in both pipelines + verifier.
sys.path.insert(0, os.path.join(PROJECT_ROOT, "hybrid_rag"))
sys.path.insert(0, PROJECT_ROOT)

_uni = None


def get_pipeline():
    """Import and cache the unified router (loads both Chroma indices + models)."""
    global _uni
    if _uni is None:
        if not os.getenv("DEEPSEEK_API_KEY"):
            raise RuntimeError("DEEPSEEK_API_KEY is not set in the environment.")
        log.info("Loading unified pipeline (statutes + judgments + models)…")
        import unified_legal_Rag as uni   # noqa: E402  (heavy import on purpose)
        _uni = uni
        log.info("Unified pipeline ready.")
    return _uni


def _norm_sec(section) -> str:
    """'103(1)(a)' -> '103' — compare on the base section number."""
    m = re.match(r"\s*([0-9]+[A-Za-z]?)", str(section))
    return m.group(1) if m else str(section).strip()


# ---------------------------------------------------------------------
# Citation hyperlinks
# ---------------------------------------------------------------------
def _statute_url(act_full: str, act_short: str, section: str) -> str:
    q = quote_plus(f"{act_full or act_short} Section {section}")
    return f"https://indiankanoon.org/search/?formInput={q}"


def _judgment_url(meta: dict) -> str:
    url = meta.get("source_url") or ""
    if url.startswith("http"):
        return url
    cite = meta.get("neutral_citation") or meta.get("case_name") or ""
    return f"https://indiankanoon.org/search/?formInput={quote_plus(cite)}"


def _evidence_item(e) -> dict:
    """Normalise an Evidence object for the UI."""
    m = e.meta or {}
    score = round(float(getattr(e, "unified_score", 0.0)), 3)
    if e.source_type == "judgment":
        return {
            "kind": "judgment",
            "case": m.get("case_name", ""),
            "citation": m.get("neutral_citation", ""),
            "title": m.get("issue", "") or m.get("short_summary", ""),
            "score": score,
            "url": _judgment_url(m),
        }
    return {
        "kind": "statute",
        "act": m.get("act_short", ""),
        "section": str(m.get("section_number", "")),
        "title": m.get("title", ""),
        "score": score,
        "url": _statute_url(m.get("act_name", ""), m.get("act_short", ""), m.get("section_number", "")),
    }


def _sse(obj: dict) -> str:
    return f"data: {json.dumps(obj, ensure_ascii=False)}\n\n"


# =====================================================================
# APP
# =====================================================================
app = FastAPI(title="themis", version="0.2.0")

_origins_env = os.getenv("FRONTEND_ORIGIN", "*")
_allow_origins = ["*"] if _origins_env.strip() == "*" else [o.strip() for o in _origins_env.split(",")]
app.add_middleware(CORSMiddleware, allow_origins=_allow_origins, allow_methods=["*"], allow_headers=["*"])


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


class AskRequest(BaseModel):
    query: str
    history: list[Message] = []
    intent: str = "AUTO"   # "AUTO" -> classify; otherwise force a PROMPTS style


@app.get("/health")
def health():
    return {"status": "ok", "service": "themis", "pipeline_loaded": _uni is not None}


@app.get("/")
def root():
    return JSONResponse({"service": "themis", "version": "0.2.0", "ask": "POST /ask (SSE)"})


@app.post("/ask")
def ask(req: AskRequest):
    def event_stream():
        t_start = time.time()
        query = req.query.strip()
        if not query:
            yield _sse({"type": "error", "message": "Empty query."})
            return

        # Flush an immediate event so the SSE connection opens before the
        # (potentially slow) first-time model + index load.
        yield _sse({"type": "step", "phase": "planning", "title": "Warming up",
                    "detail": "Loading statute + judgment indices and models (first query only)…"})

        try:
            uni = get_pipeline()
        except Exception as e:
            yield _sse({"type": "error", "message": str(e)})
            return

        history = [m.model_dump() for m in req.history]

        # ---- Route ----
        decision = uni.route_query(query)
        srcs = []
        if decision.use_statute:
            srcs.append("statutes")
        if decision.use_judgment:
            srcs.append("Supreme Court judgments")
        yield _sse({"type": "step", "phase": "planning",
                    "title": "Routing the query",
                    "detail": f"Searching: {', '.join(srcs)}.\n{decision.reason}"})

        all_ev = []

        # ---- Statute retrieval ----
        if decision.use_statute:
            yield _sse({"type": "step", "phase": "retrieval",
                        "title": "Searching statutes",
                        "detail": "Query expansion → ChromaDB `indian_statutes` → cross-encoder rerank…"})
            try:
                sev = uni.get_statute_evidence(query)
            except Exception as e:
                log.warning("statute path failed: %s", e)
                sev = []
            all_ev.extend(sev)
            yield _sse({"type": "step", "phase": "retrieval",
                        "title": f"{len(sev)} statute candidate(s)", "detail": ""})

        # ---- Judgment retrieval ----
        if decision.use_judgment:
            yield _sse({"type": "step", "phase": "retrieval",
                        "title": "Searching Supreme Court judgments",
                        "detail": "Dense + BM25 → RRF fusion → dedup → child & parent rerank…"})
            try:
                jev = uni.get_judgment_evidence(query)
            except Exception as e:
                log.warning("judgment path failed: %s", e)
                jev = []
            all_ev.extend(jev)
            yield _sse({"type": "step", "phase": "retrieval",
                        "title": f"{len(jev)} judgment candidate(s)", "detail": ""})

        if not all_ev:
            yield _sse({"type": "error", "message": "No relevant statutes or judgments found for this query."})
            return

        # ---- Unified rerank across both sources ----
        yield _sse({"type": "step", "phase": "rerank",
                    "title": "Reranking all evidence together",
                    "detail": "Scoring statute + judgment candidates on one cross-encoder for a fair merge…"})
        severity = uni._build_severity_context(query)
        unified_query = f"{query} {severity}" if severity else query
        final_ev = uni.unified_rerank(unified_query, all_ev)

        yield _sse({"type": "evidence", "items": [_evidence_item(e) for e in final_ev]})

        context = uni.build_unified_context(final_ev)

        # ---- Intent (user-forced or auto-classified) + streamed answer ----
        forced = (req.intent or "AUTO").upper()
        if forced != "AUTO" and forced in uni.PROMPTS:
            intent = forced
        else:
            intent = uni.classify_intent(query)
        yield _sse({"type": "step", "phase": "answer",
                    "title": f"Drafting the answer · {intent.replace('_', ' ').title()}",
                    "detail": "Generating a grounded answer from the retrieved evidence…"})

        system_prompt = uni.PROMPTS.get(intent, uni.PROMPTS["LEGAL_RESEARCH"])
        messages = [{"role": "system", "content": system_prompt}]
        if history:
            messages.extend(history[-8:])
        messages.append({"role": "user", "content": f"QUESTION:\n{query}\n\nEVIDENCE:\n{context}"})

        full_answer = ""
        try:
            stream = uni.stat_rag.llm_client.chat.completions.create(
                model=uni.LLM_MODEL,
                messages=messages,
                temperature=uni.LLM_TEMPERATURE,
                max_tokens=uni.LLM_MAX_TOKENS,
                stream=True,
            )
            for chunk in stream:
                delta = chunk.choices[0].delta
                if delta.content:
                    full_answer += delta.content
                    yield _sse({"type": "token", "delta": delta.content})
        except Exception as e:
            log.error("DeepSeek error: %s", e)
            yield _sse({"type": "error", "message": f"Answer generation failed: {e}"})
            return

        # ---- Citation hyperlinks (sections + cases actually cited) ----
        clean = full_answer.replace("**", "")
        ev_by_cite = {
            (e.meta.get("neutral_citation") or "").upper(): e.meta
            for e in final_ev if e.source_type == "judgment"
        }
        citations = []
        cited_cases = []          # (citation, case_name) cited in the answer
        seen = set()
        try:
            for act, sec in (uni.stat_rag.parse_section_references(clean) or []):
                key = ("s", act.upper(), str(sec))
                if key in seen:
                    continue
                seen.add(key)
                rec = uni.stat_rag.direct_lookup(act, sec)
                meta = rec["metadata"] if rec else {}
                citations.append({
                    "kind": "statute",
                    "label": f"{act.upper()} Section {sec}",
                    "title": meta.get("title", ""),
                    "url": _statute_url(meta.get("act_name", ""), act, sec),
                })
            for cite in (uni.judg_rag.extract_citations(clean) or []):
                key = ("c", cite.upper())
                if key in seen:
                    continue
                seen.add(key)
                m = ev_by_cite.get(cite.upper(), {})
                citations.append({
                    "kind": "judgment",
                    "label": (m.get("case_name") or cite),
                    "title": cite if m.get("case_name") else "",
                    "url": _judgment_url(m or {"neutral_citation": cite}),
                })
                cited_cases.append((cite, m.get("case_name", "")))
        except Exception as e:
            log.warning("citation links failed: %s", e)

        # Emit the answer immediately; verification (which may hit slow external
        # portals) runs afterwards and streams in as a follow-up `verify` event.
        yield _sse({
            "type": "done",
            "answer": full_answer,
            "intent": intent,
            "route": decision.reason,
            "citations": citations,
            "elapsed_seconds": round(time.time() - t_start, 2),
        })

        # ---- Verification ----
        yield _sse({"type": "step", "phase": "verify",
                    "title": "Verifying citations",
                    "detail": "Checking cited sections against the corpus; verifying cited cases "
                              "against the SCI/eCourts portals + Indian Kanoon…"})

        # Statute sections: in evidence / in DB (retrieval miss) / not in DB (hallucinated)
        ev_sections = {
            (e.meta.get("act_short", "").upper(), _norm_sec(e.meta.get("section_number", "")))
            for e in final_ev if e.source_type == "statute"
        }
        flagged_sections, retrieval_miss = [], []
        try:
            sec_seen = set()
            for act, sec in (uni.stat_rag.parse_section_references(clean) or []):
                k = (act.upper(), _norm_sec(sec))
                if k in sec_seen:
                    continue
                sec_seen.add(k)
                if k in ev_sections:
                    continue
                in_db = (
                    (act.upper(), str(sec)) in uni.stat_rag.section_db
                    or (act.upper(), _norm_sec(sec)) in uni.stat_rag.section_db
                )
                (retrieval_miss if in_db else flagged_sections).append(f"{act.upper()} {sec}")
        except Exception as e:
            log.warning("section verify failed: %s", e)

        # Case citations: in-corpus (real, instant) vs ungrounded (live-verify)
        cases = []
        to_live = []
        case_seen = set()
        for cite, _name in cited_cases:
            if cite.upper() in case_seen:
                continue
            case_seen.add(cite.upper())
            m = ev_by_cite.get(cite.upper())
            if m:
                cases.append({"citation": cite, "case": m.get("case_name", ""),
                              "status": "IN_CORPUS", "url": _judgment_url(m),
                              "note": "in retrieved corpus"})
            else:
                to_live.append(cite)

        if to_live:
            try:
                from verifier import verify_citations
                results = asyncio.run(verify_citations(to_live[:6], concurrency=2))
                for cite, r in zip(to_live, results):
                    link = r.ik_match_url or (next(iter(r.verify_links.values()), "") if r.verify_links else "")
                    cases.append({"citation": cite, "case": (r.ik_match_title or ""),
                                  "status": r.status.value, "url": link, "note": r.note or ""})
            except Exception as e:
                log.warning("live citation verify failed: %s", e)
                for cite in to_live:
                    cases.append({"citation": cite, "case": "", "status": "ERROR",
                                  "url": "", "note": "live verification unavailable"})

        grounded = (not flagged_sections) and all(c["status"] != "NOT_FOUND" for c in cases)
        yield _sse({
            "type": "verify",
            "grounded": grounded,
            "flagged_sections": flagged_sections,
            "retrieval_miss_sections": retrieval_miss,
            "cases": cases,
        })

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