"""
LLM verification — Gemma 4 12B on transformers + ZeroGPU.
Utilizes Gemma 4's raw image patch projection to read native Bengali without OCR encoders.
"""

import io, json, os, fitz, spaces, transformers
from PIL import Image
from transformers import AutoProcessor

# Configured for Gemma 4 12B (decoder-only multimodal)
MODEL_ID = os.environ.get("PORCHA_MODEL_ID", "google/gemma-4-12b-it")

_MAX_IMAGE_SIDE = 896
_MAX_IMAGES = 10
_PDF_DPI = 150

class LLMConfigError(RuntimeError): pass

def _model_class():
    for name in ("AutoModelForImageTextToText", "AutoModelForCausalLM"):
        cls = getattr(transformers, name, None)
        if cls is not None: return cls
    raise LLMConfigError("No suitable model class found.")

try:
    processor = AutoProcessor.from_pretrained(MODEL_ID)
    model = _model_class().from_pretrained(MODEL_ID, dtype="auto", device_map="cuda")
except Exception as e:
    processor, model, _LOAD_ERROR = None, None, e
else:
    _LOAD_ERROR = None

def _downscale(img):
    img = img.convert("RGB")
    w, h = img.size
    scale = min(1.0, _MAX_IMAGE_SIDE / max(w, h))
    if scale < 1.0: img = img.resize((max(1, int(w * scale)), max(1, int(h * scale))))
    return img

def _file_to_images(path):
    if path.lower().endswith(".pdf"):
        imgs = []
        with fitz.open(path) as doc:
            for page in doc:
                pix = page.get_pixmap(dpi=_PDF_DPI)
                imgs.append(_downscale(Image.open(io.BytesIO(pix.tobytes("png")))))
        return imgs
    try: return [_downscale(Image.open(path))]
    except: return []

def _system_prompt():
    return (
        "You are an expert West Bengal Land Revenue Officer (BL&LRO). You review "
        "property documents like Porchas and Dalils written in English and Bengali. "
        "You ingest raw images natively. Read the Bengali text, extract details, and "
        "verify against the strict guidelines provided.\n\n"
        "1. For each Document ID, check if it is present. \n"
        "2. For each Rule ID, check if the condition is met (e.g., 'Bastu' vs 'Sali'). "
        "Translate Bengali keywords internally, but quote the original Bengali word "
        "in your 'evidence' string.\n"
        "3. Output strictly as JSON. No prose. Format:\n"
        "{\n"
        '  "documents": [{"id": str, "present": bool, "confidence": "high|medium|low", "evidence": str}],\n'
        '  "rules": [{"id": str, "status": "pass|fail|unclear", "evidence": str}],\n'
        '  "overall": "ready|not_ready",\n'
        '  "summary": str\n'
        "}"
    )

def _build_messages(pkg, stage_label, docs, rules, files):
    text = f"{_system_prompt()}\n\nTRANSACTION: {pkg['name']}\nSTAGE: {stage_label}\n"
    text += "DOCUMENTS:\n" + "\n".join([f"- id={d['id']} | {d['label']}" for d in docs])
    if rules:
        text += "\nRULES:\n" + "\n".join([f"- id={r['id']} | {r['check']}" for r in rules])
    
    content = [{"type": "text", "text": text}]
    n = 0
    for path in files:
        if n >= _MAX_IMAGES: break
        for img in _file_to_images(path):
            if n >= _MAX_IMAGES: break
            content.append({"type": "image", "image": img})
            n += 1
    return [{"role": "user", "content": content}]

@spaces.GPU(duration=120)
def _generate(messages, max_new_tokens):
    inputs = processor.apply_chat_template(messages, tokenize=True, return_dict=True, return_tensors="pt", add_generation_prompt=True).to(model.device)
    out = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
    gen = out[0][inputs["input_ids"].shape[-1]:]
    return processor.decode(gen, skip_special_tokens=True)

def verify(pkg, stage_label, docs, rules, files):
    if model is None: raise LLMConfigError(f"Model error: {_LOAD_ERROR}")
    messages = _build_messages(pkg, stage_label, docs, rules, files)
    raw = _generate(messages, 1024)
    
    # Robust JSON extraction
    text = raw.strip().strip("`")
    if text.lower().startswith("json"): text = text[4:].strip()
    try:
        parsed = json.loads(text)
    except json.JSONDecodeError:
        start, end = text.find("{"), text.rfind("}")
        if start != -1 and end != -1:
            try: parsed = json.loads(text[start : end + 1])
            except: return f"## ⚠️ Output Parse Error\n\n```\n{raw}\n```"
        else:
            return f"## ⚠️ Output Parse Error\n\n```\n{raw}\n```"
            
    # Rendering Logic
    lines = [f"## Verification Report: {pkg['name']}", f"**Status:** {'✅ Looks Good' if parsed.get('overall') == 'ready' else '⚠️ Discrepancies Found'}", f"> {parsed.get('summary', '')}\n", "### Document Verification"]
    
    res_docs = {d["id"]: d for d in parsed.get("documents", [])}
    for d in docs:
        r = res_docs.get(d["id"], {})
        icon = "✅" if r.get("present") else "❌"
        lines.append(f"- {icon} **{d['label']}**: {r.get('evidence', 'Not evaluated')}")
        
    if rules:
        lines.append("\n### Legal & Classification Checks")
        res_rules = {r["id"]: r for r in parsed.get("rules", [])}
        for rule in rules:
            r = res_rules.get(rule["id"], {})
            icon = "✅" if r.get("status") == "pass" else ("❌" if r.get("status") == "fail" else "⚠️")
            lines.append(f"- {icon} {rule['check']}\n  _Found: {r.get('evidence', 'Unclear')}_")
            
    return "\n".join(lines)