import re

import numpy as np
from PIL import Image

from app.config import settings

# Indian plate format, e.g. "MH 12 AB 1234"
PLATE_PATTERN = re.compile(r"^[A-Z]{2}\d{1,2}[A-Z]{1,3}\d{4}$")
_NON_ALNUM = re.compile(r"[^A-Z0-9]")


def is_valid(plate: str) -> bool:
    return bool(PLATE_PATTERN.match(plate))


class PlateReader:
    """TrOCR wrapper for license-plate text recognition (lazy-loaded).

    TrOCR is recognition-only — it expects a single, tightly cropped plate
    image and returns the full string. Detection is handled upstream by YOLO.
    """

    def __init__(self):
        self._processor = None
        self._model = None
        self._failed = False  # set once if the model can't load — don't retry every call

    def _load(self):
        if self._model is None:
            import torch
            from transformers import TrOCRProcessor, VisionEncoderDecoderModel

            self._processor = TrOCRProcessor.from_pretrained(settings.trocr_model)
            self._model = VisionEncoderDecoderModel.from_pretrained(settings.trocr_model)
            self._model.eval()
            self._torch = torch
        return self._processor, self._model

    def read(self, plate_image: np.ndarray) -> str | None:
        # OCR is the heaviest stage; on a tiny host it can OOM. A plate read is
        # best-effort metadata — never let it crash the whole analysis request.
        if not settings.ocr_enabled or self._failed:
            return None
        try:
            processor, model = self._load()

            # TrOCR wants an RGB PIL image; crops arrive grayscale (CLAHE) or BGR.
            if plate_image.ndim == 2:
                pil = Image.fromarray(plate_image).convert("RGB")
            else:
                pil = Image.fromarray(plate_image[..., ::-1]).convert("RGB")  # BGR→RGB

            pixel_values = processor(images=pil, return_tensors="pt").pixel_values
            with self._torch.no_grad():
                generated = model.generate(pixel_values, max_new_tokens=16)
            raw = processor.batch_decode(generated, skip_special_tokens=True)[0]

            text = _NON_ALNUM.sub("", raw.upper())
            if text.startswith("IND"):  # the embossed "IND" tag, not part of the number
                text = text[3:]
            return text if len(text) >= 4 else None
        except Exception as exc:  # OOM, download failure, decode error, etc.
            self._failed = True
            print(f"[ocr] plate reading disabled after failure: {exc}")
            return None


plate_reader = PlateReader()