from __future__ import annotations

import hashlib
import json
import os
import re
import secrets
import shutil
import stat
from dataclasses import dataclass
from pathlib import Path, PurePosixPath
from typing import Any

from k1link.artifacts import write_json_atomic
from k1link.sessions import SessionIntegrityError, inspect_recorded_media_epoch

COMPUTE_JOB_SCHEMA = "missioncore.compute-job/v1"
COMPUTE_RESULT_SCHEMA = "missioncore.compute-result/v1"
COMPUTE_PROFILE = "recorded-camera-perception/v1"
MAX_JOB_MANIFEST_BYTES = 32 * 1024 * 1024
MAX_JOB_FILES = 500_003

_SAFE_COMPONENT = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$")
_SAFE_JOB_ID = re.compile(r"^recorded-camera-[a-f0-9]{24}$")
_SHA256 = re.compile(r"^[a-f0-9]{64}$")


@dataclass(frozen=True, slots=True)
class CameraComputeJob:
    job_id: str
    job_root: Path
    manifest_path: Path
    session_id: str
    source_id: str
    codec_epoch: int
    input_sha256: str
    input_byte_length: int
    segment_count: int
    timeline_start_seconds: float
    timeline_end_seconds: float


def prepare_camera_compute_job(
    *,
    session_root: Path,
    source_id: str,
    codec_epoch: int,
    origin_epoch_ns: int,
    origin_monotonic_ns: int,
    output_root: Path,
) -> CameraComputeJob:
    """Publish one immutable, bounded camera epoch as a compute job.

    The source evidence is read and validated but never rewritten. Publication
    copies the canonical epoch into a private staging directory, revalidates the
    copy, writes a path-free content-addressed manifest, then atomically renames
    the complete job into place.
    """

    if _SAFE_COMPONENT.fullmatch(source_id) is None:
        raise ValueError("camera source id is not a safe component")
    if codec_epoch < 1:
        raise ValueError("camera codec epoch must be positive")
    if origin_epoch_ns < 0 or origin_monotonic_ns < 0:
        raise ValueError("session timeline origin must be non-negative")

    session = session_root.expanduser().resolve(strict=True)
    if not session.is_dir() or _SAFE_COMPONENT.fullmatch(session.name) is None:
        raise ValueError("observation session root is invalid")
    source = (session / "media" / source_id).resolve(strict=True)
    if not source.is_dir() or not source.is_relative_to(session):
        raise SessionIntegrityError("camera source escapes its observation session")
    epoch = (source / f"epoch-{codec_epoch}").resolve(strict=True)
    if not epoch.is_dir() or epoch.parent != source:
        raise SessionIntegrityError("camera epoch escapes its source")

    inspected = inspect_recorded_media_epoch(
        epoch,
        expected_source_name=source_id,
        origin_epoch_ns=origin_epoch_ns,
        origin_monotonic_ns=origin_monotonic_ns,
    )
    file_records = _epoch_file_records(epoch, inspected.segments)
    input_document = _input_document(
        session_id=session.name,
        source_id=source_id,
        codec_epoch=codec_epoch,
        inspected=inspected,
        file_records=file_records,
    )
    input_sha256 = hashlib.sha256(_canonical_json(input_document)).hexdigest()
    job_id = f"recorded-camera-{input_sha256[:24]}"
    manifest = _job_manifest(job_id, input_sha256, input_document)

    root = _prepare_private_directory(output_root)
    final = root / job_id
    if final.exists():
        existing = validate_camera_compute_job(final)
        if existing.input_sha256 != input_sha256:
            raise SessionIntegrityError("compute job id collides with different input")
        return existing

    staging = root / f".tmp-{secrets.token_hex(16)}"
    staging.mkdir(mode=0o700)
    published = False
    try:
        destination_epoch = staging / "input" / "camera" / source_id / f"epoch-{codec_epoch}"
        destination_epoch.mkdir(mode=0o700, parents=True)
        (destination_epoch / "segments").mkdir(mode=0o700)
        for record in file_records:
            relative = PurePosixPath(str(record["path"]))
            source_file = epoch / Path(*relative.parts[4:])
            destination = staging / Path(*relative.parts)
            _copy_regular_file(source_file, destination)

        copied = inspect_recorded_media_epoch(
            destination_epoch,
            expected_source_name=source_id,
            origin_epoch_ns=origin_epoch_ns,
            origin_monotonic_ns=origin_monotonic_ns,
        )
        copied_records = _epoch_file_records(destination_epoch, copied.segments)
        copied_input = _input_document(
            session_id=session.name,
            source_id=source_id,
            codec_epoch=codec_epoch,
            inspected=copied,
            file_records=copied_records,
        )
        if _canonical_json(copied_input) != _canonical_json(input_document):
            raise SessionIntegrityError("camera compute staging copy changed")

        write_json_atomic(staging / "job.json", manifest)
        _fsync_tree(staging)
        os.replace(staging, final)
        _fsync_directory(root)
        published = True
    finally:
        if not published and staging.exists():
            shutil.rmtree(staging)

    return validate_camera_compute_job(final)


