#!/usr/bin/env python3
"""Recover camera4.json for downloaded standard-path issue cases."""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable, Optional


TARGET_CALIB_REL = Path("test_data") / "calibs" / "camera4.json"
CALIB_ATTACHMENT_NAMES = (
    "sigmastar.1/calibs/camera4.json",
    "test_data/calibs/camera4.json",
    "calibs/camera4.json",
)
CAMERA_CONFIG_FILE_NAME = "camera_config_folder.bin"
PREFERRED_CAMERA_IDS = (4, 0)
REQUIRED_FLAT_CALIB_KEYS = ("focal_u", "focal_v", "cu", "cv")


@dataclass(frozen=True)
class ExtractedCalibSource:
    payload: dict[str, Any]
    source_kind: str
    source_path: Path
    detail: str


@dataclass(frozen=True)
class CalibRecoveryResult:
    status: str
    detail: str
    target_path: Path
    source_path: Path | None = None
    source_kind: str | None = None


def _safe_int(value: Any) -> int | None:
    try:
        if value is None:
            return None
        return int(value)
    except (TypeError, ValueError):
        return None


def _sorted_unique_paths(paths: Iterable[Path]) -> list[Path]:
    unique = {path.resolve() if path.exists() else path for path in paths}
    return sorted(unique, key=lambda path: (len(path.parts), str(path)))


def _camera_attachment_matches(name: str) -> bool:
    normalized = str(name).strip().replace("\\", "/")
    if normalized in CALIB_ATTACHMENT_NAMES:
        return True
    return normalized.endswith("/camera4.json") or normalized == "camera4.json"


def iter_concatenated_json_objects(text: str) -> list[dict[str, Any]]:
    decoder = json.JSONDecoder()
    pos = 0
    payloads: list[dict[str, Any]] = []
    text_len = len(text)

    while pos < text_len:
        while pos < text_len and (text[pos].isspace() or text[pos] == "\x00"):
            pos += 1
        if pos >= text_len:
            break

        obj, end = decoder.raw_decode(text, pos)
        if isinstance(obj, dict):
            payloads.append(obj)
        pos = end

    return payloads


def _coerce_float_list(value: Any, *, min_len: int = 0) -> list[float]:
    if isinstance(value, dict):
        ordered = [value.get("x"), value.get("y"), value.get("z")]
        result = [float(item) for item in ordered if item is not None]
    elif isinstance(value, (list, tuple)):
        result = [float(item) for item in value]
    else:
        result = []
    while len(result) < min_len:
        result.append(0.0)
    return result


def _is_valid_flat_calib_payload(payload: dict[str, Any]) -> bool:
    return all(payload.get(key) is not None for key in REQUIRED_FLAT_CALIB_KEYS)


def _build_flat_camera4_payload(record: dict[str, Any], source_path: Path) -> dict[str, Any]:
    payload = dict(record)
    payload["focal_u"] = float(record["focal_u"])
    payload["focal_v"] = float(record["focal_v"])
    payload["cu"] = float(record["cu"])
    payload["cv"] = float(record["cv"])
    payload["roll"] = float(record.get("roll", 0.0))
    payload["pitch"] = float(record.get("pitch", 0.0))
    payload["yaw"] = float(record.get("yaw", 0.0))
    payload["distort_coeffs"] = _coerce_float_list(record.get("distort_coeffs"))
    payload["pos"] = _coerce_float_list(record.get("pos"), min_len=3)[:3]
    if payload.get("image_width") is not None:
        payload["image_width"] = int(payload["image_width"])
    if payload.get("image_height") is not None:
        payload["image_height"] = int(payload["image_height"])
    payload.setdefault("camera_id", 4)
    payload["recovered_from"] = CAMERA_CONFIG_FILE_NAME
    payload["recovered_from_path"] = str(source_path)
    return payload


def extract_calib_from_camera_config_folder(
    config_path: Path,
    preferred_camera_ids: tuple[int, ...] = PREFERRED_CAMERA_IDS,
) -> ExtractedCalibSource | None:
    text = config_path.read_text(encoding="utf-8")
    records = iter_concatenated_json_objects(text)
    valid_records = [record for record in records if _is_valid_flat_calib_payload(record)]
    if not valid_records:
        return None

    selected_records = valid_records
    available_camera_ids = {
        _safe_int(record.get("camera_id"))
        for record in valid_records
        if _safe_int(record.get("camera_id")) is not None
    }
    chosen_camera_id: int | None = None

    for camera_id in preferred_camera_ids:
        preferred_records = [
            record for record in valid_records if _safe_int(record.get("camera_id")) == camera_id
        ]
        if preferred_records:
            selected_records = preferred_records
            chosen_camera_id = camera_id
            break

    if chosen_camera_id is None:
        if len(available_camera_ids) == 1:
            chosen_camera_id = next(iter(available_camera_ids))
        else:
            counts: dict[int, int] = {}
            for record in valid_records:
                camera_id = _safe_int(record.get("camera_id"))
                if camera_id is None:
                    continue
                counts[camera_id] = counts.get(camera_id, 0) + 1
            if counts:
                chosen_camera_id = max(counts.items(), key=lambda item: item[1])[0]
                selected_records = [
                    record for record in valid_records if _safe_int(record.get("camera_id")) == chosen_camera_id
                ]

    best_record = max(
        enumerate(selected_records),
        key=lambda item: (_safe_int(item[1].get("utc_tick")) or -1, item[0]),
    )[1]
    payload = _build_flat_camera4_payload(best_record, config_path)
    detail = (
        f"camera_id={chosen_camera_id if chosen_camera_id is not None else 'unknown'} "
        f"records={len(valid_records)} latest_utc_tick={best_record.get('utc_tick')}"
    )
    return ExtractedCalibSource(
        payload=payload,
        source_kind="camera_config_folder",
        source_path=config_path,
        detail=detail,
    )


