"""Acquire a calibration sample.

One sample = (T_base_wrist from G1) + (T_cam_marker from ArUco detection).

Pose computation: we call marker-finder only for the 2D `corners` (unaffected
by its baked D435i intrinsics) and rerun solvePnP locally with the real Nuwa
RGB intrinsics. This sidesteps what `CLAUDE.md` flags as the dominant error
source in hand-eye calibration on this stack.
"""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass
from typing import Any

import numpy as np

from ..clients import G1Client, MarkerClient, NuwaClient
from ..geometry import T_from_pos_quat, make_T
from ..perception.aruco_pose import nuwa_intrinsics_to_matrix, solve_square_marker
from ..timing import Stopwatch

log = logging.getLogger(__name__)


@dataclass
class Sample:
    created_at: float
    side: str
    marker_id: int
    T_base_wrist: list[list[float]]  # 4x4
    T_cam_marker: list[list[float]]  # 4x4
    corners: list[list[float]]
    note: str = ""
    timings_ms: dict[str, float] | None = None
    # Diagnostics from the local pose solve (absent when we fell back to
    # marker-finder's pose).
    pose_source: str = "marker_finder"          # "nuwa_intrinsics" | "marker_finder"
    reprojection_error_px: float | None = None
    ambiguity_ratio: float | None = None

    @classmethod
    def make(cls, side: str, marker_id: int, T_base_wrist: np.ndarray,
             T_cam_marker: np.ndarray, corners: np.ndarray, note: str = "",
             timings_ms: dict[str, float] | None = None,
             pose_source: str = "marker_finder",
             reprojection_error_px: float | None = None,
             ambiguity_ratio: float | None = None) -> "Sample":
        return cls(
            created_at=time.time(),
            side=side,
            marker_id=int(marker_id),
            T_base_wrist=T_base_wrist.tolist(),
            T_cam_marker=T_cam_marker.tolist(),
            corners=np.asarray(corners).tolist(),
            note=note,
            timings_ms=timings_ms,
            pose_source=pose_source,
            reprojection_error_px=reprojection_error_px,
            ambiguity_ratio=ambiguity_ratio,
        )


class SamplerError(RuntimeError):
    pass


class CalibrationSampler:
    def __init__(self, nuwa: NuwaClient, marker: MarkerClient, g1: G1Client):
        self.nuwa = nuwa
        self.marker = marker
        self.g1 = g1

    def acquire(self, marker_id: int, size_mm: float, side: str, note: str = "") -> Sample:
        sw = Stopwatch()
        # 1) wrist pose from G1
        with sw.step("g1_wrist"):
            try:
                wrist = self.g1.wrist()
            except Exception as e:
                raise SamplerError(f"cannot read wrist pose from G1: {e}") from e
        # G1 returns {'left_wrist': {...}, 'right_wrist': {...}}; tolerate {'left'/'right'} and {'data': ...}.
        body = wrist.get("data") if isinstance(wrist.get("data"), dict) else wrist
        pose = body.get(f"{side}_wrist") or body.get(side)
        if pose is None or "position" not in pose or "quaternion" not in pose:
            raise SamplerError(f"G1 /wrist response missing '{side}' pose: {wrist}")
        T_base_wrist = T_from_pos_quat(np.asarray(pose["position"]), np.asarray(pose["quaternion"]))

        # 2) ArUco corners from marker-finder (2D corners are intrinsics-free).
        with sw.step("marker_detect"):
            result = self.marker.detect_url(
                markers=[{"id": int(marker_id), "size_mm": float(size_mm)}],
                image_url=self.nuwa.rgb_url(),
            )
        found = result.get("found", [])
        if not found:
            raise SamplerError(f"marker id={marker_id} not detected (not_found={result.get('not_found')})")
        entry = next((m for m in found if int(m["id"]) == int(marker_id)), None)
        if entry is None:
            raise SamplerError(f"marker id={marker_id} missing from 'found': {found}")
        corners = np.asarray(entry["corners"], dtype=np.float64)

        # 3) Pose: solve locally with Nuwa RGB intrinsics (primary path).
        pose_source = "marker_finder"
        reproj_err = None
        ambiguity = None
        T_cam_marker = None
        try:
            with sw.step("nuwa_intrinsics"):
                intr = self.nuwa.intrinsics()
            K, dist = nuwa_intrinsics_to_matrix(intr)
            with sw.step("local_pnp"):
                sol = solve_square_marker(corners, float(size_mm) / 1000.0, K, dist)
            T_cam_marker = make_T(sol["R"], sol["t"])
            pose_source = "nuwa_intrinsics"
            reproj_err = sol["reprojection_error_px"]
            ambiguity = sol["ambiguity_ratio"]
        except Exception as e:
            log.warning("local pose failed (%s) — falling back to marker-finder pose", e)

        # Fallback: use marker-finder's baked-in pose if local solve failed.
        if T_cam_marker is None:
            R = np.asarray(entry["rotation_matrix"], dtype=np.float64).reshape(3, 3)
            t = np.asarray(entry["position_m"], dtype=np.float64).flatten()
            T_cam_marker = make_T(R, t)

        return Sample.make(
            side=side, marker_id=int(marker_id),
            T_base_wrist=T_base_wrist, T_cam_marker=T_cam_marker,
            corners=corners, note=note, timings_ms=sw.as_dict(),
            pose_source=pose_source,
            reprojection_error_px=reproj_err,
            ambiguity_ratio=ambiguity,
        )