"""Pinned-image input and provenance contract for Wakeword Training (#163).

This module deliberately does not implement wakeword generation, augmentation,
feature extraction, training, export, or evaluation.  It validates the
commercial source lock, materializes only its declared inputs, and leaves the
actual ML stages to the released LiveKit CLI inside the isolated image.
"""

from __future__ import annotations

import hashlib
import json
import os
import re
import shutil
from dataclasses import dataclass
from enum import StrEnum
from pathlib import Path
from typing import Any

LIVEKIT_WAKEWORD_VERSION = "0.2.1"
LIVEKIT_WAKEWORD_REVISION = "1ec7f680df30ff4ca0ebae6b5983441e94b10980"
LIVEKIT_WAKEWORD_SDIST_SHA256 = "cf2d9cf4867812c06788f64c15e49abd909d9d6291f0a13f1c3f9cb649fa6127"

_COMMERCIAL_PURPOSES = frozenset({"wakeword_training", "classifier_distribution"})
_BLOCKED_LICENSE_MARKERS = ("unknown", "nc", "non-commercial", "research-only")
_BLOCKED_SOURCE_MARKERS = ("acav", "unknown_rir_mirror")
_SYNTHETIC_SOURCE_KINDS = frozenset({"piper_output", "voxcpm_output"})
_SAFE_TOKEN = re.compile(r"^[a-z0-9][a-z0-9._-]{0,127}$")
_ALLOWED_PURPOSES = frozenset(
    {
        "positive",
        "adversarial_negative",
        "near_miss",
        "general_negative_speech",
        "background",
        "rir",
    }
)
_ALLOWED_PARTITIONS = frozenset({"train", "validation", "calibration", "release_test"})
_REQUIRED_PURPOSE_PARTITIONS = frozenset(
    {
        ("positive", "train"),
        ("positive", "release_test"),
        ("general_negative_speech", "train"),
        ("general_negative_speech", "validation"),
        ("background", "train"),
        ("rir", "train"),
    }
)


class SourceLockError(ValueError):
    """A source lock cannot safely enter the requested training lane."""


class RunClassification(StrEnum):
    COMMERCIAL_REVIEW_REQUIRED = "commercial_review_required"
    LOCAL_ONLY = "local_only"
    QUARANTINED = "quarantined"

    @property
    def commercial_promotion_eligible(self) -> bool:
        return self is RunClassification.COMMERCIAL_REVIEW_REQUIRED


def classify_run(run_kind: str) -> RunClassification:
    if run_kind == "commercial":
        return RunClassification.COMMERCIAL_REVIEW_REQUIRED
    if run_kind == "local_experiment":
        return RunClassification.LOCAL_ONLY
    if run_kind == "quarantined_smoke":
        return RunClassification.QUARANTINED
    raise SourceLockError("run_kind must be commercial, local_experiment, or quarantined_smoke")


