from __future__ import annotations

import json
import re
import threading
from dataclasses import dataclass
from typing import Any

from PIL import Image

from .config import SETTINGS
from .normalization import enrich_characteristic, records_from_ocr_text
from .schema import Category, Characteristic, ReviewStatus


@dataclass
class BackendResponse:
    characteristics: list[Characteristic]
    drawing_metadata: dict[str, Any]
    general_notes: list[str]
    raw_text: str


class QwenBackend:
    def __init__(self, model_id: str | None = None) -> None:
        self.model_id = model_id or SETTINGS.qwen_model_id
        self._model: Any = None
        self._processor: Any = None
        self._lock = threading.Lock()

    def _load(self) -> None:
        if self._model is not None:
            return
        with self._lock:
            if self._model is not None:
                return
            try:
                import torch
                from transformers import AutoModelForImageTextToText, AutoProcessor
            except ImportError as exc:
                raise RuntimeError(
                    "Qwen dependencies are unavailable. Install torch and transformers "
                    "or use Demo mode."
                ) from exc
            device = "cuda" if torch.cuda.is_available() else "cpu"
            if device == "cuda":
                capability = torch.cuda.get_device_capability()
                dtype = torch.bfloat16 if capability[0] >= 8 else torch.float16
            else:
                dtype = torch.float32
            kwargs: dict[str, Any] = {
                "torch_dtype": dtype,
                "low_cpu_mem_usage": True,
                "token": SETTINGS.hf_token,
            }
            if device == "cuda":
                kwargs["device_map"] = "auto"
                kwargs["attn_implementation"] = "sdpa"
            self._processor = AutoProcessor.from_pretrained(
                self.model_id, token=SETTINGS.hf_token
            )
            self._model = AutoModelForImageTextToText.from_pretrained(
                self.model_id, **kwargs
            )
            if device == "cpu":
                self._model.to(device)
            self._model.eval()

    def extract(self, image: Image.Image, prompt: str, page: int) -> BackendResponse:
        self._load()
        import torch

        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image.convert("RGB")},
                    {"type": "text", "text": prompt},
                ],
            }
        ]
        inputs = self._processor.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_dict=True,
            return_tensors="pt",
        )
        inputs = inputs.to(self._model.device)
        with torch.inference_mode():
            generated = self._model.generate(
                **inputs,
                max_new_tokens=SETTINGS.max_new_tokens,
                do_sample=False,
                use_cache=True,
            )
        prompt_length = inputs["input_ids"].shape[-1]
        text = self._processor.decode(
            generated[0][prompt_length:], skip_special_tokens=True
        ).strip()
        payload = extract_json_payload(text)
        return response_from_payload(payload, text, page)


class PaddleBackend:
    """Experimental OCR baseline; it intentionally does not claim semantic accuracy."""

    def __init__(self, model_id: str | None = None) -> None:
        self.model_id = model_id or SETTINGS.paddle_model_id
        self._model: Any = None
        self._processor: Any = None
        self._lock = threading.Lock()

    def _load(self) -> None:
        if self._model is not None:
            return
        with self._lock:
            if self._model is not None:
                return
            try:
                import torch
                from transformers import AutoModelForImageTextToText, AutoProcessor
            except ImportError as exc:
                raise RuntimeError(
                    "PaddleOCR-VL dependencies are unavailable. Install torch and "
                    "transformers>=5 or use Demo mode."
                ) from exc
            device = "cuda" if torch.cuda.is_available() else "cpu"
            if device == "cuda":
                capability = torch.cuda.get_device_capability()
                dtype = torch.bfloat16 if capability[0] >= 8 else torch.float16
            else:
                dtype = torch.float32
            self._processor = AutoProcessor.from_pretrained(
                self.model_id, token=SETTINGS.hf_token, trust_remote_code=True
            )
            kwargs: dict[str, Any] = {
                "torch_dtype": dtype,
                "low_cpu_mem_usage": True,
                "token": SETTINGS.hf_token,
                "trust_remote_code": True,
            }
            if device == "cuda":
                kwargs["device_map"] = "auto"
            self._model = AutoModelForImageTextToText.from_pretrained(
                self.model_id, **kwargs
            )
            if device == "cpu":
                self._model.to(device)
            self._model.eval()

    def recognize(self, image: Image.Image, page: int) -> BackendResponse:
        self._load()
        import torch

        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image.convert("RGB")},
                    {"type": "text", "text": "OCR:"},
                ],
            }
        ]
        inputs = self._processor.apply_chat_template(
            messages,
            add_generation_prompt=True,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
        ).to(self._model.device)
        with torch.inference_mode():
            output = self._model.generate(
                **inputs,
                max_new_tokens=SETTINGS.max_new_tokens,
                do_sample=False,
            )
        prompt_length = inputs["input_ids"].shape[-1]
        text = self._processor.decode(
            output[0][prompt_length:], skip_special_tokens=True
        ).strip()
        return BackendResponse(records_from_ocr_text(text, page), {}, [], text)