def _extract_calib_from_mcap_with_clip_reader(mcap_path: Path) -> ExtractedCalibSource | None:
    try:
        from pdcl_pyclip.reader import ClipReader
    except ImportError:
        return None

    reader = ClipReader(str(mcap_path))
    for attachment in reader.iter_attachments():
        if _camera_attachment_matches(attachment.name):
            payload = json.loads(attachment.data.decode("utf-8"))
            return ExtractedCalibSource(
                payload=payload,
                source_kind="mcap_attachment",
                source_path=mcap_path,
                detail=f"attachment={attachment.name}",
            )
    return None


def _extract_calib_from_mcap_with_generic_reader(mcap_path: Path) -> ExtractedCalibSource | None:
    try:
        from mcap.reader import make_reader
    except ImportError:
        return None

    with mcap_path.open("rb") as file:
        reader = make_reader(file)
        for attachment in reader.iter_attachments():
            if _camera_attachment_matches(attachment.name):
                payload = json.loads(attachment.data.decode("utf-8"))
                return ExtractedCalibSource(
                    payload=payload,
                    source_kind="mcap_attachment",
                    source_path=mcap_path,
                    detail=f"attachment={attachment.name}",
                )
    return None


def extract_calib_from_mcap_attachment(mcap_path: Path) -> ExtractedCalibSource | None:
    errors: list[str] = []
    for extractor in (_extract_calib_from_mcap_with_clip_reader, _extract_calib_from_mcap_with_generic_reader):
        try:
            extracted = extractor(mcap_path)
        except Exception as exc:  # pragma: no cover - recovery fallback logging
            errors.append(f"{extractor.__name__}: {type(exc).__name__}: {exc}")
            continue
        if extracted is not None:
            return extracted

    if errors:
        raise RuntimeError("; ".join(errors))
    return None


def _find_candidate_camera_config_paths(source_root: Path) -> list[Path]:
    return _sorted_unique_paths(source_root.rglob(CAMERA_CONFIG_FILE_NAME))


def _find_candidate_mcap_paths(source_root: Path) -> list[Path]:
    return _sorted_unique_paths(source_root.rglob("*.mcap"))


def _write_camera4_json(target_path: Path, payload: dict[str, Any]) -> None:
    target_path.parent.mkdir(parents=True, exist_ok=True)
    target_path.write_text(
        json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )


def recover_camera4_json(
    source_root: Path,
    target_root: Path,
    dry_run: bool = False,
) -> CalibRecoveryResult:
    target_path = target_root / TARGET_CALIB_REL
    if target_path.is_file():
        return CalibRecoveryResult(
            status="skipped_existing_calib",
            detail=f"target already exists: {target_path}",
            target_path=target_path,
        )

    camera_config_paths = _find_candidate_camera_config_paths(source_root)
    mcap_paths = _find_candidate_mcap_paths(source_root)
    errors: list[str] = []

    for config_path in camera_config_paths:
        try:
            extracted = extract_calib_from_camera_config_folder(config_path)
        except Exception as exc:
            errors.append(f"{config_path}: {type(exc).__name__}: {exc}")
            continue
        if extracted is None:
            continue
        if dry_run:
            return CalibRecoveryResult(
                status="planned_calib_recovery_from_camera_config_folder",
                detail=f"would write {target_path} from {config_path} ({extracted.detail})",
                target_path=target_path,
                source_path=config_path,
                source_kind=extracted.source_kind,
            )
        _write_camera4_json(target_path, extracted.payload)
        return CalibRecoveryResult(
            status="recovered_calib_from_camera_config_folder",
            detail=f"wrote {target_path} from {config_path} ({extracted.detail})",
            target_path=target_path,
            source_path=config_path,
            source_kind=extracted.source_kind,
        )

    for mcap_path in mcap_paths:
        try:
            extracted = extract_calib_from_mcap_attachment(mcap_path)
        except Exception as exc:
            errors.append(f"{mcap_path}: {type(exc).__name__}: {exc}")
            continue
        if extracted is None:
            continue
        if dry_run:
            return CalibRecoveryResult(
                status="planned_calib_recovery_from_mcap_attachment",
                detail=f"would write {target_path} from {mcap_path} ({extracted.detail})",
                target_path=target_path,
                source_path=mcap_path,
                source_kind=extracted.source_kind,
            )
        _write_camera4_json(target_path, extracted.payload)
        return CalibRecoveryResult(
            status="recovered_calib_from_mcap_attachment",
            detail=f"wrote {target_path} from {mcap_path} ({extracted.detail})",
            target_path=target_path,
            source_path=mcap_path,
            source_kind=extracted.source_kind,
        )

    checked_sources = (
        f"camera_config_folder.bin={len(camera_config_paths)} mcap={len(mcap_paths)}"
    )
    if errors:
        checked_sources += f"; errors: {' | '.join(errors[:3])}"
    return CalibRecoveryResult(
        status="failed_calib_recovery",
        detail=f"no recoverable calibration found under {source_root} ({checked_sources})",
        target_path=target_path,
    )