"""LLM-based contract analysis against a user-supplied checklist."""
import io
import json
import logging
import os
from pathlib import Path

import fitz  # PyMuPDF
from PIL import Image
import pytesseract
from openai import AsyncOpenAI

logger = logging.getLogger(__name__)

_client: AsyncOpenAI | None = None


def _get_client() -> AsyncOpenAI:
    global _client
    if _client is None:
        _client = AsyncOpenAI(
            base_url=os.getenv("LITELLM_BASE_URL", "http://host.docker.internal:4000/v1"),
            api_key=os.getenv("LITELLM_API_KEY", "sk-dummy"),
        )
    return _client


MODEL = os.getenv("LLM_MODEL", "anthropic/claude-sonnet-4-20250514")


SYSTEM_PROMPT = """Jste expert na české obchodní právo a pomáháte podnikateli posoudit smlouvu, kterou má podepsat.

Vaším úkolem je analyzovat smlouvu z pohledu PŘÍJEMCE smlouvy (toho, kdo ji má podepsat), nikoliv toho, kdo ji sepsal. Hledáte rizika, nevýhodná ustanovení a věci, na které by si měl dát příjemce pozor.

Text smlouvy obdržíte rozdělený podle stran — každá strana začíná značkou "--- Strana N ---". U každého citátu MUSÍTE uvést, na které straně se nachází.

Pro každý bod kontrolního seznamu posoudíte smlouvu a vrátíte JSON s následující strukturou:

{
  "items": [
    {
      "id": "id_polozky",
      "status": "ok" | "warning" | "problem" | "missing",
      "title": "Krátký nadpis nálezu (3-7 slov)",
      "summary": "1-3 věty vysvětlující nález z pohledu příjemce smlouvy",
      "excerpts": [
        {
          "text": "Přesný citát ze smlouvy (krátký, 5-15 slov, ideálně dohledatelný v PDF)",
          "comment": "Stručný komentář k úryvku — co znamená / proč je to relevantní",
          "page": <číslo strany, kde se citát nachází (celé číslo)>
        }
      ]
    }
  ],
  "overall_summary": "Celkové shrnutí 3-5 vět: hlavní rizika a hlavní výhody. ŽÁDNÁ doporučení co dělat.",
  "risk_level": "low" | "medium" | "high"
}

Pravidla:
- Status "ok" = ustanovení je vyvážené a bezpečné pro příjemce
- Status "warning" = stojí za pozornost, ne nutně problém
- Status "problem" = jasně nevýhodné nebo riskantní pro příjemce
- Status "missing" = bod se ve smlouvě nepokrývá (chybí ustanovení)
- Excerpts musí být DOSLOVNÉ citáty ze smlouvy (kvůli zvýraznění v PDF). Pokud nelze citovat (např. status=missing), vraťte prázdné pole [].
- KAŽDÝ citát MUSÍ obsahovat číslo strany ("page"). Stranu poznáte podle značky "--- Strana N ---" před textem.
- NEDÁVEJTE doporučení co s nálezem dělat — uživatel sám ví, co bude řešit. Pouze konstatujte fakta.
- Komentáře piště česky, věcně a stručně.
- Nevymýšlejte si — pokud něco ve smlouvě není, řekněte to.
- Vraťte POUZE JSON, žádný markdown, žádný text okolo."""


async def analyze_contract(pdf_path: Path, checklist_items: list[dict]) -> dict:
    """Run LLM analysis on a contract PDF against the given checklist."""
    text, used_ocr = _extract_text(pdf_path)
    if not text.strip():
        raise RuntimeError(
            "Z PDF se nepodařilo získat žádný text — ani standardním "
            "způsobem ani OCR. Zkuste lepší kvalitu skenu nebo PDF s textovou vrstvou."
        )

    checklist_block = "\n".join(
        f"- id={it['id']}: {it['label']} — {it.get('hint', '')}"
        for it in checklist_items
    )

    user_msg = f"""KONTROLNÍ SEZNAM (co posoudit):
{checklist_block}

TEXT SMLOUVY:
{text}

Vraťte JSON s analýzou každé položky kontrolního seznamu."""

    resp = await _get_client().chat.completions.create(
        model=MODEL,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_msg},
        ],
        temperature=0.1,
        max_tokens=8000,
    )
    raw = (resp.choices[0].message.content or "").strip()
    logger.info("Analysis raw length: %d chars", len(raw))
    # Strip markdown fences if model wrapped it anyway
    raw = raw.removeprefix("```json").removeprefix("```").removesuffix("```").strip()
    try:
        data = json.loads(raw)
    except json.JSONDecodeError as e:
        logger.error("JSON parse failed: %s\n%s", e, raw[:1000])
        raise RuntimeError(f"Nepodařilo se zpracovat odpověď AI: {e}")
    # Tell the caller whether we had to OCR — UI can show a notice and
    # we'll skip in-page highlighting since the original PDF has no text layer.
    data["_used_ocr"] = used_ocr
    return data


