"""Reviewed, request-bound artifacts for local and offloaded imitation training."""

from __future__ import annotations

import hashlib
import hmac
import json
import math
import os
import tempfile
from pathlib import Path
from typing import Any, List, Mapping, Optional, Sequence, Tuple, Union

if __package__:
    from .functions.trajectory_training import (
        EvaluationMetrics,
        NearestNeighborPolicy,
        TrainingConfig,
        TrajectoryError,
    )
else:  # Direct ``main/runtime.py`` execution on the Raspberry Pi.
    from functions.trajectory_training import (  # type: ignore[no-redef]
        EvaluationMetrics,
        NearestNeighborPolicy,
        TrainingConfig,
        TrajectoryError,
    )


POLICY_SCHEMA = "labcobot.imitation-policy.v1"
REMOTE_RESULT_SCHEMA = "labcobot.training-result.v1"
MAXIMUM_ARTIFACT_BYTES = 32 * 1024 * 1024
MAXIMUM_EXAMPLES = 100_000


class TrainingArtifactError(TrajectoryError):
    """Raised when a training result is not safe to review or deploy."""


def canonical_training_request_sha256(payload: Mapping[str, Any]) -> str:
    """Return the deterministic digest a remote result must echo."""

    try:
        serialized = json.dumps(
            dict(payload),
            allow_nan=False,
            separators=(",", ":"),
            sort_keys=True,
        ).encode("utf-8")
    except (TypeError, ValueError) as exc:
        raise TrainingArtifactError("Training request is not canonical JSON") from exc
    return hashlib.sha256(serialized).hexdigest()


def policy_payload(
    policy: NearestNeighborPolicy,
    *,
    config: TrainingConfig,
    evaluation: EvaluationMetrics,
    request_sha256: str,
) -> Mapping[str, Any]:
    """Build one unit-explicit policy artifact after evaluation."""

    request_digest = _sha256(request_sha256, "training request")
    return _validated_policy_payload(
        {
            "schema": POLICY_SCHEMA,
            "algorithm": "nearest-neighbor",
            "units": {
                "joints": "radian",
                "time": "second",
            },
            "request_sha256": request_digest,
            "config": {
                "feature_names": list(config.feature_names),
                "include_current_joints": bool(config.include_current_joints),
                "action_mode": config.action_mode,
                "neighbors": int(config.neighbors),
                "evaluation_tolerance": float(config.evaluation_tolerance),
            },
            "evaluation": {
                "sample_count": int(evaluation.sample_count),
                "mean_absolute_error": float(evaluation.mean_absolute_error),
                "root_mean_squared_error": float(
                    evaluation.root_mean_squared_error
                ),
                "max_absolute_error": float(evaluation.max_absolute_error),
                "within_tolerance_fraction": float(
                    evaluation.within_tolerance_fraction
                ),
            },
            "policy": {
                "neighbors": int(policy.neighbors),
                "features": [list(row) for row in policy.features],
                "actions": [list(row) for row in policy.actions],
            },
        }
    )


def validate_remote_training_result(
    result: Any,
    *,
    expected_request_sha256: str,
) -> Mapping[str, Any]:
    """Accept only a correlated, checksummed policy returned by ShortcutJob."""

    if not isinstance(result, Mapping):
        raise TrainingArtifactError("Remote training result must be an object")
    if result.get("schema") != REMOTE_RESULT_SCHEMA:
        raise TrainingArtifactError("Unsupported remote training-result schema")
    expected = _sha256(expected_request_sha256, "expected training request")
    returned = _sha256(
        str(result.get("request_sha256") or ""),
        "remote training request",
    )
    if not hmac.compare_digest(returned, expected):
        raise TrainingArtifactError(
            "Remote training result belongs to a different request"
        )
    raw_policy = result.get("artifact")
    if not isinstance(raw_policy, Mapping):
        raise TrainingArtifactError("Remote training result has no policy artifact")
    validated = _validated_policy_payload(raw_policy)
    serialized = _canonical_json(validated)
    expected_artifact = _sha256(
        str(result.get("artifact_sha256") or ""),
        "remote policy artifact",
    )
    actual_artifact = hashlib.sha256(serialized).hexdigest()
    if not hmac.compare_digest(actual_artifact, expected_artifact):
        raise TrainingArtifactError(
            "Remote policy artifact SHA-256 does not match its content"
        )
    if not hmac.compare_digest(
        str(validated["request_sha256"]),
        expected,
    ):
        raise TrainingArtifactError(
            "Remote policy artifact belongs to a different request"
        )
    return validated