def validate_camera_compute_job(job_root: Path) -> CameraComputeJob:
    """Validate a published compute job without trusting local paths in JSON."""

    root = job_root.expanduser().resolve(strict=True)
    if not root.is_dir() or _SAFE_JOB_ID.fullmatch(root.name) is None:
        raise SessionIntegrityError("compute job root is invalid")
    manifest_path = root / "job.json"
    manifest = _read_json_object(manifest_path, root)
    if (
        manifest.get("schema_version") != COMPUTE_JOB_SCHEMA
        or manifest.get("job_id") != root.name
        or manifest.get("profile") != COMPUTE_PROFILE
        or manifest.get("result_contract")
        != {
            "schema_version": COMPUTE_RESULT_SCHEMA,
            "timestamp_basis": "session-time-seconds",
        }
    ):
        raise SessionIntegrityError("compute job manifest is incompatible")
    input_sha256 = manifest.get("input_sha256")
    input_document = manifest.get("input")
    if (
        not isinstance(input_sha256, str)
        or _SHA256.fullmatch(input_sha256) is None
        or not isinstance(input_document, dict)
        or hashlib.sha256(_canonical_json(input_document)).hexdigest() != input_sha256
        or root.name != f"recorded-camera-{input_sha256[:24]}"
    ):
        raise SessionIntegrityError("compute job input generation is inconsistent")

    source_id = input_document.get("source_id")
    session_id = input_document.get("session_id")
    codec_epoch = input_document.get("codec_epoch")
    timeline = input_document.get("timeline")
    files = input_document.get("files")
    if (
        not isinstance(session_id, str)
        or _SAFE_COMPONENT.fullmatch(session_id) is None
        or not isinstance(source_id, str)
        or _SAFE_COMPONENT.fullmatch(source_id) is None
        or not isinstance(codec_epoch, int)
        or isinstance(codec_epoch, bool)
        or codec_epoch < 1
        or not isinstance(timeline, dict)
        or not isinstance(files, list)
        or not 4 <= len(files) <= MAX_JOB_FILES
    ):
        raise SessionIntegrityError("compute job input descriptor is invalid")

    expected_prefix = PurePosixPath("input") / "camera" / source_id / f"epoch-{codec_epoch}"
    seen: set[str] = set()
    total_bytes = 0
    for record in files:
        if not isinstance(record, dict):
            raise SessionIntegrityError("compute job file descriptor is invalid")
        relative = _safe_relative_path(record.get("path"), expected_prefix)
        encoded = relative.as_posix()
        if encoded in seen:
            raise SessionIntegrityError("compute job file descriptor is duplicated")
        seen.add(encoded)
        expected_bytes = record.get("byte_length")
        expected_sha256 = record.get("sha256")
        if (
            not isinstance(expected_bytes, int)
            or isinstance(expected_bytes, bool)
            or expected_bytes < 1
            or not isinstance(expected_sha256, str)
            or _SHA256.fullmatch(expected_sha256) is None
        ):
            raise SessionIntegrityError("compute job file identity is invalid")
        payload_path = root / Path(*relative.parts)
        metadata = _confined_regular_file(payload_path, root)
        if metadata.st_size != expected_bytes or _sha256_file(payload_path) != expected_sha256:
            raise SessionIntegrityError("compute job payload identity changed")
        total_bytes += expected_bytes

    required = {
        (expected_prefix / "summary.json").as_posix(),
        (expected_prefix / "index.jsonl").as_posix(),
        (expected_prefix / "init.mp4").as_posix(),
    }
    segment_count = input_document.get("segment_count")
    if not isinstance(segment_count, int) or isinstance(segment_count, bool) or segment_count < 1:
        raise SessionIntegrityError("compute job segment count is invalid")
    required.update(
        (expected_prefix / "segments" / f"{sequence}.m4s").as_posix()
        for sequence in range(1, segment_count + 1)
    )
    if seen != required or input_document.get("byte_length") != total_bytes:
        raise SessionIntegrityError("compute job payload set is inconsistent")

    start = timeline.get("start_seconds")
    end = timeline.get("end_seconds")
    if not isinstance(start, (int, float)) or not isinstance(end, (int, float)):
        raise SessionIntegrityError("compute job timeline is invalid")
    return CameraComputeJob(
        job_id=root.name,
        job_root=root,
        manifest_path=manifest_path,
        session_id=session_id,
        source_id=source_id,
        codec_epoch=codec_epoch,
        input_sha256=input_sha256,
        input_byte_length=total_bytes,
        segment_count=segment_count,
        timeline_start_seconds=float(start),
        timeline_end_seconds=float(end),
    )


