import base64
import csv
import io
import itertools
import json
import os
import re
import threading
import time
import requests
from html import escape
from PIL import Image, ImageOps
import inventory

LLAMA_SERVER = os.environ.get("LLAMA_SERVER", "https://inference.llmops.pl")
TIMEOUT = 300  # seconds; CPU inference is slow

# ---------- Humbrol → Tamiya lookup ----------

def _load_humbrol_tamiya() -> dict[str, tuple[str, str]]:
    path = os.path.join(os.path.dirname(__file__), "data", "humbrol_tamyia.csv")
    result: dict[str, tuple[str, str]] = {}
    with open(path, newline="") as f:
        for row in csv.DictReader(f):
            h = row["humbrol"].strip()
            t = row["tamiya"].strip()
            n = row["tamiya_name"].strip()
            if h and h not in result:
                result[h] = (t, n)
    return result

HUMBROL_TO_TAMIYA = _load_humbrol_tamiya()

# ---------- Format configurations ----------

# One prompt + one regex shared by every document type. The model returns a plain
# `- CODE: NAME` list regardless of source layout, so a single tolerant regex parses
# all of them — no per-format coupling to break when the model varies its output.
LIST_PROMPT = (
    # "List every paint code visible in this image."
    "List all paint codes shown. 'code' = the numeric identifier (e.g. '33'). 'name' = colour name (e.g. 'Matt Black')"
    # "Output one per line, exactly as `- CODE: NAME`, and nothing else."
    # "For a paint given as a mix ratio, join the codes with + (for example 118+34)."
)

# Tolerant by design: the 2B model formats the same prompt differently per input —
# a clean "- 11: Silver" for sparse instruction sheets, but an echoed, noisy
# "- Acrylic Paint 11 - Silver - Metallic - CODE: 11 - Silver" for dense shop
# screenshots. So we don't anchor to line start or a fixed separator; we grab the
# first code-like token, then a ':' or '-' separator, then the name up to the next
# separator/newline. Duplicates (shop lines repeat the code) are deduped downstream.
CODE_REGEX = re.compile(
    r'(?P<code>\d+(?:\+\d+)?)\s*[:\-]\s*(?P<name>[^-:\n]+?)\s*(?=[-–\n]|$)'
)

# ---------- Output constraint (per-request, switchable) ----------
#
# A 2B model can't reliably *follow* a format instruction (longer prompt = more
# random), so instead of asking we constrain the sampler. The constraint is a
# per-request payload field — no model/server change, no restart. Flip OUTPUT_MODE
# and re-run the same image to A/B which mode gives the best recall + stability.
#
#   "json" — response_format json_schema; output parsed with json.loads
#   "line" — GBNF grammar forcing exactly the "- CODE: NAME" lines CODE_REGEX expects
#   "free" — no constraint; free text + tolerant regex (legacy baseline)
OUTPUT_MODE = "json"

# No `pattern` on `code`: this build of llama.cpp accepts the pattern at schema
# conversion but its structured-output response parser then 500s on real input
# (the json_schema→GBNF→parser path is fragile here). So the schema enforces only
# shape; the model tends to dump the whole echoed product line into `code`, which
# we clean up downstream (see _clean_code). The `line` GBNF mode is the stronger
# alternative — it forces digits in the code position without this machinery.
JSON_SCHEMA = {
    "type": "array",
    "items": {
        "type": "object",
        "properties": {
            "code": {"type": "string"},
            "name": {"type": "string"},
        },
        "required": ["code"],
    },
}

# Forces exactly the line shape CODE_REGEX already parses, so the "line" path
# reuses the existing regex downstream with zero new parsing code. The name is
# length-capped ({1,40}): unbounded [^\n]+ let the model loop a phrase forever
# (e.g. "(14ml) each. (14ml) each. …") until max_tokens — the cap forces the
# newline and lets generation terminate.
LINE_GRAMMAR = r'''root ::= item+
item ::= "- " code ": " name "\n"
code ::= [0-9]+ ("+" [0-9]+)?
name ::= [^\n]{1,40}'''

FORMATS = {
    "airfix_shop": {
        "label": "Airfix shop screenshot",
        "prompt": LIST_PROMPT,
        "regex": CODE_REGEX,
    },
    "airfix_instruction": {
        "label": "Airfix paper instruction sheet",
        "prompt": LIST_PROMPT,
        "regex": CODE_REGEX,
    },
}

# ---------- Image & inference ----------

def encode_image(img: Image.Image) -> str:
    """Correct EXIF orientation, resize to max 960px (longest side), encode as base64 JPEG."""
    img = ImageOps.exif_transpose(img)
    img = img.convert("RGB")
    # 960px: A/B test showed 640px is insufficient for small-print instruction sheets
    # (model sees colour names but not numeric codes). 960px recovers 100% recall on
    # paper photos (3000+px originals) with ~60s extra vs 640px. 1280px gave identical
    # recall at +90s — no gain. Shop screenshots are already ≤640px so unchanged.
    img.thumbnail((960, 960))
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=90)
    return base64.b64encode(buf.getvalue()).decode()


