"""Dual-camera and Newton validation for a taught Lite6 TCP workspace.

This module is deliberately an assessment boundary.  It never imports the
xArm SDK, never changes controller state and never returns a motion permit.
The operator-recorded workspace proposal remains the source of the TCP box;
the calibrated dual-depth preflight and Newton Physics only validate evidence
bound to that exact proposal.
"""

from __future__ import annotations

import hashlib
import importlib
import json
import math
import re
import stat
import threading
from pathlib import Path
from typing import Any, Callable, Mapping, Optional, Tuple

_HEX_32 = re.compile(r"[0-9a-f]{32}\Z")
_HEX_64 = re.compile(r"[0-9a-f]{64}\Z")
_VERSION = re.compile(r"[0-9]+\.[0-9]+\.[0-9]+\Z")
_MAX_MODEL_BYTES = 512 * 1024
_MAX_TCP_DEVIATION_MM = 2.0

WORKSPACE_BOUNDARY_BLOCKERS = (
    "proposal_invalid",
    "dual_depth_unavailable",
    "dual_depth_blocked",
    "newton_unavailable",
    "newton_mismatch",
    "stopped",
    "result_invalid",
)


class NewtonWorkspaceError(RuntimeError):
    """A pinned Newton model or validation result is unavailable/invalid."""


def _canonical_bytes(value: Mapping[str, Any]) -> bytes:
    try:
        return json.dumps(
            value,
            allow_nan=False,
            ensure_ascii=True,
            separators=(",", ":"),
            sort_keys=True,
        ).encode("utf-8")
    except (TypeError, ValueError, OverflowError) as exc:
        raise NewtonWorkspaceError("workspace validation is not canonical") from exc


def _sha256(value: Any, name: str) -> str:
    if not isinstance(value, str) or _HEX_64.fullmatch(value) is None:
        raise NewtonWorkspaceError("{} SHA-256 is invalid".format(name))
    return value


def _finite_vector(value: Any, length: int, name: str) -> Tuple[float, ...]:
    if not isinstance(value, (tuple, list)) or len(value) != length:
        raise NewtonWorkspaceError(
            "{} must contain exactly {} finite values".format(name, length)
        )
    result = []
    for item in value:
        if isinstance(item, bool):
            raise NewtonWorkspaceError("{} contains a boolean".format(name))
        try:
            number = float(item)
        except (TypeError, ValueError, OverflowError) as exc:
            raise NewtonWorkspaceError("{} contains a non-number".format(name)) from exc
        if not math.isfinite(number):
            raise NewtonWorkspaceError("{} contains a non-finite number".format(name))
        result.append(number)
    return tuple(result)


def _proposal_input(proposal: Any) -> Mapping[str, Any]:
    document = getattr(proposal, "document", proposal)
    digest = getattr(proposal, "sha256", "")
    if not isinstance(document, Mapping):
        raise NewtonWorkspaceError("workspace proposal is unavailable")
    if not isinstance(digest, str) or _HEX_64.fullmatch(digest) is None:
        digest = hashlib.sha256(_canonical_bytes(document)).hexdigest()
    if hashlib.sha256(_canonical_bytes(document)).hexdigest() != digest:
        raise NewtonWorkspaceError("workspace proposal digest changed")
    if (
        document.get("schema") != "labcobot.workspace-teach-proposal.v1"
        or document.get("schema_version") != 1
        or document.get("status") != "review_proposal"
        or document.get("coordinate_frame") != "xarm_base_mm"
    ):
        raise NewtonWorkspaceError("workspace proposal contract is invalid")
    session_id = document.get("session_id")
    if not isinstance(session_id, str) or _HEX_32.fullmatch(session_id) is None:
        raise NewtonWorkspaceError("workspace proposal session is invalid")
    bounds = document.get("proposed_bounds_mm")
    if not isinstance(bounds, Mapping) or set(bounds) != {"lower", "upper"}:
        raise NewtonWorkspaceError("workspace proposal bounds are invalid")
    lower = _finite_vector(bounds.get("lower"), 3, "workspace lower bound")
    upper = _finite_vector(bounds.get("upper"), 3, "workspace upper bound")
    if any(low >= high for low, high in zip(lower, upper)):
        raise NewtonWorkspaceError("workspace proposal bounds are reversed")
    samples = document.get("samples")
    if not isinstance(samples, list) or len(samples) != 6:
        raise NewtonWorkspaceError("workspace proposal needs six samples")
    joints = []
    tcp_positions = []
    pose_sequences = []
    for index, sample in enumerate(samples):
        if not isinstance(sample, Mapping):
            raise NewtonWorkspaceError("workspace sample is invalid")
        joints.append(
            _finite_vector(sample.get("joint_angles_rad"), 6, "sample joints")
        )
        tcp = _finite_vector(sample.get("tcp_pose_mm_rad"), 6, "sample TCP pose")
        tcp_positions.append(tcp[:3])
        sequence = sample.get("pose_sequence")
        if isinstance(sequence, bool) or not isinstance(sequence, int) or sequence < 0:
            raise NewtonWorkspaceError("workspace pose sequence is invalid")
        pose_sequences.append(sequence)
    if any(
        after <= before for before, after in zip(pose_sequences, pose_sequences[1:])
    ):
        raise NewtonWorkspaceError("workspace pose sequences are not increasing")
    bindings = document.get("bindings")
    if not isinstance(bindings, Mapping):
        raise NewtonWorkspaceError("workspace proposal bindings are invalid")
    robot_sha256 = _sha256(bindings.get("robot_sha256"), "robot")
    return {
        "document": dict(document),
        "proposal_sha256": digest,
        "session_id": session_id,
        "bounds": {"lower": list(lower), "upper": list(upper)},
        "joints": tuple(joints),
        "tcp_positions": tuple(tcp_positions),
        "robot_sha256": robot_sha256,
    }


