from __future__ import annotations

import csv
import json
from dataclasses import dataclass
from datetime import datetime, timezone
from math import degrees, radians
from pathlib import Path
from typing import Any, Dict, List, Tuple

import mujoco
import numpy as np
from PIL import Image

from .mujoco_robot import MujocoRobot
from .realman_reference import RealManFkReference
from .stage2_validation import Stage2Settings, Stage2Validator
from .teleop_config import load_dual_arm_profiles
from .teleop_control import (
    ControllerSample,
    DualArmQpTeleopController,
    SafetyState,
)


@dataclass(frozen=True)
class Stage3Settings:
    seed: int = 20260630
    fk_samples: int = 2_000
    ik_samples: int = 200
    trajectories: int = 5
    trajectory_points: int = 200
    render_width: int = 1280
    render_height: int = 720

    @classmethod
    def quick(cls, seed: int = 20260630) -> "Stage3Settings":
        return cls(
            seed=seed,
            fk_samples=100,
            ik_samples=30,
            trajectories=2,
            trajectory_points=25,
            render_width=640,
            render_height=360,
        )

    def stage2_settings(self) -> Stage2Settings:
        return Stage2Settings(
            seed=self.seed,
            fk_samples=self.fk_samples,
            ik_samples=self.ik_samples,
            trajectories=self.trajectories,
            trajectory_points=self.trajectory_points,
            render_width=self.render_width,
            render_height=self.render_height,
        )


class Stage3Validator:
    def __init__(
        self,
        reference: RealManFkReference,
        teleop_config_path: Path | str,
        peripheral_config_path: Path | str,
        settings: Stage3Settings = Stage3Settings(),
    ) -> None:
        self.reference = reference
        self.settings = settings
        self.profiles = load_dual_arm_profiles(
            teleop_config_path, peripheral_config_path
        )
        self.checks: Dict[str, Dict[str, Any]] = {}
        self.failures: List[Dict[str, Any]] = []
        self.images: Dict[str, np.ndarray] = {}

    def run(self) -> Dict[str, Any]:
        right_validator = Stage2Validator(
            self.reference,
            controlled_arm="right",
            settings=self.settings.stage2_settings(),
        )
        right_summary = right_validator.run()
        self.failures.extend(right_validator.failures)
        for label, (rgb, segmentation) in right_validator.images.items():
            self.images[f"right_{label}_rgb"] = rgb
            self.images[f"right_{label}_segmentation"] = segmentation
        self.checks["right_arm_validation"] = {
            "passed": right_summary["passed"],
            "required": True,
            "checks": right_summary["checks"],
        }

        robot = MujocoRobot(self.profiles)
        try:
            self._initial_tcp_check(robot)
            self._dual_grip_control_check(robot)
            image = robot.render(
                self.settings.render_width, self.settings.render_height
            )
            self.images["dual_control_final_rgb"] = image
            self.checks["dual_control_render"] = {
                "passed": float(image.std()) > 5.0,
                "required": True,
                "rgb_std": float(image.std()),
            }
        finally:
            robot.close()

        passed = all(check["passed"] for check in self.checks.values())
        return {
            "schema_version": 1,
            "generated_at": datetime.now(timezone.utc).isoformat(),
            "seed": self.settings.seed,
            "mujoco_version": mujoco.__version__,
            "realman_api_version": self.reference.api_version,
            "passed": passed,
            "checks": self.checks,
            "failure_count": len(self.failures),
        }

    def _initial_tcp_check(self, robot: MujocoRobot) -> None:
        metrics: Dict[str, Dict[str, Any]] = {}
        passed = True
        for arm in ("left", "right"):
            diagnostic = robot.initial_pose_diagnostics[arm]
            arm_passed = (
                diagnostic.solved_position_error_m <= 1e-3
                and diagnostic.solved_orientation_error_rad <= radians(0.1)
            )
            passed &= arm_passed
            metrics[arm] = {
                "active_tool": self.profiles[arm].active_tool_name,
                "solved_q_deg": np.rad2deg(diagnostic.solved_q_rad).tolist(),
                "solved_position_error_m": diagnostic.solved_position_error_m,
                "solved_orientation_error_deg": degrees(
                    diagnostic.solved_orientation_error_rad
                ),
                "configured_joint_position_error_m": (
                    diagnostic.configured_joint_position_error_m
                ),
                "configured_joint_orientation_error_deg": degrees(
                    diagnostic.configured_joint_orientation_error_rad
                ),
            }
            if not arm_passed:
                self._record_failure(
                    "initial_tcp",
                    0 if arm == "left" else 1,
                    f"{arm} initial TCP solve exceeded tolerance",
                    diagnostic.solved_q_rad,
                    diagnostic.solved_position_error_m,
                    diagnostic.solved_orientation_error_rad,
                )
        self.checks["initial_tcp_pose"] = {
            "passed": passed,
            "required": True,
            "arms": metrics,
            "note": "configured initial_joint_pose mismatch is diagnostic only",
        }

    def _dual_grip_control_check(self, robot: MujocoRobot) -> None:
        controller = DualArmQpTeleopController(robot, self.profiles)
        initial = robot.read_joint_positions().positions_rad
        identity = np.array([0.0, 0.0, 0.0, 1.0])

        def sample(arm: str, grip: bool, position) -> ControllerSample:
            return ControllerSample(
                hand=arm,
                grip=grip,
                trigger=0.0,
                position_m=np.asarray(position, dtype=float),
                quaternion_xyzw=identity,
            )

        controller.update_sample(sample("left", False, [0, 0, 0]), 0.0)
        controller.update_sample(sample("right", False, [0, 0, 0]), 0.0)
        idle = controller.step(0.0)
        controller.update_sample(sample("left", True, [0, 0, 0]), 0.01)
        controller.update_sample(sample("right", True, [0, 0, 0]), 0.01)
        engaged = controller.step(0.01)
        after_engage = robot.read_joint_positions().positions_rad
        no_engage_jump = all(
            np.max(np.abs(after_engage[arm] - initial[arm])) < 1e-10
            for arm in ("left", "right")
        )

        last_result = engaged
        for index in range(1, 11):
            now = 0.01 + index / 90.0
            hand_delta = 0.001 * index
            controller.update_sample(
                sample("left", True, [0, hand_delta, 0]), now
            )
            controller.update_sample(
                sample("right", True, [0, hand_delta, 0]), now
            )
            last_result = controller.step(now)
            if last_result.state is SafetyState.FAULT:
                break
        moved = robot.read_joint_positions().positions_rad
        joint_delta = {
            arm: float(np.max(np.abs(moved[arm] - initial[arm])))
            for arm in ("left", "right")
        }
        both_moved = all(value > 1e-5 for value in joint_delta.values())

        release_time = 0.15
        controller.update_sample(sample("left", False, [0, 0.01, 0]), release_time)
        controller.update_sample(sample("right", False, [0, 0.01, 0]), release_time)
        released = controller.step(release_time)
        passed = (
            idle.state is SafetyState.IDLE
            and engaged.state is SafetyState.ACTIVE
            and no_engage_jump
            and last_result.state is SafetyState.ACTIVE
            and both_moved
            and released.state is SafetyState.IDLE
        )
        if not passed:
            self._record_failure(
                "dual_grip_control",
                0,
                f"idle={idle.state.value}, engaged={engaged.state.value}, "
                f"last={last_result.state.value}, released={released.state.value}",
            )
        self.checks["dual_grip_control"] = {
            "passed": passed,
            "required": True,
            "no_engage_jump": no_engage_jump,
            "max_joint_delta_rad": joint_delta,
            "final_state": released.state.value,
        }

    def _record_failure(
        self,
        category: str,
        index: int,
        reason: str,
        q: np.ndarray | None = None,
        position_error_m: float = float("nan"),
        orientation_error_rad: float = float("nan"),
    ) -> None:
        self.failures.append(
            {
                "category": category,
                "sample": index,
                "reason": reason,
                "position_error_m": position_error_m,
                "orientation_error_rad": orientation_error_rad,
                "q_rad": json.dumps(q.tolist()) if q is not None else "",
            }
        )