def _input_document(
    *,
    session_id: str,
    source_id: str,
    codec_epoch: int,
    inspected: Any,
    file_records: list[dict[str, Any]],
) -> dict[str, Any]:
    summary_record = next(
        record for record in file_records if record["path"].endswith("summary.json")
    )
    index_record = next(record for record in file_records if record["path"].endswith("index.jsonl"))
    return {
        "kind": "canonical-camera-epoch",
        "session_id": session_id,
        "source_id": source_id,
        "codec_epoch": codec_epoch,
        "synchronization": "host-arrival-best-effort",
        "media_type": inspected.media_type,
        "timeline": {
            "basis": "session-time-seconds",
            "start_seconds": inspected.timeline_start_seconds,
            "end_seconds": inspected.timeline_end_seconds,
        },
        "segment_count": len(inspected.segments),
        "byte_length": sum(int(record["byte_length"]) for record in file_records),
        "archive_summary_sha256": summary_record["sha256"],
        "archive_index_sha256": index_record["sha256"],
        "files": file_records,
    }


def _job_manifest(job_id: str, input_sha256: str, input_document: dict[str, Any]) -> dict[str, Any]:
    return {
        "schema_version": COMPUTE_JOB_SCHEMA,
        "job_id": job_id,
        "profile": COMPUTE_PROFILE,
        "input_sha256": input_sha256,
        "input": input_document,
        "result_contract": {
            "schema_version": COMPUTE_RESULT_SCHEMA,
            "timestamp_basis": "session-time-seconds",
        },
    }


def _epoch_file_records(epoch: Path, segments: tuple[Any, ...]) -> list[dict[str, Any]]:
    prefix = PurePosixPath("input") / "camera" / epoch.parent.name / epoch.name
    paths = [epoch / "summary.json", epoch / "index.jsonl", epoch / "init.mp4"]
    paths.extend(segment.path for segment in segments)
    records: list[dict[str, Any]] = []
    for path in paths:
        metadata = _confined_regular_file(path, epoch)
        relative_inside_epoch = path.relative_to(epoch)
        records.append(
            {
                "path": (prefix / relative_inside_epoch.as_posix()).as_posix(),
                "byte_length": metadata.st_size,
                "sha256": _sha256_file(path),
            }
        )
    return records


