"""
The Bullshit Tribunal — Backend proxy.

Routes:
  POST /api/trial  : text -> structured courtroom verdict (JSON)
  POST /api/speak  : text -> WAV audio of the verdict (Kokoro-82M, local)
  GET  /health
  GET  /traces

The model (Qwen2.5-1.5B fine-tuned, served by llama.cpp) returns a single JSON
object. We constrain it with a json_schema + json_object response_format and then
defensively repair/validate the result so a 1.5B model can't break the UI.
All trials are logged to JSONL for the "Sharing is Caring" badge.
"""

import io
import json
import os
import re
import time

import httpx
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response
from pydantic import BaseModel

app = FastAPI(title="The Bullshit Tribunal")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

LLAMA_URL = "http://localhost:8080/v1/chat/completions"
TRACE_FILE = os.environ.get("TRACE_FILE", "/app/agent_traces.jsonl")

SYSTEM_PROMPT = """You are Judge Verity, presiding over The Bullshit Tribunal — a courtroom where manipulative text stands trial. You are jaded, brilliant, and very funny. You have seen every lie ever written and you are unimpressed.

The user submits a piece of text (a job ad, dating profile, terms of service, political speech, marketing claim, corporate email — anything). That text is the DEFENDANT. You prosecute and judge it.

Respond with ONLY a single JSON object, no prose before or after, in exactly this shape:

{
  "case_title": "The People vs. <short label for the text>",
  "bs_score": <integer 0-100, how much bullshit this text contains>,
  "charges": [
    {"count": 1, "name": "<short, creative crime name>", "definition": "<one sentence: what this manipulation is>"}
  ],
  "exhibits": [
    {"label": "A", "quote": "<exact phrase copied from the text>", "severity": "<misdemeanour | felony | capital>"}
  ],
  "cross_exam": [
    {"their_line": "<what they wrote>", "actual_meaning": "<what they actually mean, blunt and funny>"}
  ],
  "verdict": "<ONE punchy memorable sentence. This gets read aloud in court.>",
  "sentence": ["<concrete thing the reader should do>", "<another concrete action>"],
  "hung_jury": <true only if the text is genuinely honest and fair, otherwise false>
}

Rules:
- 2-4 charges, 2-4 exhibits, 2-3 cross_exam pairs, 2-3 sentence actions.
- Exhibit quotes MUST be copied verbatim from the submitted text.
- severity: "misdemeanour" (mildly slimy), "felony" (seriously manipulative), "capital" (egregious).
- Be specific and savage but never slur or attack protected groups; you attack the text, not people.
- If the text is genuinely honest, set hung_jury true, bs_score low, and let the verdict be wry about its rarity.
- Output valid JSON only."""

# llama.cpp can constrain generation to this schema (greatly improves 1.5B reliability).
VERDICT_SCHEMA = {
    "type": "object",
    "properties": {
        "case_title": {"type": "string"},
        "bs_score": {"type": "integer", "minimum": 0, "maximum": 100},
        "charges": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "count": {"type": "integer"},
                    "name": {"type": "string"},
                    "definition": {"type": "string"},
                },
                "required": ["name", "definition"],
            },
        },
        "exhibits": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "label": {"type": "string"},
                    "quote": {"type": "string"},
                    "severity": {"type": "string", "enum": ["misdemeanour", "felony", "capital"]},
                },
                "required": ["quote", "severity"],
            },
        },
        "cross_exam": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "their_line": {"type": "string"},
                    "actual_meaning": {"type": "string"},
                },
                "required": ["their_line", "actual_meaning"],
            },
        },
        "verdict": {"type": "string"},
        "sentence": {"type": "array", "items": {"type": "string"}},
        "hung_jury": {"type": "boolean"},
    },
    "required": ["case_title", "bs_score", "charges", "exhibits", "cross_exam", "verdict", "sentence"],
}


class TrialRequest(BaseModel):
    text: str


class SpeakRequest(BaseModel):
    text: str
    voice: str | None = None


def log_trace(entry: dict):
    """Append a trial trace for the Sharing is Caring badge."""
    try:
        with open(TRACE_FILE, "a") as f:
            f.write(json.dumps(entry) + "\n")
    except OSError:
        pass


def extract_json(content: str) -> dict:
    """Pull the first JSON object out of the model output, tolerating stray prose."""
    content = content.strip()
    # Strip ```json fences if present.
    content = re.sub(r"^```(?:json)?\s*|\s*```$", "", content, flags=re.MULTILINE).strip()
    try:
        return json.loads(content)
    except json.JSONDecodeError:
        pass
    start = content.find("{")
    end = content.rfind("}")
    if start != -1 and end != -1 and end > start:
        try:
            return json.loads(content[start : end + 1])
        except json.JSONDecodeError:
            pass
    raise ValueError("No valid JSON object found in model output")