def _sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as file:
        for block in iter(lambda: file.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def _require_mapping(value: object, label: str) -> dict[str, Any]:
    if not isinstance(value, dict):
        raise SourceLockError(f"{label} must be an object")
    return value


def _require_string(value: object, label: str) -> str:
    if not isinstance(value, str) or not value.strip():
        raise SourceLockError(f"{label} must be a non-empty string")
    return value.strip()


def _require_sha256(value: object, label: str) -> str:
    digest = _require_string(value, label).removeprefix("sha256:").lower()
    if len(digest) != 64 or any(character not in "0123456789abcdef" for character in digest):
        raise SourceLockError(f"{label} must be a SHA-256 digest")
    return digest


def _safe_input_path(root: Path, declared_path: str) -> Path:
    candidate = (root / declared_path).resolve(strict=True)
    if not candidate.is_file() or os.path.commonpath((root, candidate)) != str(root):
        raise SourceLockError("source lock file path must identify a regular file below input_root")
    return candidate


@dataclass(frozen=True)
class MaterializedInputs:
    provenance_manifest_path: Path
    recorded_manifest_path: Path


class InputMaterializer:
    """Verify and copy a source lock's declared inputs into one sealed run.

    The input root is read-only when this code runs in the training image.  The
    run directory is the sole persistence target, so no upstream ``setup``
    download can become an undeclared training input.
    """

    def __init__(self, lock_path: Path, input_root: Path, run_dir: Path, *, run_kind: str) -> None:
        self.lock_path = lock_path.resolve(strict=True)
        self.input_root = input_root.resolve(strict=True)
        self.run_dir = run_dir.resolve()
        self.run_kind = run_kind

    def materialize(self) -> MaterializedInputs:
        classification = classify_run(self.run_kind)
        lock = self._read_lock()
        if lock.get("run_kind") != self.run_kind:
            raise SourceLockError("source lock run_kind does not match the requested run")
        sources = lock.get("sources")
        if not isinstance(sources, list) or not sources:
            raise SourceLockError("source lock must contain at least one source")

        materialized_root = self.run_dir / "materialized"
        materialized_root.mkdir(parents=True, exist_ok=False)
        source_manifest: list[dict[str, object]] = []
        supplied_purposes: set[tuple[str, str]] = set()
        seen_source_ids: set[str] = set()
        materialized_destinations: set[Path] = set()
        for index, raw_source in enumerate(sources):
            source = _require_mapping(raw_source, f"sources[{index}]")
            source_id, purpose, partition, file_records = self._validate_source(
                source, index, commercial=classification.commercial_promotion_eligible
            )
            if source_id in seen_source_ids:
                raise SourceLockError(f"duplicate source_id: {source_id}")
            seen_source_ids.add(source_id)
            supplied_purposes.add((purpose, partition))
            target_root = materialized_root / partition / purpose / source_id
            target_root.mkdir(parents=True)
            materialized_files: list[dict[str, str]] = []
            for record in file_records:
                source_path = _safe_input_path(self.input_root, record["path"])
                if _sha256_file(source_path) != record["sha256"]:
                    raise SourceLockError(f"source hash mismatch: {record['path']}")
                target = target_root / source_path.name
                if target in materialized_destinations:
                    raise SourceLockError(f"duplicate materialized destination: {target.name}")
                materialized_destinations.add(target)
                shutil.copyfile(source_path, target)
                target.chmod(0o444)
                materialized_files.append(
                    {
                        "source_path": record["path"],
                        "materialized_path": target.relative_to(self.run_dir).as_posix(),
                        "sha256": record["sha256"],
                    }
                )
            source_manifest.append(
                {
                    "source_id": source_id,
                    "purpose": purpose,
                    "partition": partition,
                    "files": materialized_files,
                }
            )

        if classification.commercial_promotion_eligible:
            missing = _REQUIRED_PURPOSE_PARTITIONS - supplied_purposes
            if missing:
                formatted = ", ".join(
                    f"{purpose}:{partition}" for purpose, partition in sorted(missing)
                )
                raise SourceLockError(
                    f"commercial source lock is missing required inputs: {formatted}"
                )

        provenance_path = self.run_dir / "sources.lock.json"
        if self.lock_path != provenance_path:
            shutil.copyfile(self.lock_path, provenance_path)
        provenance_path.chmod(0o444)
        attribution_dir = self.run_dir / "attribution"
        attribution_dir.mkdir(exist_ok=True)
        self._materialize_attribution(lock["sources"], attribution_dir)
        recorded_manifest_path = self.run_dir / "recorded-inputs.json"
        recorded_manifest_path.write_text(
            json.dumps(
                {
                    "schema_version": 1,
                    "classification": classification.value,
                    "sources": source_manifest,
                },
                sort_keys=True,
                separators=(",", ":"),
            )
            + "\n",
            encoding="utf-8",
        )
        recorded_manifest_path.chmod(0o444)
        return MaterializedInputs(
            provenance_manifest_path=provenance_path,
            recorded_manifest_path=recorded_manifest_path,
        )

    def _read_lock(self) -> dict[str, Any]:
        try:
            lock = json.loads(self.lock_path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError) as exc:
            raise SourceLockError("source lock must be readable JSON") from exc
        lock = _require_mapping(lock, "source lock")
        if lock.get("schema_version") != 1:
            raise SourceLockError("source lock schema_version must be 1")
        return lock

    def _validate_source(
        self, source: dict[str, Any], index: int, *, commercial: bool
    ) -> tuple[str, str, str, list[dict[str, str]]]:
        source_id = _require_string(source.get("source_id"), f"sources[{index}].source_id")
        purpose = _require_string(source.get("purpose"), f"sources[{index}].purpose")
        partition = _require_string(source.get("partition"), f"sources[{index}].partition")
        source_kind = _require_string(source.get("source_kind"), f"sources[{index}].source_kind")
        if not _SAFE_TOKEN.fullmatch(source_id):
            raise SourceLockError(f"sources[{index}].source_id must be a safe token")
        if purpose not in _ALLOWED_PURPOSES:
            raise SourceLockError(f"sources[{index}].purpose is not an allowed purpose")
        if partition not in _ALLOWED_PARTITIONS:
            raise SourceLockError(f"sources[{index}].partition is not an allowed partition")
        publisher = _require_mapping(
            source.get("original_publisher"), f"sources[{index}].original_publisher"
        )
        _require_string(publisher.get("name"), f"sources[{index}].original_publisher.name")
        _require_string(publisher.get("url"), f"sources[{index}].original_publisher.url")
        _require_string(source.get("version"), f"sources[{index}].version")
        _require_string(source.get("revision"), f"sources[{index}].revision")
        _require_string(source.get("attribution"), f"sources[{index}].attribution")
        license_data = _require_mapping(source.get("license"), f"sources[{index}].license")
        license_id = _require_string(
            license_data.get("spdx_id"), f"sources[{index}].license.spdx_id"
        )
        _require_string(license_data.get("terms_url"), f"sources[{index}].license.terms_url")
        captured_text = _require_string(
            license_data.get("captured_text_path"),
            f"sources[{index}].license.captured_text_path",
        )
        captured_hash = _require_sha256(
            license_data.get("captured_text_sha256"),
            f"sources[{index}].license.captured_text_sha256",
        )
        license_path = _safe_input_path(self.input_root, captured_text)
        if _sha256_file(license_path) != captured_hash:
            raise SourceLockError(f"captured license hash mismatch: {captured_text}")
        files = source.get("files")
        if not isinstance(files, list) or not files:
            raise SourceLockError(f"sources[{index}].files must be a non-empty list")
        records: list[dict[str, str]] = []
        for file_index, raw_file in enumerate(files):
            file_data = _require_mapping(raw_file, f"sources[{index}].files[{file_index}]")
            path = _require_string(
                file_data.get("path"), f"sources[{index}].files[{file_index}].path"
            )
            digest = _require_sha256(
                file_data.get("sha256"), f"sources[{index}].files[{file_index}].sha256"
            )
            size = file_data.get("size_bytes")
            if not isinstance(size, int) or size < 1:
                raise SourceLockError(
                    f"sources[{index}].files[{file_index}].size_bytes must be positive"
                )
            if _safe_input_path(self.input_root, path).stat().st_size != size:
                raise SourceLockError(f"source size mismatch: {path}")
            records.append({"path": path, "sha256": digest})
        if commercial:
            self._validate_commercial_source(source_id, source_kind, license_id, source, index)
        return source_id, purpose, partition, records

    @staticmethod
    def _validate_commercial_source(
        source_id: str, source_kind: str, license_id: str, source: dict[str, Any], index: int
    ) -> None:
        normalized_id = source_id.lower()
        normalized_license = license_id.lower()
        if any(marker in normalized_id for marker in _BLOCKED_SOURCE_MARKERS) or any(
            marker in normalized_license for marker in _BLOCKED_LICENSE_MARKERS
        ):
            raise SourceLockError(f"blocked source for commercial mode: {source_id}")
        commercial_data = _require_mapping(source.get("commercial"), f"sources[{index}].commercial")
        if commercial_data.get("disposition") != "approved":
            raise SourceLockError(f"commercial source is not approved: {source_id}")
        allowed_purposes = commercial_data.get("allowed_purposes")
        if not isinstance(allowed_purposes, list) or not _COMMERCIAL_PURPOSES.issubset(
            set(allowed_purposes)
        ):
            raise SourceLockError(f"commercial source has incomplete allowed_purposes: {source_id}")
        _require_string(commercial_data.get("approver"), f"sources[{index}].commercial.approver")
        _require_string(
            commercial_data.get("approved_at"), f"sources[{index}].commercial.approved_at"
        )
        if (
            source_kind in _SYNTHETIC_SOURCE_KINDS
            and commercial_data.get("tts_output_approved") is not True
        ):
            raise SourceLockError(f"unapproved synthetic source for commercial mode: {source_id}")

    def _materialize_attribution(self, sources: object, attribution_dir: Path) -> None:
        assert isinstance(sources, list)
        for raw_source in sources:
            source = _require_mapping(raw_source, "source")
            source_id = _require_string(source.get("source_id"), "source.source_id")
            license_data = _require_mapping(source.get("license"), "source.license")
            license_path = _safe_input_path(
                self.input_root,
                _require_string(
                    license_data.get("captured_text_path"), "source.license.captured_text_path"
                ),
            )
            target = attribution_dir / f"{source_id}-license.txt"
            shutil.copyfile(license_path, target)
            target.chmod(0o444)


def build_stage_command(run_kind: str) -> tuple[str, ...]:
    """Return the only command an isolated training container may execute."""

    classification = classify_run(run_kind)
    if classification is RunClassification.COMMERCIAL_REVIEW_REQUIRED:
        return (
            "python",
            "/opt/jr-wakeword/run-commercial-pipeline.py",
            "/run/resolved-config.yaml",
        )
    if classification is RunClassification.LOCAL_ONLY:
        return (
            "python",
            "/opt/jr-wakeword/run-local-experiment.py",
            "/run/resolved-config.yaml",
        )
    return (
        "python",
        "/opt/jr-wakeword/run-quarantined-smoke.py",
        "/run/resolved-config.yaml",
    )


def collect_image_provenance(
    *,
    dependency_lock_path: Path,
    frontend_asset_paths: list[Path],
    installed_distributions: dict[str, str],
    image_digest: str,
    python_version: str,
    cuda_version: str,
    os_packages: list[str],
) -> dict[str, object]:
    """Return manifest-safe image evidence; callers persist it in the run bundle."""

    if not image_digest.startswith("sha256:"):
        raise ValueError("image_digest must be sha256-pinned")
    lock_path = dependency_lock_path.resolve(strict=True)
    assets = []
    for path in frontend_asset_paths:
        resolved = path.resolve(strict=True)
        assets.append({"path": resolved.name, "sha256": _sha256_file(resolved)})
    return {
        "schema_version": 1,
        "image_digest": image_digest,
        "python_version": python_version,
        "cuda_version": cuda_version,
        "os_packages": sorted(os_packages),
        "livekit_wakeword": {
            "distribution": "livekit-wakeword",
            "version": LIVEKIT_WAKEWORD_VERSION,
            "source_revision": LIVEKIT_WAKEWORD_REVISION,
            "source_distribution_sha256": LIVEKIT_WAKEWORD_SDIST_SHA256,
            "installed_version": installed_distributions.get("livekit-wakeword"),
        },
        "installed_distributions": dict(sorted(installed_distributions.items())),
        "dependency_lock": {"path": lock_path.name, "sha256": _sha256_file(lock_path)},
        "frontend_assets": assets,
    }