class NewtonPhysicsKinematicsValidator:
    """Recompute the six taught Lite6 poses with pinned Newton Physics.

    Newton is imported lazily so the normal Raspberry Pi runtime remains
    usable without the optional simulator environment.  An explicitly enabled
    but unavailable/mismatched Newton installation fails closed.
    """

    def __init__(
        self,
        *,
        urdf_path: str,
        expected_urdf_sha256: str,
        expected_version: str = "1.4.0",
        device: str = "cpu",
        module_loader: Callable[[str], Any] = importlib.import_module,
        kinematics: Any = None,
    ) -> None:
        self.urdf_path = Path(urdf_path)
        self.expected_urdf_sha256 = _sha256(expected_urdf_sha256, "Newton URDF")
        if (
            not isinstance(expected_version, str)
            or _VERSION.fullmatch(expected_version) is None
        ):
            raise NewtonWorkspaceError("Newton version is invalid")
        if device not in {"cpu", "cuda"}:
            raise NewtonWorkspaceError("Newton device must be cpu or cuda")
        self.expected_version = expected_version
        self.device = device
        self._module_loader = module_loader
        if kinematics is None:
            try:
                from main.digital_twin import Lite6VisualizationKinematics
            except ImportError:  # Direct ``main/runtime.py`` execution on the Pi.
                from digital_twin import (  # type: ignore[no-redef]
                    Lite6VisualizationKinematics,
                )

            kinematics = Lite6VisualizationKinematics()
        self.kinematics = kinematics
        self._verify_model_file()

    def _verify_model_file(self) -> None:
        try:
            metadata = self.urdf_path.lstat()
        except OSError as exc:
            raise NewtonWorkspaceError("Newton URDF is unavailable") from exc
        if (
            stat.S_ISLNK(metadata.st_mode)
            or not stat.S_ISREG(metadata.st_mode)
            or metadata.st_size <= 0
            or metadata.st_size > _MAX_MODEL_BYTES
        ):
            raise NewtonWorkspaceError("Newton URDF must be a bounded regular file")
        try:
            payload = self.urdf_path.read_bytes()
        except OSError as exc:
            raise NewtonWorkspaceError("Newton URDF cannot be read") from exc
        actual = hashlib.sha256(payload).hexdigest()
        if actual != self.expected_urdf_sha256:
            raise NewtonWorkspaceError("Newton URDF digest does not match")

    def validate(
        self,
        proposal: Any,
        *,
        stop_requested: Optional[Callable[[], bool]] = None,
    ) -> Mapping[str, Any]:
        normalized = _proposal_input(proposal)
        if stop_requested is not None and stop_requested():
            raise NewtonWorkspaceError("stopped")
        try:
            newton = self._module_loader("newton")
        except Exception as exc:
            raise NewtonWorkspaceError("Newton Physics is unavailable") from exc
        version = str(getattr(newton, "__version__", ""))
        if version != self.expected_version:
            raise NewtonWorkspaceError("Newton Physics version does not match")
        try:
            builder = newton.ModelBuilder()
            builder.add_urdf(
                str(self.urdf_path),
                floating=False,
                hide_visuals=True,
                parse_visuals_as_colliders=False,
                enable_self_collisions=True,
                collapse_fixed_joints=False,
            )
            model = builder.finalize(device=self.device)
            state = model.state()
            if len(state.joint_q) != 6 or len(state.joint_qd) != 6:
                raise NewtonWorkspaceError("Newton Lite6 joint layout is invalid")
            maximum_deviation_mm = 0.0
            for joints in normalized["joints"]:
                if stop_requested is not None and stop_requested():
                    raise NewtonWorkspaceError("stopped")
                state.joint_q.assign(joints)
                state.joint_qd.assign((0.0,) * 6)
                newton.eval_fk(model, state.joint_q, state.joint_qd, state)
                body_q = state.body_q.numpy()
                if len(body_q) < 6:
                    raise NewtonWorkspaceError("Newton Lite6 body layout is invalid")
                newton_tcp_mm = tuple(float(value) * 1000.0 for value in body_q[-1][:3])
                local_tcp_mm = tuple(self.kinematics.link_points_mm(joints)[-1])
                # The recorded TCP may include a separately commissioned tool
                # offset that is intentionally not part of this visualization
                # URDF.  Newton therefore cross-checks the pinned Lite6 model,
                # while the observed TCP samples remain the sole source of the
                # taught workspace box.
                deviation = math.sqrt(
                    sum(
                        (left - right) ** 2
                        for left, right in zip(local_tcp_mm, newton_tcp_mm)
                    )
                )
                maximum_deviation_mm = max(maximum_deviation_mm, deviation)
        except NewtonWorkspaceError:
            raise
        except Exception as exc:
            raise NewtonWorkspaceError("Newton validation failed") from exc
        consistent = maximum_deviation_mm <= _MAX_TCP_DEVIATION_MM
        result = {
            "engine": "newton",
            "version": version,
            "device": self.device,
            "model_sha256": self.expected_urdf_sha256,
            "robot_sha256": normalized["robot_sha256"],
            "proposal_sha256": normalized["proposal_sha256"],
            "sample_count": 6,
            "max_tcp_deviation_mm": round(maximum_deviation_mm, 6),
            "kinematics_consistent": consistent,
        }
        result["result_sha256"] = hashlib.sha256(_canonical_bytes(result)).hexdigest()
        return result