def normalize_verdict(data: dict, source_text: str) -> dict:
    """Coerce the model output into a guaranteed-valid shape for the UI."""
    out = {}
    out["case_title"] = str(data.get("case_title") or "The People vs. This Text").strip()

    try:
        score = int(float(data.get("bs_score", 50)))
    except (TypeError, ValueError):
        score = 50
    out["bs_score"] = max(0, min(100, score))

    valid_sev = {"misdemeanour", "felony", "capital"}
    charges = []
    for i, c in enumerate(data.get("charges", []) or [], start=1):
        if not isinstance(c, dict):
            continue
        name = str(c.get("name", "")).strip()
        if not name:
            continue
        charges.append({
            "count": int(c.get("count", i)) if str(c.get("count", i)).isdigit() else i,
            "name": name,
            "definition": str(c.get("definition", "")).strip(),
        })
    out["charges"] = charges

    exhibits = []
    for idx, e in enumerate(data.get("exhibits", []) or []):
        if not isinstance(e, dict):
            continue
        quote = str(e.get("quote", "")).strip()
        if not quote:
            continue
        sev = str(e.get("severity", "felony")).strip().lower()
        if sev not in valid_sev:
            sev = "felony"
        label = str(e.get("label") or chr(65 + idx)).strip()
        exhibits.append({"label": label, "quote": quote, "severity": sev})
    out["exhibits"] = exhibits

    cross = []
    for x in data.get("cross_exam", []) or []:
        if not isinstance(x, dict):
            continue
        their = str(x.get("their_line", "")).strip()
        mean = str(x.get("actual_meaning", "")).strip()
        if their or mean:
            cross.append({"their_line": their, "actual_meaning": mean})
    out["cross_exam"] = cross

    out["verdict"] = str(data.get("verdict") or "The court finds this text guilty of wasting everyone's time.").strip()

    sentence = data.get("sentence", []) or []
    if isinstance(sentence, str):
        sentence = [sentence]
    out["sentence"] = [str(s).strip() for s in sentence if str(s).strip()]

    hj = data.get("hung_jury", False)
    if isinstance(hj, str):
        hj = hj.strip().lower() in ("true", "1", "yes")
    out["hung_jury"] = bool(hj)
    return out


@app.post("/api/trial")
async def hold_trial(request: TrialRequest):
    text = request.text.strip()
    if len(text) < 20:
        raise HTTPException(status_code=400, detail="The defendant is too short to stand trial (min 20 characters).")
    if len(text) > 4000:
        raise HTTPException(status_code=400, detail="The defendant is too long (max 4000 characters).")

    payload = {
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": f"The defendant submits the following text. Hold the trial:\n\n{text}"},
        ],
        "temperature": 0.8,
        "top_p": 0.9,
        "max_tokens": 1200,
        # json_object guarantees syntactically valid JSON from llama.cpp; the
        # fine-tune supplies the correct shape, and normalize_verdict() repairs the rest.
        "response_format": {"type": "json_object"},
    }

    try:
        async with httpx.AsyncClient(timeout=120.0) as client:
            resp = await client.post(LLAMA_URL, json=payload)
            resp.raise_for_status()
            raw = resp.json()["choices"][0]["message"]["content"]
    except Exception as e:  # noqa: BLE001
        raise HTTPException(status_code=502, detail=f"The court is in recess (model error): {e}")

    try:
        parsed = extract_json(raw)
        verdict = normalize_verdict(parsed, text)
    except Exception as e:  # noqa: BLE001
        raise HTTPException(status_code=500, detail=f"Mistrial — could not read the verdict: {e}")

    log_trace({
        "timestamp": time.time(),
        "text_preview": text[:160],
        "verdict": verdict,
        "raw_response": raw,
    })
    return verdict


# ----- Kokoro TTS (lazy-loaded; the judge reads the verdict aloud) -----

_kokoro = None
# af_heart is Kokoro's highest-quality voice (grade A). Voices starting with "b"
# are British (en-gb), others American (en-us).
DEFAULT_VOICE = os.environ.get("KOKORO_VOICE", "af_heart")
SPEAK_SPEED = float(os.environ.get("KOKORO_SPEED", "0.92"))  # slightly slow = judicial gravitas


def voice_lang(voice: str) -> str:
    return "en-gb" if voice[:1] == "b" else "en-us"


def get_kokoro():
    global _kokoro
    if _kokoro is None:
        from kokoro_onnx import Kokoro

        model_path = os.environ.get("KOKORO_MODEL", "/app/kokoro/kokoro-v1.0.int8.onnx")
        voices_path = os.environ.get("KOKORO_VOICES", "/app/kokoro/voices-v1.0.bin")
        _kokoro = Kokoro(model_path, voices_path)
    return _kokoro


@app.post("/api/speak")
async def speak(request: SpeakRequest):
    text = request.text.strip()
    if not text:
        raise HTTPException(status_code=400, detail="Nothing to read aloud.")
    text = text[:600]
    try:
        import soundfile as sf

        kokoro = get_kokoro()
        voice = request.voice or DEFAULT_VOICE
        samples, sample_rate = kokoro.create(
            text, voice=voice, speed=SPEAK_SPEED, lang=voice_lang(voice)
        )
        buf = io.BytesIO()
        sf.write(buf, samples, sample_rate, format="WAV")
        return Response(content=buf.getvalue(), media_type="audio/wav")
    except Exception as e:  # noqa: BLE001
        raise HTTPException(status_code=503, detail=f"The judge has lost their voice: {e}")


@app.on_event("startup")
async def warm_kokoro():
    """Load + warm the TTS model in the background so the first verdict isn't slow."""
    import threading

    def _warm():
        try:
            k = get_kokoro()
            k.create("The court is now in session.", voice=DEFAULT_VOICE,
                     speed=SPEAK_SPEED, lang=voice_lang(DEFAULT_VOICE))
        except Exception:  # noqa: BLE001
            pass

    threading.Thread(target=_warm, daemon=True).start()


@app.get("/health")
@app.get("/api/health")
async def health():
    return {"status": "ok", "service": "bullshit-tribunal"}


@app.get("/traces")
@app.get("/api/traces")
async def get_traces():
    if not os.path.exists(TRACE_FILE):
        return {"traces": [], "count": 0}
    traces = []
    with open(TRACE_FILE) as f:
        for line in f:
            if line.strip():
                traces.append(json.loads(line))
    return {"traces": traces, "count": len(traces)}


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=5000)