"""Supplier-offer extraction via Claude.

Inputs are supplier price offers (Czech "nabídka") in wildly varied layouts —
PDF (text or scanned) and XLSX. We extract the SUPPLIER's identity plus the
line items into a fixed schema the CFO defined. Colsys s.r.o. (IČO 14799634)
is always the BUYER, never the supplier.
"""
import base64
import io
import json
import logging
import os
from pathlib import Path

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

logger = logging.getLogger(__name__)

_client: AsyncOpenAI | None = None
MODEL = os.getenv("LLM_MODEL", "anthropic/claude-sonnet-4-20250514")
MAX_PAGES = 8
MAX_IMG_LONG_EDGE = 2200
MAX_XLSX_ROWS = 400          # guard against huge sheets blowing the prompt


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


SYSTEM_PROMPT = """Jste přesný extraktor dat z českých dodavatelských nabídek (cenových nabídek).
Z poskytnutého dokumentu vyextrahujete údaje o DODAVATELI a seznam položek do pevného JSON tvaru.

DŮLEŽITÉ — kdo je dodavatel:
- DODAVATEL = firma, která nabídku VYSTAVILA a nabízí zboží (prodávající).
- Firma „Colsys s.r.o." (IČO 14799634, DIČ CZ14799634) je VŽDY ODBĚRATEL (kupující), NIKDY dodavatel. Pokud ji uvidíte, NEzapisujte ji jako dodavatele.
- Dodavatel je tedy ta druhá strana (např. Rasel, ELKOV elektro, ABBAS apod.).

Vraťte POUZE platný JSON, žádné markdown obaly, žádný komentář mimo JSON:

{
  "is_offer": true,
  "offer_number": "string|null",
  "supplier": {
    "name": "string|null",
    "ico": "string|null",
    "dic": "string|null",
    "email": "string|null",
    "address": "string|null"
  },
  "currency": "CZK",
  "line_items": [
    {
      "name": "string",
      "supplier_id": "string|null",
      "quantity": "number|null",
      "unit": "string|null",
      "unit_price": "number|null",
      "total_price": "number|null"
    }
  ]
}

Pravidla:
- is_offer = false, pokud dokument NENÍ dodavatelská nabídka/cenová nabídka zboží (např. mzdová tabulka, výkaz směn, obecný dokument). V takovém případě nech supplier prázdné a line_items jako [].
- "offer_number" = číslo nabídky dodavatele (např. „Nabídka č. NP250703175", „25NA00023"). Pokud není, null.
- "name" = název zboží přesně jak je uvedeno.
- "supplier_id" = kód/ID zboží u dodavatele (sloupce „Kód", „Obj. kód", „Kód zboží", „Katalogové číslo"). Pokud není, null.
- "quantity" = množství (číslo).
- "unit" = měrná jednotka (m, ks, kg, bal, ...). Zachovej přesně.
- "unit_price" = jednotková cena BEZ DPH za jednotku (sloupce „J.cena", „Cena/MJ", „Cena/mj", „Jednotková cena"). Pokud je uvedena jen cena s DPH, použij ji a poznamenej to není potřeba — prostě vrať dostupnou jednotkovou cenu.
- "total_price" = cena za položku celkem BEZ DPH (sloupce „Celkem bez DPH", „Cena celkem", „Celkem"). Pokud není uvedena, dopočítej quantity × unit_price.
- Čísla zapisuj s tečkou jako desetinným oddělovačem, BEZ mezer a BEZ symbolu měny (např. 1289.40, ne „1 289,40 Kč").
- Zachovej PŘESNÉ pořadí položek tak, jak jsou v dokumentu.
- Vynech řádky, které nejsou skutečné položky zboží (mezisoučty, „Součet položek", „CELKEM K ÚHRADĚ", rekapitulace DPH, prázdné řádky).
- Pokud údaj chybí, použij null. NEVYMÝŠLEJ si hodnoty."""