class ImitationPolicyArtifactStore:
    """Atomically persist and restore a bounded reviewed policy artifact."""

    def save(
        self,
        path: Union[os.PathLike, str],
        payload: Mapping[str, Any],
    ) -> Path:
        validated = _validated_policy_payload(payload)
        serialized = json.dumps(
            validated,
            allow_nan=False,
            indent=2,
            sort_keys=True,
        ).encode("utf-8")
        if len(serialized) > MAXIMUM_ARTIFACT_BYTES:
            raise TrainingArtifactError("Policy artifact exceeds the size limit")
        destination = Path(path)
        destination.parent.mkdir(parents=True, exist_ok=True)
        descriptor, temporary_name = tempfile.mkstemp(
            prefix="." + destination.name + ".",
            suffix=".tmp",
            dir=str(destination.parent),
        )
        try:
            with os.fdopen(descriptor, "wb") as handle:
                handle.write(serialized)
                handle.flush()
                os.fsync(handle.fileno())
            os.chmod(temporary_name, 0o600)
            os.replace(temporary_name, destination)
        finally:
            try:
                os.unlink(temporary_name)
            except FileNotFoundError:
                pass
        return destination

    def load(
        self,
        path: Union[os.PathLike, str],
        *,
        expected_sha256: Optional[str] = None,
    ) -> Tuple[NearestNeighborPolicy, Mapping[str, Any]]:
        source = Path(path)
        if source.is_symlink() or not source.is_file():
            raise TrainingArtifactError(
                "Policy artifact must be a regular, non-symlink file"
            )
        raw = source.read_bytes()
        if not raw or len(raw) > MAXIMUM_ARTIFACT_BYTES:
            raise TrainingArtifactError("Policy artifact has an invalid size")
        if expected_sha256 is not None:
            expected = _sha256(expected_sha256, "expected policy artifact")
            if not hmac.compare_digest(hashlib.sha256(raw).hexdigest(), expected):
                raise TrainingArtifactError(
                    "Policy artifact SHA-256 does not match the reviewed value"
                )
        try:
            value = json.loads(raw.decode("utf-8"))
        except (UnicodeError, json.JSONDecodeError) as exc:
            raise TrainingArtifactError("Policy artifact is not valid JSON") from exc
        if not isinstance(value, Mapping):
            raise TrainingArtifactError("Policy artifact must contain an object")
        validated = _validated_policy_payload(value)
        policy = validated["policy"]
        return (
            NearestNeighborPolicy(
                features=tuple(tuple(row) for row in policy["features"]),
                actions=tuple(tuple(row) for row in policy["actions"]),
                neighbors=int(policy["neighbors"]),
            ),
            validated,
        )


def _validated_policy_payload(value: Mapping[str, Any]) -> Mapping[str, Any]:
    if value.get("schema") != POLICY_SCHEMA:
        raise TrainingArtifactError("Unsupported imitation-policy schema")
    if value.get("algorithm") != "nearest-neighbor":
        raise TrainingArtifactError("Unsupported imitation-policy algorithm")
    if value.get("units") != {"joints": "radian", "time": "second"}:
        raise TrainingArtifactError("Policy units are missing or unsupported")
    request_digest = _sha256(
        str(value.get("request_sha256") or ""),
        "policy training request",
    )
    config = _mapping(value.get("config"), "policy config")
    action_mode = str(config.get("action_mode") or "")
    if action_mode not in {"joint_delta", "joint_position"}:
        raise TrainingArtifactError("Policy action mode is unsupported")
    neighbors = _positive_integer(config.get("neighbors"), "config neighbors")
    tolerance = _finite_positive(
        config.get("evaluation_tolerance"),
        "evaluation tolerance",
    )
    feature_names_raw = config.get("feature_names")
    if not isinstance(feature_names_raw, list) or any(
        not isinstance(name, str) or not name
        for name in feature_names_raw
    ):
        raise TrainingArtifactError("Policy feature names are invalid")
    include_current = config.get("include_current_joints")
    if not isinstance(include_current, bool):
        raise TrainingArtifactError(
            "Policy include_current_joints must be true or false"
        )

    evaluation = _mapping(value.get("evaluation"), "policy evaluation")
    sample_count = _positive_integer(
        evaluation.get("sample_count"),
        "evaluation sample count",
    )
    mae = _finite_nonnegative(
        evaluation.get("mean_absolute_error"),
        "mean absolute error",
    )
    rmse = _finite_nonnegative(
        evaluation.get("root_mean_squared_error"),
        "root mean squared error",
    )
    maximum = _finite_nonnegative(
        evaluation.get("max_absolute_error"),
        "maximum absolute error",
    )
    within = _finite_number(
        evaluation.get("within_tolerance_fraction"),
        "within-tolerance fraction",
    )
    if not 0.0 <= within <= 1.0:
        raise TrainingArtifactError(
            "Within-tolerance fraction must be between zero and one"
        )

    policy = _mapping(value.get("policy"), "nearest-neighbor policy")
    policy_neighbors = _positive_integer(
        policy.get("neighbors"),
        "policy neighbors",
    )
    if policy_neighbors != neighbors:
        raise TrainingArtifactError(
            "Policy neighbors do not match the training configuration"
        )
    features = _matrix(policy.get("features"), "policy features")
    actions = _matrix(policy.get("actions"), "policy actions")
    if len(features) != len(actions):
        raise TrainingArtifactError(
            "Policy feature and action row counts must match"
        )
    if not features or len(features) > MAXIMUM_EXAMPLES:
        raise TrainingArtifactError("Policy example count is outside limits")
    if policy_neighbors > len(features):
        raise TrainingArtifactError(
            "Policy neighbors exceed the available training examples"
        )
    if sample_count < 1:
        raise TrainingArtifactError("Policy evaluation is empty")

    return {
        "schema": POLICY_SCHEMA,
        "algorithm": "nearest-neighbor",
        "units": {"joints": "radian", "time": "second"},
        "request_sha256": request_digest,
        "config": {
            "feature_names": list(feature_names_raw),
            "include_current_joints": include_current,
            "action_mode": action_mode,
            "neighbors": neighbors,
            "evaluation_tolerance": tolerance,
        },
        "evaluation": {
            "sample_count": sample_count,
            "mean_absolute_error": mae,
            "root_mean_squared_error": rmse,
            "max_absolute_error": maximum,
            "within_tolerance_fraction": within,
        },
        "policy": {
            "neighbors": policy_neighbors,
            "features": features,
            "actions": actions,
        },
    }


