"""Template discovery and annotation resolution.

Scans ``image/`` recursively for scanned bank document templates, assigns
each a stable ``template_id`` (``bank_doc_01``, ``bank_doc_02``, ...) based
on a deterministic sort of their paths, and resolves signature-box
annotations for each using, in order of preference:

1. **Mode 2 (preferred when present):** a manually authored
   ``annotations/<template_id>.json`` file.
2. **Mode 1 (fallback):** automatic CV-based detection
   (:mod:`src.templates.box_detector`), whose result is cached to
   ``annotations/<template_id>.json`` so subsequent runs reuse it without
   re-running detection.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path

import numpy as np

from src.templates.annotation_schema import TemplateAnnotation, TemplatePageAnnotation, SignatureBoxAnnotation
from src.templates.box_detector import detect_signature_boxes
from src.utils.image_io import discover_images, load_image_rgb
from src.utils.logging_utils import get_logger

logger = get_logger(__name__)


@dataclass
class Template:
    """A discovered template document: its image path and resolved annotation."""

    template_id: str
    image_path: Path
    annotation: TemplateAnnotation


class TemplateLoader:
    """Discovers bank document templates and resolves their signature-box annotations.

    Example:
        >>> loader = TemplateLoader(image_dir="image", annotations_dir="annotations")
        >>> templates = loader.discover()
        >>> tmpl = loader.get("bank_doc_01")
    """

    def __init__(self, image_dir: "str | Path", annotations_dir: "str | Path") -> None:
        self.image_dir = Path(image_dir)
        self.annotations_dir = Path(annotations_dir)
        self.annotations_dir.mkdir(parents=True, exist_ok=True)
        self._templates: "dict[str, Template] | None" = None

    def discover(self) -> "dict[str, Template]":
        """Discover all templates under ``image_dir`` and resolve their annotations.

        Returns:
            An ordered mapping of ``template_id -> Template``, sorted by id.
        """
        if self._templates is not None:
            return self._templates

        image_paths = discover_images(self.image_dir)
        if not image_paths:
            raise FileNotFoundError(
                f"No supported template images found under '{self.image_dir}'. "
                "Expected PNG/JPG/JPEG/TIFF/BMP files."
            )

        templates: "dict[str, Template]" = {}
        for idx, image_path in enumerate(image_paths, start=1):
            template_id = f"bank_doc_{idx:02d}"
            annotation = self._resolve_annotation(template_id, image_path)
            templates[template_id] = Template(template_id, image_path, annotation)
            logger.info(
                "Registered template '%s' -> %s (%d signature box(es), auto_detected=%s)",
                template_id, image_path, sum(len(p.signature_boxes) for p in annotation.pages),
                annotation.auto_detected,
            )

        self._templates = templates
        return templates

    def get(self, template_id: str) -> Template:
        """Return a discovered template by id, raising ``KeyError`` if unknown."""
        templates = self.discover()
        if template_id not in templates:
            raise KeyError(f"Unknown template id '{template_id}'. Known: {sorted(templates)}")
        return templates[template_id]

    def list_ids(self) -> list[str]:
        """Return all discovered template ids, sorted."""
        return sorted(self.discover().keys())

    def random_template(self, rng: np.random.Generator) -> Template:
        """Pick a uniformly random template using the given seeded generator."""
        ids = self.list_ids()
        choice = ids[int(rng.integers(0, len(ids)))]
        return self.get(choice)

    def _resolve_annotation(self, template_id: str, image_path: Path) -> TemplateAnnotation:
        annotation_path = self.annotations_dir / f"{template_id}.json"
        if annotation_path.exists():
            data = json.loads(annotation_path.read_text(encoding="utf-8"))
            return TemplateAnnotation.model_validate(data)

        logger.warning(
            "No manual annotation found for '%s'; falling back to automatic CV detection.",
            template_id,
        )
        annotation = self._auto_detect_annotation(template_id, image_path)
        annotation_path.write_text(
            json.dumps(annotation.model_dump(), indent=2), encoding="utf-8"
        )
        logger.info("Cached auto-detected annotation to %s", annotation_path)
        return annotation

    def _auto_detect_annotation(self, template_id: str, image_path: Path) -> TemplateAnnotation:
        image = load_image_rgb(image_path)
        h, w = image.shape[:2]
        boxes = detect_signature_boxes(image)

        if not boxes:
            # Last-resort fallback: a generic box in the lower-right quadrant,
            # where signature fields conventionally sit on Western bank forms.
            boxes = [(w * 0.55, h * 0.80, w * 0.92, h * 0.90)]
            logger.warning(
                "'%s': no lines detected; using generic lower-right fallback box.", template_id
            )

        signature_boxes = [
            SignatureBoxAnnotation(
                id=f"sig{i + 1}", x1=x1, y1=y1, x2=x2, y2=y2, role="Signer", name="Signer",
            )
            for i, (x1, y1, x2, y2) in enumerate(boxes)
        ]
        page = TemplatePageAnnotation(page=1, page_size=(w, h), dpi=300, signature_boxes=signature_boxes)
        return TemplateAnnotation(
            id=template_id,
            source_image=str(image_path.as_posix()),
            document_type="unknown",
            pages=[page],
            auto_detected=True,
        )