async def extract_offer(path: Path) -> dict:
    """Extract structured supplier-offer data. Routes PDF→vision, XLSX→text."""
    suffix = path.suffix.lower()
    if suffix == ".pdf":
        content = _pdf_content(path)
    elif suffix in (".xlsx", ".xlsm", ".xls"):
        content = _xlsx_content(path)
    else:
        raise RuntimeError(f"Nepodporovaný formát: {suffix}")

    if not content:
        raise RuntimeError("Soubor neobsahuje žádný čitelný obsah")

    try:
        resp = await _get_client().chat.completions.create(
            model=MODEL,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": content},
            ],
            temperature=0.0,
            max_tokens=8000,
        )
    except Exception as exc:
        logger.exception("LLM call failed")
        raise RuntimeError(f"AI extrakce selhala: {exc}")

    raw = (resp.choices[0].message.content or "").strip()
    raw = raw.removeprefix("```json").removeprefix("```").removesuffix("```").strip()
    try:
        data = json.loads(raw)
    except json.JSONDecodeError as exc:
        logger.error("JSON parse failed: %s\n%s", exc, raw[:500])
        raise RuntimeError(f"Nepodařilo se zpracovat odpověď AI: {exc}")

    return _normalize(data)


ORDER_EMAIL_PROMPT = """Jste extraktor dat z českého e-mailu, kterým zaměstnanec Colsys žádá o vystavení objednávky.
Z textu e-mailu vytáhněte údaje objednávky do pevného JSON tvaru. Text může pocházet z HTML e-mailu (různé zalomení, mezery).

Vraťte POUZE platný JSON, žádné markdown, žádný komentář:

{
  "dodavatel": {
    "name": "string|null",
    "ico": "string|null",
    "dic": "string|null",
    "email": "string|null",
    "kontakt": "string|null"
  },
  "nazev_zakazky": "string|null",
  "kod_zakazky": "string|null",
  "text": "string|null",
  "cena_celkem_text": "string|null",
  "termin_dodani": "string|null",
  "kontaktni_osoba": "string|null",
  "misto_dodani": "string|null",
  "splatnost": "string|null",
  "zaruka": "string|null",
  "prilohy": "string|null",
  "zpusob_odeslani": "string|null"
}

Pravidla:
- Hodnoty hledejte u odpovídajících štítků: „Dodavatel", „IČ"/„IČO", „DIČ", „KONTAKTNÍ OSOBA DODAVATELE" → dodavatel.kontakt (jméno + tel + email; email i do dodavatel.email), „NÁZEV ZAKÁZKY", „KÓD ZAKÁZKY", „TEXT", „CENA CELKEM" → cena_celkem_text (přepište přesně, např. „-" pokud není), „TERMÍN DODÁNÍ", „KONTAKTNÍ OSOBA V MÍSTĚ DODÁNÍ", „MÍSTO DODÁNÍ", „SPLATNOST", „ZÁRUKA", „PŘÍLOHY", „ZPŮSOB ODESLÁNÍ OBJEDNÁVKY".
- TERMÍN DODÁNÍ zapište jako DD.MM.RRRR pokud je to datum.
- Ignorujte patičkový právní text o VNP/VSP a větu „tento text výše je obsažen v zápatí…".
- Pokud údaj chybí, použijte null. NEVYMÝŠLEJTE si."""


async def parse_order_email(body_text: str) -> dict:
    """Extract order-form fields from the structured request e-mail body."""
    body_text = (body_text or "").strip()
    if not body_text:
        return {}
    try:
        resp = await _get_client().chat.completions.create(
            model=MODEL,
            messages=[
                {"role": "system", "content": ORDER_EMAIL_PROMPT},
                {"role": "user", "content": body_text[:20000]},
            ],
            temperature=0.0,
            max_tokens=2000,
        )
    except Exception as exc:
        logger.exception("order-email LLM call failed")
        raise RuntimeError(f"AI zpracování e-mailu selhalo: {exc}")
    raw = (resp.choices[0].message.content or "").strip()
    raw = raw.removeprefix("```json").removeprefix("```").removesuffix("```").strip()
    try:
        return json.loads(raw)
    except json.JSONDecodeError as exc:
        logger.error("order-email JSON parse failed: %s\n%s", exc, raw[:400])
        return {}