def _matrix(value: Any, name: str) -> List[List[float]]:
    if not isinstance(value, list) or not value:
        raise TrainingArtifactError("{} must be a non-empty list".format(name))
    rows = []
    width = None
    for row in value:
        if not isinstance(row, list) or not row:
            raise TrainingArtifactError("{} rows must be non-empty lists".format(name))
        numbers = [_finite_number(item, name) for item in row]
        if width is None:
            width = len(numbers)
        elif len(numbers) != width:
            raise TrainingArtifactError("{} rows have inconsistent size".format(name))
        rows.append(numbers)
    return rows


def _mapping(value: Any, name: str) -> Mapping[str, Any]:
    if not isinstance(value, Mapping):
        raise TrainingArtifactError("{} must be an object".format(name))
    return value


def _positive_integer(value: Any, name: str) -> int:
    if isinstance(value, bool):
        raise TrainingArtifactError("{} must be a positive integer".format(name))
    try:
        number = int(value)
    except (TypeError, ValueError) as exc:
        raise TrainingArtifactError(
            "{} must be a positive integer".format(name)
        ) from exc
    if number < 1 or number != value:
        raise TrainingArtifactError("{} must be a positive integer".format(name))
    return number


def _finite_number(value: Any, name: str) -> float:
    try:
        number = float(value)
    except (TypeError, ValueError) as exc:
        raise TrainingArtifactError("{} must be numeric".format(name)) from exc
    if not math.isfinite(number):
        raise TrainingArtifactError("{} must be finite".format(name))
    return number


def _finite_positive(value: Any, name: str) -> float:
    number = _finite_number(value, name)
    if number <= 0:
        raise TrainingArtifactError("{} must be positive".format(name))
    return number


def _finite_nonnegative(value: Any, name: str) -> float:
    number = _finite_number(value, name)
    if number < 0:
        raise TrainingArtifactError("{} cannot be negative".format(name))
    return number


def _sha256(value: str, name: str) -> str:
    cleaned = str(value).strip().lower()
    if len(cleaned) != 64 or any(
        character not in "0123456789abcdef" for character in cleaned
    ):
        raise TrainingArtifactError("{} SHA-256 is invalid".format(name))
    return cleaned


def _canonical_json(value: Mapping[str, Any]) -> bytes:
    return json.dumps(
        dict(value),
        allow_nan=False,
        separators=(",", ":"),
        sort_keys=True,
    ).encode("utf-8")


__all__ = [
    "ImitationPolicyArtifactStore",
    "POLICY_SCHEMA",
    "REMOTE_RESULT_SCHEMA",
    "TrainingArtifactError",
    "canonical_training_request_sha256",
    "policy_payload",
    "validate_remote_training_result",
]