async def extract_counterparty(text: str) -> dict:
    """Find the COUNTERPARTY (not Colsys) — its name + IČO — from the contract."""
    prompt = (
        "Z textu smlouvy urči PROTISTRANU (druhou smluvní stranu), tedy NE "
        "společnost Colsys s.r.o. (IČO 14799634), ale toho, s kým Colsys "
        "smlouvu uzavírá. Vrať POUZE JSON: "
        '{"name": "obchodní jméno protistrany|null", "ico": "IČO (8 číslic)|null"}. '
        "Pokud je stran víc, vyber hlavního dodavatele/zhotovitele/poskytovatele. "
        "IČO uveď jen číslice. Pokud nevíš, dej null."
    )
    try:
        resp = await _get_client().chat.completions.create(
            model=MODEL,
            messages=[{"role": "system", "content": prompt},
                      {"role": "user", "content": text[:30000]}],
            temperature=0.0, max_tokens=300,
        )
        raw = (resp.choices[0].message.content or "").strip()
        raw = raw.removeprefix("```json").removeprefix("```").removesuffix("```").strip()
        d = json.loads(raw)
        return {"name": d.get("name"), "ico": d.get("ico")}
    except Exception as exc:
        logger.warning("counterparty extraction failed: %s", exc)
        return {"name": None, "ico": None}


async def suggest_changes(text: str, finding: dict) -> str:
    """Concrete suggested wording / negotiation points for one finding."""
    fin = (f"Bod: {finding.get('title') or finding.get('label') or ''}\n"
           f"Status: {finding.get('status')}\n"
           f"Shrnutí nálezu: {finding.get('summary') or ''}\n"
           f"Citace: " + "; ".join(e.get("text", "")
                                   for e in (finding.get("excerpts") or [])))
    prompt = (
        "Jste expert na české obchodní právo a vyjednávání smluv. Pro daný "
        "nález ze smlouvy navrhněte z pohledu PŘÍJEMCE smlouvy konkrétní úpravy: "
        "(1) co konkrétně v ustanovení změnit a proč, (2) návrh přesnější/"
        "bezpečnější formulace, (3) vyjednávací argument. Stručně, česky, "
        "v odrážkách. Žádný JSON, jen čitelný text."
    )
    resp = await _get_client().chat.completions.create(
        model=MODEL,
        messages=[{"role": "system", "content": prompt},
                  {"role": "user", "content": f"{fin}\n\nKontext (text smlouvy, zkráceno):\n{text[:15000]}"}],
        temperature=0.3, max_tokens=900,
    )
    return (resp.choices[0].message.content or "").strip()


def _chat_messages(text: str, messages: list[dict]) -> list[dict]:
    system = (
        "Jste asistent, který odpovídá na dotazy o konkrétní smlouvě. "
        "Odpovídejte POUZE na základě textu smlouvy níže. Pokud odpověď ve "
        "smlouvě není, řekněte to. Citujte relevantní pasáže. Česky, věcně.\n\n"
        f"TEXT SMLOUVY:\n{text[:60000]}"
    )
    convo = [{"role": "system", "content": system}]
    for m in messages[-12:]:
        if m.get("role") in ("user", "assistant") and m.get("content"):
            convo.append({"role": m["role"], "content": str(m["content"])[:4000]})
    return convo


async def chat_contract(text: str, messages: list[dict]) -> str:
    """Answer a user question grounded in the contract text (non-streaming)."""
    resp = await _get_client().chat.completions.create(
        model=MODEL, messages=_chat_messages(text, messages),
        temperature=0.2, max_tokens=1200,
    )
    return (resp.choices[0].message.content or "").strip()


async def chat_contract_stream(text: str, messages: list[dict]):
    """Async generator yielding answer text deltas (for SSE streaming)."""
    stream = await _get_client().chat.completions.create(
        model=MODEL, messages=_chat_messages(text, messages),
        temperature=0.2, max_tokens=1200, stream=True,
    )
    async for chunk in stream:
        try:
            delta = chunk.choices[0].delta.content or ""
        except (IndexError, AttributeError):
            delta = ""
        if delta:
            yield delta


def extract_text(pdf_path: Path) -> tuple[str, bool]:
    """Public alias for text extraction (text, used_ocr)."""
    return _extract_text(pdf_path)


def _extract_text(pdf_path: Path) -> tuple[str, bool]:
    """Extract text from PDF. Falls back to OCR if the PDF has no text layer.

    Returns (text, used_ocr).
    """
    doc = fitz.open(str(pdf_path))
    parts = []
    total_chars = 0
    for i, page in enumerate(doc):
        t = page.get_text()
        parts.append(f"--- Strana {i + 1} ---\n{t}")
        total_chars += len(t.strip())
    # Heuristic: <20 chars per page on average → looks like a scan
    needs_ocr = doc.page_count > 0 and (total_chars / max(doc.page_count, 1)) < 20
    if not needs_ocr:
        doc.close()
        return "\n\n".join(parts), False

    logger.info("PDF appears to be scanned (~%d chars across %d pages) — running OCR",
                total_chars, doc.page_count)
    parts = []
    for i, page in enumerate(doc):
        # Rasterize page at 250 DPI for OCR quality
        pix = page.get_pixmap(dpi=250, alpha=False)
        img = Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
        try:
            ocr_text = pytesseract.image_to_string(img, lang="ces+eng")
        except pytesseract.TesseractError as e:
            logger.error("Tesseract failed on page %d: %s", i + 1, e)
            ocr_text = ""
        parts.append(f"--- Strana {i + 1} ---\n{ocr_text}")
    doc.close()
    return "\n\n".join(parts), True