def call_llama(img_b64: str, prompt: str) -> str:
    payload = {
        "messages": [{
            "role": "user",
            "content": [
                {"type": "text", "text": prompt},
                {"type": "image_url",
                 "image_url": {"url": f"data:image/jpeg;base64,{img_b64}"}}
            ]
        }],
        "temperature": 0.0,
        # Dense shop screenshots (e.g. dual enamel+acrylic listings) overran 500
        # tokens and truncated mid-JSON, which json.loads can't parse at all. 1024
        # fits them; short lists still stop early on EOS, so only long ones pay.
        "max_tokens": 1024,
        # Critical for multimodal: llama-server reuses KV/prompt cache across
        # consecutive requests (LRU slots), bleeding a previous image's encoding
        # into the next and producing runaway garbage on back-to-back calls. Force
        # a fresh evaluation every request so results are order-independent.
        "cache_prompt": False,
    }
    if OUTPUT_MODE == "json":
        payload["response_format"] = {
            "type": "json_schema",
            "json_schema": {"name": "paint_codes", "schema": JSON_SCHEMA},
        }
    elif OUTPUT_MODE == "line":
        payload["grammar"] = LINE_GRAMMAR
    payload["stream"] = True
    r = requests.post(
        f"{LLAMA_SERVER}/v1/chat/completions",
        json=payload,
        timeout=TIMEOUT,
        stream=True,
    )
    r.raise_for_status()
    content = []
    for line in r.iter_lines():
        if not line:
            continue
        line = line.decode("utf-8") if isinstance(line, bytes) else line
        if line.startswith("data: "):
            data = line[6:]
            if data.strip() == "[DONE]":
                break
            try:
                chunk = json.loads(data)
                delta = chunk["choices"][0]["delta"].get("content", "")
                if delta:
                    content.append(delta)
            except (json.JSONDecodeError, KeyError, IndexError):
                pass
    return "".join(content)

# ---------- Parsing ----------

def strip_parens(text: str) -> str:
    return re.sub(r'\([^)]*\)', '', text)


def parse_entries(text: str, pattern: re.Pattern) -> list[dict]:
    return [m.groupdict() for m in pattern.finditer(text)]


def _loads_tolerant(raw: str):
    """json.loads, but salvage a truncated array (model hit max_tokens mid-entry).
    A cut-off "[ {...}, {...}, {...incomplete" still has complete objects before
    the break: keep through the last '}' and close the array. Returns None if even
    that fails."""
    try:
        return json.loads(raw)
    except (json.JSONDecodeError, TypeError):
        pass
    start, end = raw.find("["), raw.rfind("}")
    if start != -1 and end > start:
        try:
            return json.loads(raw[start:end + 1] + "]")
        except json.JSONDecodeError:
            return None
    return None


def parse_json_entries(raw: str) -> list[dict]:
    """Parse the json_schema-constrained array into {code, name} entries.

    json_schema enforces the array *shape* but not field semantics, so the 2B
    model often dumps the whole echoed product line into `code`
    ("Acrylic Paint 11 - Silver - Metallic"). We pull the first numeric code
    (optionally a mix like "118+34") out of that field — reliable because it
    operates on the already-isolated entry, not the raw response. Items with no
    code-like token are dropped; a broken payload yields []."""
    data = _loads_tolerant(raw)
    if not isinstance(data, list):
        return []
    entries = []
    for item in data:
        if not isinstance(item, dict):
            continue
        m = re.search(r'\d+(?:\+\d+)?', str(item.get("code", "")))
        if m:
            entries.append({"code": m.group(0), "name": item.get("name")})
    return entries


# Instruction sheets mix paint codes with decal numbers, part numbers and dates,
# and the 2B model can't reliably tell them apart — it sometimes emits decals as
# if they were paints (non-deterministic across layouts, though stable per image
# at temp 0). This is a perception problem the format constraint can't fix, so we
# drop entries whose *name* gives them away: a decal/stencil/transfer reference,
# a bare number, or a 4-digit year. A missing name is NOT evidence (the code still
# matters for the Tamiya lookup), so those are kept.
NON_PAINT_NAME = re.compile(r'decal|stencil|transfer|\d{4}', re.IGNORECASE)


def is_paint_name(name) -> bool:
    if name is None:
        return True
    s = str(name).strip()
    if not s:
        return True
    if s.replace("+", "").isdigit():  # name is just a number → likely a misread reference
        return False
    return not NON_PAINT_NAME.search(s)