def write_stage3_report(
    output_dir: Path | str,
    summary: Dict[str, Any],
    failures: List[Dict[str, Any]],
    images: Dict[str, np.ndarray],
) -> Tuple[Path, Path, Path, List[Path]]:
    directory = Path(output_dir)
    directory.mkdir(parents=True, exist_ok=True)
    json_path = directory / "stage3_summary.json"
    csv_path = directory / "stage3_failures.csv"
    markdown_path = directory / "stage3_report.md"
    json_path.write_text(
        json.dumps(summary, ensure_ascii=True, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    fields = [
        "category",
        "sample",
        "reason",
        "position_error_m",
        "orientation_error_rad",
        "q_rad",
    ]
    with csv_path.open("w", encoding="utf-8", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=fields)
        writer.writeheader()
        writer.writerows(failures)
    image_paths = []
    for label, image in images.items():
        path = directory / f"stage3_{label}.png"
        Image.fromarray(image).save(path)
        image_paths.append(path)
    lines = [
        "# RM75-B Stage 3 Dual-Arm QP Validation",
        "",
        f"- Overall: **{'PASS' if summary['passed'] else 'FAIL'}**",
        f"- Seed: `{summary['seed']}`",
        f"- MuJoCo: `{summary['mujoco_version']}`",
        f"- RealMan API: `{summary['realman_api_version']}`",
        f"- Failures recorded: `{summary['failure_count']}`",
        "",
        "| Check | Result | Metrics |",
        "|---|---:|---|",
    ]
    for name, check in summary["checks"].items():
        metrics = {
            key: value
            for key, value in check.items()
            if key not in {"passed", "required"}
        }
        lines.append(
            f"| `{name}` | {'PASS' if check['passed'] else 'FAIL'} | "
            f"`{json.dumps(metrics, ensure_ascii=True, sort_keys=True)}` |"
        )
    markdown_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
    return json_path, csv_path, markdown_path, image_paths
