"""JSON schema for per-image feature detection results."""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

FEATURE_JSON_VERSION = 1


@dataclass
class FeatureRecord:
    image_path: Path
    json_path: Path
    camera_folder: str
    feature_type: str
    success: bool
    board_size: Optional[Tuple[int, int]] = None
    square_size: Optional[float] = None
    corners: Optional[np.ndarray] = None  # Nx1x2 float32
    center: Optional[Tuple[float, float]] = None
    ellipse: Optional[Dict[str, Any]] = None
    timestamp_sec: Optional[float] = None
    pair_key: Optional[str] = None
    preprocessing: Optional[str] = None
    error: Optional[str] = None

    @property
    def is_chessboard(self) -> bool:
        return self.success and self.feature_type == "chessboard" and self.corners is not None

    @property
    def corner_count(self) -> int:
        if self.corners is None:
            return 0
        return int(self.corners.shape[0])


def corners_to_list(corners: np.ndarray) -> List[List[float]]:
    flat = corners.reshape(-1, 2)
    return [[float(x), float(y)] for x, y in flat]


def corners_from_list(data: List[List[float]]) -> np.ndarray:
    arr = np.array(data, dtype=np.float32).reshape(-1, 1, 2)
    return arr


def save_feature_json(record: FeatureRecord) -> None:
    payload: Dict[str, Any] = {
        "version": FEATURE_JSON_VERSION,
        "image": record.image_path.name,
        "camera_folder": record.camera_folder,
        "feature_type": record.feature_type,
        "success": record.success,
        "preprocessing": record.preprocessing,
        "timestamp_sec": record.timestamp_sec,
        "pair_key": record.pair_key,
    }
    if record.board_size is not None:
        payload["board_size"] = [int(record.board_size[0]), int(record.board_size[1])]
    if record.square_size is not None:
        payload["square_size"] = float(record.square_size)
    if record.corners is not None:
        payload["corners"] = corners_to_list(record.corners)
    if record.center is not None:
        payload["center"] = [float(record.center[0]), float(record.center[1])]
    if record.ellipse is not None:
        payload["ellipse"] = record.ellipse
    if record.error:
        payload["error"] = record.error

    record.json_path.parent.mkdir(parents=True, exist_ok=True)
    with open(record.json_path, "w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2)


def load_feature_json(json_path: Path, image_path: Optional[Path] = None) -> FeatureRecord:
    with open(json_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    if image_path is not None:
        img = Path(image_path)
    else:
        stem = json_path.stem
        parent = json_path.parent
        img = parent / data.get("image", stem)
        if not img.exists():
            for ext in (".bmp", ".png", ".jpg", ".jpeg"):
                candidate = parent / f"{stem}{ext}"
                if candidate.exists():
                    img = candidate
                    break

    board_size = None
    if "board_size" in data and data["board_size"]:
        board_size = (int(data["board_size"][0]), int(data["board_size"][1]))

    corners = None
    if data.get("corners"):
        corners = corners_from_list(data["corners"])

    center = None
    if data.get("center"):
        center = (float(data["center"][0]), float(data["center"][1]))

    return FeatureRecord(
        image_path=Path(img),
        json_path=Path(json_path),
        camera_folder=data.get("camera_folder", ""),
        feature_type=data.get("feature_type", "unknown"),
        success=bool(data.get("success", False)),
        board_size=board_size,
        square_size=data.get("square_size"),
        corners=corners,
        center=center,
        ellipse=data.get("ellipse"),
        timestamp_sec=data.get("timestamp_sec"),
        pair_key=data.get("pair_key"),
        preprocessing=data.get("preprocessing"),
        error=data.get("error"),
    )


def load_folder_features(camera_dir: Path) -> List[FeatureRecord]:
    records = []
    for json_path in sorted(camera_dir.glob("*.json")):
        try:
            records.append(load_feature_json(json_path))
        except (json.JSONDecodeError, OSError) as exc:
            print(f"[WARN] Skipping invalid JSON {json_path}: {exc}")
    return records