import os
import sys
import re
import spaces
import torch
from unittest.mock import patch
from transformers import AutoModel, AutoTokenizer
from transformers.dynamic_module_utils import get_imports as _original_get_imports

# Capture original BEFORE patching — importing inside the patched fn causes
# infinite recursion because it would resolve to the patched version.
def _strip_flash_attn(filename, *args, **kwargs):
    return [i for i in _original_get_imports(filename, *args, **kwargs) if i != "flash_attn"]

# Block flash_attn at runtime — HF ZeroGPU has it installed, so without this
# the model's own "try: import flash_attn" succeeds and forces FA2 on.
sys.modules["flash_attn"] = None  # None = ImportError on any "import flash_attn"

# ---------------------------------------------------------------------------
# Global Initialization — Nemotron Parse (Nvidia Spatial Expert)
# Specializes in dense printed documents: tables, grids, bounding boxes.
# Loaded into CPU memory on startup. GPU is claimed inside @spaces.GPU only.
# ---------------------------------------------------------------------------
print("Initializing Nemotron Parse (Nvidia)... This may take a moment.")

parser_model_id = "nvidia/NVIDIA-Nemotron-Parse-v1.2"  # Verified correct ID
hf_token = os.environ.get("HF_TOKEN")

def _patch_flash_attn_imports(filename, *args, **kwargs):
    from transformers.dynamic_module_utils import get_imports as _orig
    return [imp for imp in _orig(filename, *args, **kwargs) if imp != "flash_attn"]

try:
    with patch(
        "transformers.dynamic_module_utils.get_imports",
        side_effect=_strip_flash_attn
    ):
        parser_tokenizer = AutoTokenizer.from_pretrained(
            parser_model_id,
            trust_remote_code=True,
            token=hf_token
        )
        parser_model = AutoModel.from_pretrained(
            parser_model_id,
            trust_remote_code=True,
            torch_dtype=torch.bfloat16,
            low_cpu_mem_usage=True,
            token=hf_token
        )
    print("Nemotron Parse loaded successfully.")
except Exception as e:
    print(f"Warning: Nemotron Parse failed to load: {e}")
    parser_tokenizer = None
    parser_model = None


@spaces.GPU
def run_nemotron_parse(image_path: str) -> dict:
    """
    Spatial Expert: Runs Nemotron Parse on a raw receipt image.
    Best at: dense printed bills, hospital invoices, utility statements.
    """
    if parser_model is None or parser_tokenizer is None:
        return {"source": "nemotron_parse", "text": "", "confidence": 0.0}

    try:
        from PIL import Image as PILImage
        from transformers import AutoProcessor

        parser_model.to("cuda")
        parser_model.eval()

        image = PILImage.open(image_path).convert("RGB")

        # Nemotron Parse v1.2 requires a specific 4-token prompt format:
        # <predict_bbox><predict_classes><output_markdown><predict_text_in_pic>
        # The last token controls whether text-in-image is extracted.
        prompt = "<predict_no_bbox><predict_no_classes><output_markdown><predict_text_in_pic>"

        processor = AutoProcessor.from_pretrained(
            parser_model_id, trust_remote_code=True,
            token=os.environ.get("HF_TOKEN")
        )

        inputs = processor(text=prompt, images=image, return_tensors="pt").to("cuda")

        with torch.no_grad():
            out_ids = parser_model.generate(**inputs, max_new_tokens=600)

        input_length = inputs["input_ids"].shape[-1]
        raw = processor.decode(out_ids[0][input_length:], skip_special_tokens=True)

        confidence = _score_confidence(raw)
        return {"source": "nemotron_parse", "text": raw.strip(), "confidence": confidence}

    except Exception as e:
        print(f"Nemotron Parse inference error: {e}")
        return {"source": "nemotron_parse", "text": "", "confidence": 0.0}


def _score_confidence(text: str) -> float:
    """
    Scores how structured an extracted text looks.
    Higher score = more price-like tokens = better quality output.
    This is the 'arbiter' that picks the winner in the parallel engine.
    """
    if not text or len(text.strip()) < 20:
        return 0.0

    # Count how many price-like patterns appear (e.g. $4.99 / ₹799 / 12.00)
    price_hits = len(re.findall(r"[\$₹€£¥]\s*\d+[\.,]\d{2}|\d+[\.,]\d{2}", text))

    # Count structured line markers (dashes, colons, tabs — signs of a table)
    structure_hits = len(re.findall(r"[-:|\t]", text))

    # Normalize to a 0–1 range, capped at 1.0
    score = min(1.0, (price_hits * 0.15) + (structure_hits * 0.02))
    return round(score, 3)