class WorkspaceBoundaryValidationCoordinator:
    """Run dual-depth and Newton validation in a bounded background thread."""

    def __init__(
        self,
        *,
        dual_depth_preflight: Any,
        newton_validator: Any,
    ) -> None:
        self.dual_depth_preflight = dual_depth_preflight
        self.newton_validator = newton_validator
        self._lock = threading.Lock()
        self._cancel = threading.Event()
        self._thread: Optional[threading.Thread] = None
        self._run_id = 0
        self._status = self._empty_status(configured=True, state="idle")

    @staticmethod
    def _empty_status(*, configured: bool, state: str) -> Mapping[str, Any]:
        return {
            "schema_version": 1,
            "configured": configured,
            "state": state,
            "session_id": "",
            "proposal_sha256": "",
            "robot_sha256": "",
            "bounds_mm": {"lower": [], "upper": []},
            "cameras": {
                "tool": {
                    "source_id": "",
                    "frame_sequence": 0,
                    "calibration_sha256": "",
                },
                "workspace": {
                    "source_id": "",
                    "frame_sequence": 0,
                    "calibration_sha256": "",
                },
            },
            "camera_evidence_sha256": "",
            "newton_version": "",
            "newton_model_sha256": "",
            "newton_result_sha256": "",
            "blockers": [],
            "motion_authority": "none",
            "motion_authorized": False,
        }

    @classmethod
    def unavailable(cls) -> Mapping[str, Any]:
        return cls._empty_status(configured=False, state="unavailable")

    def snapshot(self) -> Mapping[str, Any]:
        with self._lock:
            return json.loads(_canonical_bytes(self._status).decode("utf-8"))

    def start(self, proposal: Any) -> bool:
        try:
            normalized = _proposal_input(proposal)
        except NewtonWorkspaceError:
            with self._lock:
                value = self._empty_status(configured=True, state="blocked")
                value["blockers"] = ["proposal_invalid"]
                self._status = value
            return False
        with self._lock:
            if self._thread is not None and self._thread.is_alive():
                return False
            self._cancel.set()
            self._cancel = threading.Event()
            self._run_id += 1
            run_id = self._run_id
            value = self._empty_status(configured=True, state="validating")
            value.update(
                {
                    "session_id": normalized["session_id"],
                    "proposal_sha256": normalized["proposal_sha256"],
                    "robot_sha256": normalized["robot_sha256"],
                    "bounds_mm": normalized["bounds"],
                }
            )
            self._status = value
            thread = threading.Thread(
                target=self._run,
                args=(run_id, proposal, normalized, self._cancel),
                name="workspace-boundary-validation",
                daemon=True,
            )
            self._thread = thread
            thread.start()
            return True

    def invalidate(self) -> None:
        with self._lock:
            self._cancel.set()
            if self._status.get("state") in ("validating", "valid"):
                value = dict(self._status)
                value["state"] = "invalidated"
                value["blockers"] = ["stopped"]
                self._status = value

    def close(self, timeout_seconds: float = 1.0) -> bool:
        self.invalidate()
        with self._lock:
            thread = self._thread
        if thread is not None and thread is not threading.current_thread():
            thread.join(max(0.0, float(timeout_seconds)))
        return thread is None or not thread.is_alive()

    @staticmethod
    def _camera_contract(evidence: Any) -> Tuple[Mapping[str, Any], str]:
        if not isinstance(evidence, Mapping):
            raise NewtonWorkspaceError("dual-depth evidence is invalid")
        cameras = evidence.get("cameras")
        if not isinstance(cameras, Mapping) or set(cameras) != {"tool", "workspace"}:
            raise NewtonWorkspaceError("dual-depth camera evidence is invalid")
        result = {}
        for role in ("tool", "workspace"):
            camera = cameras.get(role)
            if not isinstance(camera, Mapping):
                raise NewtonWorkspaceError("dual-depth camera evidence is invalid")
            source_id = camera.get("source_id")
            sequence = camera.get("frame_sequence")
            digest = camera.get("calibration_sha256")
            if (
                not isinstance(source_id, str)
                or not source_id
                or len(source_id) > 255
                or isinstance(sequence, bool)
                or not isinstance(sequence, int)
                or sequence <= 0
                or not isinstance(digest, str)
                or _HEX_64.fullmatch(digest) is None
                or camera.get("depth_ready") is not True
            ):
                raise NewtonWorkspaceError("dual-depth camera evidence is incomplete")
            result[role] = {
                "source_id": source_id,
                "frame_sequence": sequence,
                "calibration_sha256": digest,
            }
        if result["tool"]["source_id"] == result["workspace"]["source_id"]:
            raise NewtonWorkspaceError("dual-depth cameras are not distinct")
        return result, hashlib.sha256(_canonical_bytes(evidence)).hexdigest()

    def _run(
        self,
        run_id: int,
        proposal: Any,
        normalized: Mapping[str, Any],
        cancel: threading.Event,
    ) -> None:
        blockers = []
        camera_contract = None
        camera_digest = ""
        newton_result: Mapping[str, Any] = {}
        try:
            decision = self.dual_depth_preflight(
                normalized["joints"],
                stop_requested=cancel.is_set,
            )
            if cancel.is_set():
                blockers.append("stopped")
            elif getattr(decision, "clear", False) is not True:
                blockers.append("dual_depth_blocked")
            else:
                evidence = getattr(decision, "evidence", None)
                camera_contract, camera_digest = self._camera_contract(evidence)
        except Exception:
            blockers.append("dual_depth_unavailable")
        if not blockers:
            try:
                newton_result = self.newton_validator.validate(
                    proposal,
                    stop_requested=cancel.is_set,
                )
                if cancel.is_set():
                    blockers.append("stopped")
                elif (
                    not isinstance(newton_result, Mapping)
                    or newton_result.get("engine") != "newton"
                    or newton_result.get("proposal_sha256")
                    != normalized["proposal_sha256"]
                    or newton_result.get("kinematics_consistent") is not True
                    or not isinstance(newton_result.get("result_sha256"), str)
                    or _HEX_64.fullmatch(newton_result["result_sha256"]) is None
                ):
                    blockers.append("newton_mismatch")
            except NewtonWorkspaceError as exc:
                blockers.append(
                    "stopped" if str(exc) == "stopped" else "newton_unavailable"
                )
            except Exception:
                blockers.append("newton_unavailable")
        blockers = [value for value in WORKSPACE_BOUNDARY_BLOCKERS if value in blockers]
        with self._lock:
            if run_id != self._run_id or cancel is not self._cancel:
                return
            value = dict(self._status)
            if cancel.is_set() and "stopped" not in blockers:
                blockers.append("stopped")
            value["state"] = (
                "valid"
                if not blockers
                else ("invalidated" if blockers == ["stopped"] else "blocked")
            )
            value["blockers"] = blockers
            if camera_contract is not None:
                value["cameras"] = camera_contract
                value["camera_evidence_sha256"] = camera_digest
            if newton_result:
                value["newton_version"] = str(newton_result.get("version") or "")
                value["newton_model_sha256"] = str(
                    newton_result.get("model_sha256") or ""
                )
                value["newton_result_sha256"] = str(
                    newton_result.get("result_sha256") or ""
                )
            value["motion_authority"] = "none"
            value["motion_authorized"] = False
            self._status = value


__all__ = [
    "NewtonPhysicsKinematicsValidator",
    "NewtonWorkspaceError",
    "WORKSPACE_BOUNDARY_BLOCKERS",
    "WorkspaceBoundaryValidationCoordinator",
]
