"""OpenCV template-matching cache for sovereignty perception.

Implements "See Once, Template Forever" from the Sovereignty Loop (#953).

First encounter: VLM analyses screenshot (3-6 sec) → structured JSON.
Crystallized as: OpenCV template + bounding box → templates.json (3 ms).

The ``crystallize_perception()`` function converts VLM output into
reusable OpenCV templates, and ``PerceptionCache.match()`` retrieves
them without calling the VLM again.

Refs: #955, #953 (Section III.1 — Perception)
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import cv2
import numpy as np

logger = logging.getLogger(__name__)


@dataclass
class Template:
    """A reusable visual template extracted from VLM analysis."""

    name: str
    image: np.ndarray
    threshold: float = 0.85
    bbox: tuple[int, int, int, int] | None = None  # (x1, y1, x2, y2)
    metadata: dict[str, Any] | None = None


@dataclass
class CacheResult:
    """Result of a template match against a screenshot."""

    confidence: float
    state: Any | None


class PerceptionCache:
    """OpenCV-based visual template cache.

    Stores templates extracted from VLM responses and matches them
    against future screenshots using template matching, eliminating
    the need for repeated VLM calls on known visual patterns.
    """

    def __init__(self, templates_path: Path | str = "data/templates.json") -> None:
        self.templates_path = Path(templates_path)
        self.templates: list[Template] = []
        self.load()

    def match(self, screenshot: np.ndarray) -> CacheResult:
        """Match stored templates against a screenshot.

        Returns the highest-confidence match. If confidence exceeds
        the template's threshold, the cached state is returned.

        Parameters
        ----------
        screenshot:
            The current frame as a numpy array (BGR or grayscale).

        Returns
        -------
        CacheResult
            Confidence score and cached state (or None if no match).
        """
        best_match_confidence = 0.0
        best_match_name = None
        best_match_metadata = None

        for template in self.templates:
            if template.image.size == 0:
                continue

            try:
                # Convert to grayscale if needed for matching
                if len(screenshot.shape) == 3 and len(template.image.shape) == 2:
                    frame = cv2.cvtColor(screenshot, cv2.COLOR_BGR2GRAY)
                elif len(screenshot.shape) == 2 and len(template.image.shape) == 3:
                    frame = screenshot
                    # skip mismatched template
                    continue
                else:
                    frame = screenshot

                # Ensure template is smaller than frame
                if (
                    template.image.shape[0] > frame.shape[0]
                    or template.image.shape[1] > frame.shape[1]
                ):
                    continue

                res = cv2.matchTemplate(frame, template.image, cv2.TM_CCOEFF_NORMED)
                _, max_val, _, _ = cv2.minMaxLoc(res)

                if max_val > best_match_confidence:
                    best_match_confidence = max_val
                    best_match_name = template.name
                    best_match_metadata = template.metadata
            except cv2.error:
                logger.debug("Template match failed for '%s'", template.name)
                continue

        if best_match_confidence >= 0.85 and best_match_name is not None:
            return CacheResult(
                confidence=best_match_confidence,
                state={"template_name": best_match_name, **(best_match_metadata or {})},
            )
        return CacheResult(confidence=best_match_confidence, state=None)

    def add(self, templates: list[Template]) -> None:
        """Add new templates to the cache."""
        self.templates.extend(templates)

    def persist(self) -> None:
        """Write template metadata to disk.

        Note: actual template images are stored alongside as .npy files
        for fast loading. The JSON file stores metadata only.
        """
        self.templates_path.parent.mkdir(parents=True, exist_ok=True)

        entries = []
        for t in self.templates:
            entry: dict[str, Any] = {"name": t.name, "threshold": t.threshold}
            if t.bbox is not None:
                entry["bbox"] = list(t.bbox)
            if t.metadata:
                entry["metadata"] = t.metadata

            # Save non-empty template images as .npy
            if t.image.size > 0:
                img_path = self.templates_path.parent / f"template_{t.name}.npy"
                try:
                    np.save(str(img_path), t.image)
                    entry["image_path"] = str(img_path.name)
                except Exception as exc:
                    logger.warning("Failed to save template image for '%s': %s", t.name, exc)

            entries.append(entry)

        with self.templates_path.open("w") as f:
            json.dump(entries, f, indent=2)
        logger.debug("Persisted %d templates to %s", len(entries), self.templates_path)

    def load(self) -> None:
        """Load templates from disk."""
        if not self.templates_path.exists():
            return

        try:
            with self.templates_path.open("r") as f:
                templates_data = json.load(f)
        except (json.JSONDecodeError, OSError) as exc:
            logger.warning("Failed to load templates: %s", exc)
            return

        self.templates = []
        for t in templates_data:
            # Try to load the image from .npy if available
            image = np.array([])
            image_path = t.get("image_path")
            if image_path:
                full_path = self.templates_path.parent / image_path
                if full_path.exists():
                    try:
                        image = np.load(str(full_path))
                    except Exception:
                        pass

            bbox = tuple(t["bbox"]) if "bbox" in t else None

            self.templates.append(
                Template(
                    name=t["name"],
                    image=image,
                    threshold=t.get("threshold", 0.85),
                    bbox=bbox,
                    metadata=t.get("metadata"),
                )
            )

    def clear(self) -> None:
        """Remove all templates."""
        self.templates.clear()

    def __len__(self) -> int:
        return len(self.templates)


def crystallize_perception(
    screenshot: np.ndarray,
    vlm_response: Any,
) -> list[Template]:
    """Extract reusable OpenCV templates from a VLM response.

    Converts VLM-identified UI elements into cropped template images
    that can be matched in future frames without calling the VLM.

    Parameters
    ----------
    screenshot:
        The full screenshot that was analysed by the VLM.
    vlm_response:
        Structured VLM output. Expected formats:
        - dict with ``"items"`` list, each having ``"name"`` and ``"bounding_box"``
        - dict with ``"elements"`` list (same structure)
        - list of dicts with ``"name"`` and ``"bbox"`` or ``"bounding_box"``

    Returns
    -------
    list[Template]
        Extracted templates ready to be added to a PerceptionCache.
    """
    templates: list[Template] = []

    # Normalize the response format
    items: list[dict[str, Any]] = []
    if isinstance(vlm_response, dict):
        items = vlm_response.get("items", vlm_response.get("elements", []))
    elif isinstance(vlm_response, list):
        items = vlm_response

    for item in items:
        name = item.get("name") or item.get("label") or item.get("type")
        bbox = item.get("bounding_box") or item.get("bbox")

        if not name or not bbox:
            continue

        try:
            if len(bbox) == 4:
                x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])
            else:
                continue

            # Validate bounds
            h, w = screenshot.shape[:2]
            x1 = max(0, min(x1, w - 1))
            y1 = max(0, min(y1, h - 1))
            x2 = max(x1 + 1, min(x2, w))
            y2 = max(y1 + 1, min(y2, h))

            template_image = screenshot[y1:y2, x1:x2].copy()

            if template_image.size == 0:
                continue

            metadata = {
                k: v for k, v in item.items() if k not in ("name", "label", "bounding_box", "bbox")
            }

            templates.append(
                Template(
                    name=name,
                    image=template_image,
                    bbox=(x1, y1, x2, y2),
                    metadata=metadata if metadata else None,
                )
            )
            logger.debug(
                "Crystallized perception template '%s' (%dx%d)",
                name,
                x2 - x1,
                y2 - y1,
            )

        except (ValueError, IndexError, TypeError) as exc:
            logger.debug("Failed to crystallize item '%s': %s", name, exc)
            continue

    if templates:
        logger.info(
            "Crystallized %d perception template(s) from VLM response",
            len(templates),
        )

    return templates
