"""
CompLAI Python serving layer — backend-agnostic RAG over EU compliance regulations.

Run:
  INFERENCE_BACKEND=laplateforme python serving/app.py     # dev (API)
  INFERENCE_BACKEND=vllm        python serving/app.py     # on-prem GPU
  INFERENCE_BACKEND=mlx         python serving/app.py     # Apple-Silicon sovereign demo

This mirrors the Rust service but in Python — the language Mistral's Applied AI team
actually uses — and proves the same code path serves managed-API and self-hosted
open-weight models identically (the air-gapped story).
"""
from __future__ import annotations
import math
import pathlib
from contextlib import asynccontextmanager

from fastapi import FastAPI
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

import inference
from rag import build_index, PolicyIndex

# ── Calibration: confidence from TOP RETRIEVAL COSINE (see eval/eval_calibration.py) ──
# Fitted by logistic regression of cosine against "answerable-from-corpus", LOO-CV ECE.
PLATT_A, PLATT_B = 0.3322, 1.5689

def calibrate(top_cosine: float) -> float:
    x = max(0.0, min(1.0, top_cosine))
    return 1.0 / (1.0 + math.exp(-(PLATT_A * x + PLATT_B)))

SYSTEM = (
    "You are a precise EU regulatory compliance assistant. Your corpus covers ONLY three "
    "regimes: CBAM (Reg 2023/956), CSRD (Dir 2022/2464), and the EU Taxonomy (Reg 2020/852). "
    "Answer ONLY from the provided context. Be concise and specific. "
    "If a question concerns a different regime (e.g. VAT, customs duties, AML/KYC, MiFID) or "
    "the context does not support an answer, set disposition to needs_review and say it is "
    "outside your covered regulations."
)

REVIEW_THRESHOLD = 0.70
GROUNDING_FLOOR = 0.74
HIGH_RISK = {"prohibited", "needs_review"}


def review_gate(confidence: float, disposition: str, retrieval_score: float) -> tuple[bool, str]:
    if retrieval_score < GROUNDING_FLOOR:
        return True, (f"Weak grounding — closest passage matches at {retrieval_score*100:.0f}% "
                      f"(floor {GROUNDING_FLOOR*100:.0f}%); regulations likely don't cover this.")
    if disposition in HIGH_RISK:
        return True, f"High-risk disposition ({disposition}) — confirm with a compliance officer before relying on this."
    if confidence < REVIEW_THRESHOLD:
        return True, (f"Calibrated confidence {confidence*100:.0f}% is below the "
                      f"{REVIEW_THRESHOLD*100:.0f}% auto-answer threshold — route to human review.")
    return False, "Grounded, low-risk, and above the confidence threshold — safe to auto-answer."


import re
def extract_reference(text: str) -> str | None:
    m = re.search(r"[Aa]rticle\s+(\d+[a-z]?)", text)
    if m:
        return f"Article {m.group(1)}"
    m = re.search(r"ESRS\s+([ESG]\d|\d)", text)
    if m:
        return f"ESRS {m.group(1)}"
    m = re.search(r"[Aa]nnex\s+([IVX]+|\d+)", text)
    if m:
        return f"Annex {m.group(1)}"
    return None

INDEX: PolicyIndex | None = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    global INDEX
    info = inference.backend_info()
    print(f"[app] backend={info['backend']} model={info['chat_model']} url={info['base_url']}")
    INDEX = build_index(inference.embed)
    print(f"[app] index ready: {len(INDEX.chunks)} chunks")
    yield


app = FastAPI(title="CompLAI (Python serving)", lifespan=lifespan)


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

class AskRequest(BaseModel):
    question: str
    history: list[HistoryItem] = []


def _context_and_citations(question: str):
    top, top_score = INDEX.retrieve_scored(question, k=5)
    context = "\n\n".join(f"[{i+1}] ({c.source}) {c.text}" for i, c in enumerate(top))
    citations = [
        {"source": c.source, "excerpt": _best_sentence(c.text), "reference": extract_reference(c.text)}
        for c in top
    ]
    return context, citations, top_score


def _best_sentence(text: str) -> str:
    kws = ("article", "regulation", "directive", "cbam", "csrd", "must", "shall", "€", "%", "tonne")
    cands = [s.strip() for s in text.replace("\n", ". ").split(".") if len(s.strip()) >= 40]
    if not cands:
        return text[:280] + "…"
    best = max(cands, key=lambda s: sum(k in s.lower() for k in kws) * 10
               + (5 if any(ch.isdigit() for ch in s) else 0) + min(len(s), 200) // 20)
    return best[:280] + ("…" if len(best) > 280 else "")


@app.get("/health")
def health():
    return {"status": "ok", **inference.backend_info()}


@app.post("/ask")
def ask(req: AskRequest):
    context, citations, top_score = _context_and_citations(req.question)
    user = f"Context:\n{context}\n\nQuestion: {req.question}"
    history = [h.model_dump() for h in req.history]
    out = inference.structured_answer(SYSTEM, history, user)
    confidence = round(calibrate(top_score), 4)
    route, reason = review_gate(confidence, out["disposition"], top_score)
    return {
        "answer": out["answer"],
        "citations": citations,
        "confidence": confidence,
        "disposition": out["disposition"],
        "route_to_human": route,
        "review_reason": reason,
        "retrieval_score": round(top_score, 4),
        "backend": inference.backend_info()["backend"],
    }


@app.post("/ask/compare")
def compare(req: AskRequest):
    context, citations, top_score = _context_and_citations(req.question)
    user = f"Context:\n{context}\n\nQuestion: {req.question}"
    rag = inference.structured_answer(SYSTEM, [], user)
    base = inference.plain_answer(req.question)
    return {
        "rag": {
            "answer": rag["answer"], "citations": citations,
            "confidence": round(calibrate(top_score), 4),
            "disposition": rag["disposition"],
        },
        "base": {"answer": base, "citations": [], "confidence": 0.5, "disposition": "factual"},
    }


# Serve the existing web UI unchanged
_web = pathlib.Path(__file__).parent.parent / "web"
if _web.exists():
    app.mount("/", StaticFiles(directory=str(_web), html=True), name="web")


if __name__ == "__main__":
    import os, uvicorn
    port = int(os.environ.get("PORT", "8080"))
    uvicorn.run(app, host="0.0.0.0", port=port)