def extract_entries(raw: str, pattern: re.Pattern) -> list[dict]:
    """Branch parsing on the active output mode (JSON for the json_schema
    constraint, tolerant regex otherwise), then drop non-paint entries."""
    if OUTPUT_MODE == "json":
        entries = parse_json_entries(raw)
    else:
        entries = parse_entries(strip_parens(raw), pattern)
    result = []
    for e in entries:
        name = e.get("name")
        # finetune model echoes the code as name when it has no colour name — treat as unnamed
        if isinstance(name, str) and name.strip() == e.get("code", "").strip():
            e = {**e, "name": None}
        if is_paint_name(e.get("name")):
            result.append(e)
    return result


def expand_codes(entries: list[dict]) -> list[dict]:
    """Split composite codes (e.g. '118+34') into separate items; dedupe."""
    seen = set()
    result = []
    for entry in entries:
        code = entry["code"]
        if "+" in code:
            for c in code.split("+"):
                if c not in seen:
                    seen.add(c)
                    result.append({"code": c, "name": None})
        else:
            if code not in seen:
                seen.add(code)
                result.append(entry)
    return result

# ---------- Main flow ----------

def tamiya_shop_url(code: str, name: str) -> str:
    slug_code = re.sub(r'([A-Za-z]+)(\d+)', r'\1-\2', code)
    slug_name = name.replace(".", "").replace(" ", "-")
    return f"https://www.mojehobby.pl/products/{slug_code}-{slug_name}.html"


LOADING_INTERVAL = 8.0   # seconds between status lines; ~90s inference ≈ one pass of 10 lines

LOADING_LINES = [
    "▸ Establishing uplink to field unit…",
    "▸ Calibrating optical sensors…",
    "▸ Enhancing image resolution…",
    "▸ Scanning paint tin labels…",
    "▸ Identifying colour swatches…",
    "▸ Cross-referencing Humbrol registry…",
    "▸ Computing Tamiya equivalents…",
    "▸ Consulting quartermaster database…",
    "▸ Verifying stock availability…",
    "▸ Assembling resupply manifest…",
]


def _run_analysis(image: Image.Image, format_id: str) -> tuple[str, str]:
    cfg = FORMATS[format_id]
    try:
        img_b64 = encode_image(image)
        raw = call_llama(img_b64, cfg["prompt"])
        entries = extract_entries(raw, cfg["regex"])
        expanded = expand_codes(entries)
        if expanded:
            items = []
            for e in expanded:
                code = e["code"]
                tamiya_code, tamiya_name = HUMBROL_TO_TAMIYA.get(code, (None, ""))
                if tamiya_code is None:
                    items.append(f"<li>Humbrol {escape(code)} – ?</li>")
                elif tamiya_code == "":
                    items.append(f"<li>Humbrol {escape(code)} – (no Tamiya equivalent)</li>")
                elif tamiya_name:
                    if tamiya_code.upper() in inventory.OWNED:
                        items.append(
                            f'<li class="owned">Humbrol {escape(code)} –'
                            f' Tamiya {tamiya_code} {escape(tamiya_name)}'
                            f' <span class="owned-badge">&#10003; ISSUED</span>'
                            f'<span class="owned-status">On personal charge — no resupply required</span></li>'
                        )
                    else:
                        url = tamiya_shop_url(tamiya_code, tamiya_name)
                        items.append(
                            f'<li>Humbrol {escape(code)} –'
                            f' <a href="{url}" target="_blank">Tamiya {tamiya_code} {tamiya_name} ↗</a></li>'
                        )
                else:
                    items.append(f"<li>Humbrol {escape(code)} – Tamiya {tamiya_code}</li>")
            codes_str = "<ul>" + "".join(items) + "</ul>"
        else:
            codes_str = "<p>(none parsed)</p>"
        return raw, codes_str
    except requests.Timeout:
        return "Timeout — model took too long.", ""
    except requests.RequestException as e:
        return f"Server error: {e}", ""


def analyze(image: Image.Image, format_id: str):
    """Generator: stream rotating status lines into Field report while inference
    runs on a worker thread, then yield the final (raw report, manifest HTML)."""
    if image is None:
        yield "No image provided.", ""
        return
    box: dict[str, tuple[str, str]] = {}

    def worker():
        try:
            box["result"] = _run_analysis(image, format_id)
        except Exception as e:  # noqa: BLE001 — never leave box empty
            box["result"] = (f"Error: {e}", "")

    t = threading.Thread(target=worker, daemon=True)
    t.start()
    for line in itertools.cycle(LOADING_LINES):
        if not t.is_alive():
            break
        yield line, ""  # status into Field report; manifest stays empty (collapsed)
        time.sleep(LOADING_INTERVAL)
    t.join()
    yield box.get("result", ("Unknown error.", ""))