def extract_json_payload(text: str) -> dict[str, Any]:
    cleaned = text.strip()
    cleaned = re.sub(r"^```(?:json)?\s*", "", cleaned, flags=re.I)
    cleaned = re.sub(r"\s*```$", "", cleaned)
    decoder = json.JSONDecoder()
    starts = [index for index in (cleaned.find("{"), cleaned.find("[")) if index >= 0]
    if not starts:
        raise ValueError("The model response did not contain JSON.")
    start = min(starts)
    try:
        payload, _ = decoder.raw_decode(cleaned[start:])
    except json.JSONDecodeError as exc:
        raise ValueError(f"The model returned malformed JSON: {exc.msg}.") from exc
    if isinstance(payload, list):
        return {"characteristics": payload}
    if not isinstance(payload, dict):
        raise ValueError("The model JSON must be an object or a list of characteristics.")
    return payload


def response_from_payload(
    payload: dict[str, Any], raw_text: str, page: int
) -> BackendResponse:
    raw_records = payload.get("characteristics", payload.get("parameters", []))
    if not isinstance(raw_records, list):
        raw_records = []
    records: list[Characteristic] = []
    for raw in raw_records:
        if not isinstance(raw, dict):
            continue
        normalized = _normalize_record_dict(raw, page)
        try:
            records.append(enrich_characteristic(Characteristic.model_validate(normalized)))
        except Exception:
            # Preserve partially useful content instead of failing a complete drawing.
            fallback = Characteristic(
                balloon_id=str(raw.get("balloon_id", raw.get("balloon", ""))),
                page=page,
                requirement=str(raw.get("requirement", raw.get("text", ""))),
                source_text=str(raw.get("source_text", raw.get("raw_text", ""))),
                category=Category.OTHER,
                confidence=0.2,
                status=ReviewStatus.NEEDS_REVIEW,
                validation_issues=["Model record required schema recovery."],
            )
            records.append(enrich_characteristic(fallback))
    metadata = payload.get("drawing_metadata", {})
    if not isinstance(metadata, dict):
        metadata = {}
    notes = payload.get("general_notes", [])
    if isinstance(notes, str):
        notes = [notes]
    notes = [str(note).strip() for note in notes if str(note).strip()]
    return BackendResponse(records, metadata, notes, raw_text)


def _normalize_record_dict(raw: dict[str, Any], page: int) -> dict[str, Any]:
    item = dict(raw)
    aliases = {
        "balloon": "balloon_id",
        "balloon_number": "balloon_id",
        "type": "category",
        "parameter_type": "category",
        "raw_text": "source_text",
        "bbox": "source_bbox",
        "center": "balloon_center",
        "upper_tol": "upper_tolerance",
        "lower_tol": "lower_tolerance",
        "primary_datum": "datum_primary",
        "secondary_datum": "datum_secondary",
        "tertiary_datum": "datum_tertiary",
    }
    for old, new in aliases.items():
        if new not in item and old in item:
            item[new] = item[old]
    item["page"] = _integer(item.get("page"), page)
    item["category"] = _category(item.get("category", "Other"))
    item["status"] = ReviewStatus.NEEDS_REVIEW
    return item


def _category(value: Any) -> str:
    text = re.sub(r"[^a-z]+", " ", str(value).casefold()).strip()
    mappings = {
        "linear dimension": Category.LINEAR_DIMENSION,
        "dimension": Category.LINEAR_DIMENSION,
        "measure": Category.LINEAR_DIMENSION,
        "diameter": Category.DIAMETER,
        "radius": Category.RADIUS,
        "angle": Category.ANGLE,
        "gd t": Category.GD_AND_T,
        "gdt": Category.GD_AND_T,
        "geometric tolerance": Category.GD_AND_T,
        "thread": Category.THREAD,
        "surface finish": Category.SURFACE_FINISH,
        "surface roughness": Category.SURFACE_FINISH,
        "general tolerance": Category.GENERAL_TOLERANCE,
        "material": Category.MATERIAL,
        "note": Category.NOTE,
    }
    return mappings.get(text, Category.OTHER).value


def _integer(value: Any, default: int) -> int:
    try:
        return int(value)
    except (TypeError, ValueError):
        return default


class BackendRegistry:
    """Process-local model cache; Spaces normally run one worker per GPU."""

    def __init__(self) -> None:
        self._qwen: QwenBackend | None = None
        self._paddle: PaddleBackend | None = None

    def qwen(self) -> QwenBackend:
        if self._qwen is None:
            self._qwen = QwenBackend()
        return self._qwen

    def paddle(self) -> PaddleBackend:
        if self._paddle is None:
            self._paddle = PaddleBackend()
        return self._paddle


BACKENDS = BackendRegistry()