def _prepare_private_directory(path: Path) -> Path:
    candidate = path.expanduser()
    candidate.mkdir(mode=0o700, parents=True, exist_ok=True)
    metadata = candidate.lstat()
    if stat.S_ISLNK(metadata.st_mode) or not stat.S_ISDIR(metadata.st_mode):
        raise ValueError("compute output root must be a real directory")
    root = candidate.resolve(strict=True)
    os.chmod(root, 0o700)
    return root


def _copy_regular_file(source: Path, destination: Path) -> None:
    confinement_root = source.parents[1] if source.parent.name == "segments" else source.parent
    _confined_regular_file(source, confinement_root)
    destination.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
    with source.open("rb") as reader, destination.open("xb") as writer:
        shutil.copyfileobj(reader, writer, length=1024 * 1024)
        writer.flush()
        os.fsync(writer.fileno())
    os.chmod(destination, 0o600)


def _read_json_object(path: Path, root: Path) -> dict[str, Any]:
    metadata = _confined_regular_file(path, root)
    if not 0 < metadata.st_size <= MAX_JOB_MANIFEST_BYTES:
        raise SessionIntegrityError("compute job manifest is outside bounds")
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise SessionIntegrityError("compute job manifest is unavailable") from exc
    if not isinstance(value, dict):
        raise SessionIntegrityError("compute job manifest is not an object")
    return value


def _safe_relative_path(value: object, expected_prefix: PurePosixPath) -> PurePosixPath:
    if not isinstance(value, str) or not value or "\\" in value:
        raise SessionIntegrityError("compute job payload path is invalid")
    relative = PurePosixPath(value)
    if (
        relative.is_absolute()
        or ".." in relative.parts
        or relative.parts[: len(expected_prefix.parts)] != expected_prefix.parts
    ):
        raise SessionIntegrityError("compute job payload path escapes its root")
    return relative


def _confined_regular_file(path: Path, root: Path) -> os.stat_result:
    try:
        resolved_root = root.resolve(strict=True)
        resolved = path.resolve(strict=True)
        metadata = path.lstat()
    except OSError as exc:
        raise SessionIntegrityError("compute job payload is unavailable") from exc
    if (
        stat.S_ISLNK(metadata.st_mode)
        or not stat.S_ISREG(metadata.st_mode)
        or not resolved.is_relative_to(resolved_root)
    ):
        raise SessionIntegrityError("compute job payload is not a confined regular file")
    return metadata


def _sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        while chunk := stream.read(1024 * 1024):
            digest.update(chunk)
    return digest.hexdigest()


def _canonical_json(value: object) -> bytes:
    try:
        payload = json.dumps(
            value,
            ensure_ascii=False,
            sort_keys=True,
            separators=(",", ":"),
            allow_nan=False,
        ).encode("utf-8")
    except (TypeError, ValueError) as exc:
        raise SessionIntegrityError("compute job descriptor cannot be encoded") from exc
    if not 0 < len(payload) <= MAX_JOB_MANIFEST_BYTES:
        raise SessionIntegrityError("compute job descriptor is outside bounds")
    return payload


def _fsync_tree(root: Path) -> None:
    directories = sorted(
        (path for path in root.rglob("*") if path.is_dir()),
        key=lambda path: len(path.parts),
        reverse=True,
    )
    for directory in directories:
        _fsync_directory(directory)
    _fsync_directory(root)


def _fsync_directory(path: Path) -> None:
    descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0))
    try:
        os.fsync(descriptor)
    finally:
        os.close(descriptor)