def _pdf_content(path: Path) -> list[dict]:
    """PDF → page images for the vision model, plus the extracted text layer
    (helps with exact codes/prices that can be hard to OCR from the raster)."""
    doc = fitz.open(str(path))
    content: list[dict] = [{
        "type": "text",
        "text": "Níže je dodavatelská nabídka jako obrázky stránek "
                "(a pod nimi textová vrstva, pokud existuje). "
                "Vyextrahujte data podle definovaného JSON tvaru.",
    }]
    text_parts = []
    for i, page in enumerate(doc):
        if i >= MAX_PAGES:
            break
        pix = page.get_pixmap(dpi=200, alpha=False)
        img = Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
        content.append({
            "type": "image_url",
            "image_url": {"url": f"data:image/png;base64,{_compress(img)}"},
        })
        txt = page.get_text("text") or ""
        if txt.strip():
            text_parts.append(f"--- strana {i + 1} ---\n{txt.strip()}")
    doc.close()
    if text_parts:
        content.append({
            "type": "text",
            "text": "Textová vrstva PDF:\n" + "\n\n".join(text_parts),
        })
    return content


def _xlsx_content(path: Path) -> list[dict]:
    """XLSX → a compact text dump of every sheet (rows as tab-joined cells)."""
    wb = openpyxl.load_workbook(str(path), data_only=True, read_only=True)
    blocks = []
    for ws in wb.worksheets:
        lines = [f"=== List: {ws.title} ==="]
        for r, row in enumerate(ws.iter_rows(values_only=True)):
            if r >= MAX_XLSX_ROWS:
                lines.append("… (zbytek listu vynechán)")
                break
            cells = ["" if c is None else str(c) for c in row]
            if any(c.strip() for c in cells):
                lines.append("\t".join(cells).rstrip())
        blocks.append("\n".join(lines))
    wb.close()
    dump = "\n\n".join(blocks)
    return [{
        "type": "text",
        "text": "Níže je obsah excelové dodavatelské nabídky (buňky listů). "
                "Vyextrahujte data podle definovaného JSON tvaru.\n\n" + dump,
    }]


def _compress(img: Image.Image) -> str:
    if max(img.size) > MAX_IMG_LONG_EDGE:
        ratio = MAX_IMG_LONG_EDGE / max(img.size)
        img = img.resize((int(img.size[0] * ratio), int(img.size[1] * ratio)),
                         Image.LANCZOS)
    buf = io.BytesIO()
    img.save(buf, format="PNG", optimize=True)
    return base64.b64encode(buf.getvalue()).decode("ascii")


def _normalize(data: dict) -> dict:
    def _num(v):
        if v is None or v == "":
            return None
        if isinstance(v, (int, float)):
            return float(v)
        s = str(v).strip()
        # Czech number formatting: "1 289,40" → "1289.40"
        s = s.replace(" ", "").replace(" ", "").replace(",", ".")
        # strip trailing currency/units the model may have left in
        s = "".join(ch for ch in s if ch.isdigit() or ch in ".-")
        try:
            return float(s) if s not in ("", ".", "-") else None
        except ValueError:
            return None

    if not isinstance(data, dict):
        return {"is_offer": False, "supplier": {}, "line_items": []}

    data.setdefault("is_offer", True)
    data.setdefault("supplier", {})
    data.setdefault("currency", "CZK")
    items = data.get("line_items") or []
    for item in items:
        for k in ("quantity", "unit_price", "total_price"):
            if k in item:
                item[k] = _num(item[k])
        # backfill total if missing but quantity & unit price present
        if item.get("total_price") is None and \
                item.get("quantity") is not None and \
                item.get("unit_price") is not None:
            item["total_price"] = round(item["quantity"] * item["unit_price"], 2)
    data["line_items"] = items
    return data