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
from typing import Any, cast

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

from .jobs import CameraComputeJob, validate_camera_compute_job

QUALIFICATION_SLICE_SCHEMA = "missioncore.recorded-qualification-slice/v1"
QUALIFICATION_SLICE_IDENTITY_SCHEMA = "missioncore.recorded-qualification-slice-identity/v1"
QUALIFICATION_POLICY = "uniform-frame-index-full-epoch/v1"
DEFAULT_QUALIFICATION_FRAME_COUNT = 256
MAX_QUALIFICATION_FRAME_COUNT = 4096
MAX_SLICE_MANIFEST_BYTES = 16 * 1024 * 1024
MAX_INDEX_LINE_BYTES = 16 * 1024

_SHA256 = re.compile(r"^[a-f0-9]{64}$")
_SAFE_GENERATION = re.compile(r"^qualification-slice-[a-f0-9]{64}$")


class RecordedQualificationSliceError(RuntimeError):
    """A deterministic recorded-camera qualification slice is invalid."""


@dataclass(frozen=True, slots=True)
class QualificationFrame:
    frame_index: int
    sequence: int
    segment_sha256: str
    host_epoch_ns: int
    host_monotonic_ns: int
    archive_session_monotonic_ns: int


@dataclass(frozen=True, slots=True)
class RecordedQualificationSlice:
    generation_id: str
    root: Path
    manifest_path: Path
    job_id: str
    input_sha256: str
    policy: str
    source_frame_count: int
    frames: tuple[QualificationFrame, ...]


def prepare_recorded_qualification_slice(
    *,
    job_root: Path,
    output_root: Path,
    sample_count: int = DEFAULT_QUALIFICATION_FRAME_COUNT,
) -> RecordedQualificationSlice:
    """Seal an evenly distributed, exact-repeat frame-index slice.

    The selection includes both ends of the epoch and is keyed by the complete
    compute-job input identity. It deliberately selects by decoded frame index;
    exact PTS are attached by the worker after decoding, while archive arrival
    timestamps remain diagnostic metadata only.
    """

    job = validate_camera_compute_job(job_root)
    if (
        not isinstance(sample_count, int)
        or isinstance(sample_count, bool)
        or not 1 <= sample_count <= MAX_QUALIFICATION_FRAME_COUNT
    ):
        raise RecordedQualificationSliceError("qualification sample count is outside bounds")
    admitted_count = min(sample_count, job.segment_count)
    selected_indices = _uniform_indices(job.segment_count, admitted_count)
    index_rows = _read_archive_index(job)
    frames = tuple(_frame_from_index(index_rows[index], index) for index in selected_indices)

    identity = {
        "schema_version": QUALIFICATION_SLICE_IDENTITY_SCHEMA,
        "job_id": job.job_id,
        "input_sha256": job.input_sha256,
        "session_id": job.session_id,
        "source_id": job.source_id,
        "codec_epoch": job.codec_epoch,
        "source_frame_count": job.segment_count,
        "policy": QUALIFICATION_POLICY,
        "requested_sample_count": sample_count,
        "admitted_sample_count": len(frames),
        "selected_frames": [
            {
                "frame_index": frame.frame_index,
                "sequence": frame.sequence,
                "segment_sha256": frame.segment_sha256,
            }
            for frame in frames
        ],
    }
    identity_sha256 = hashlib.sha256(_canonical_json(identity)).hexdigest()
    generation_id = f"qualification-slice-{identity_sha256}"
    root = _prepare_private_directory(output_root)
    final = root / generation_id
    if final.exists():
        existing = validate_recorded_qualification_slice(final, job_root=job.job_root)
        if existing.input_sha256 != job.input_sha256:
            raise RecordedQualificationSliceError(
                "qualification generation collides with another input"
            )
        return existing

    staging = root / f".{generation_id}.{secrets.token_hex(12)}.incomplete"
    published = False
    try:
        staging.mkdir(mode=0o700)
        manifest = {
            "schema_version": QUALIFICATION_SLICE_SCHEMA,
            "generation_id": generation_id,
            "identity_sha256": identity_sha256,
            "identity": identity,
            "created_at_utc": utc_now_iso(),
            "source_timeline": {
                "basis": "session-time-seconds",
                "start_seconds": job.timeline_start_seconds,
                "end_seconds": job.timeline_end_seconds,
                "duration_seconds": job.timeline_end_seconds - job.timeline_start_seconds,
            },
            "frames": [
                {
                    "frame_index": frame.frame_index,
                    "sequence": frame.sequence,
                    "segment_sha256": frame.segment_sha256,
                    "archive_host_epoch_ns": frame.host_epoch_ns,
                    "archive_host_monotonic_ns": frame.host_monotonic_ns,
                    "archive_session_monotonic_ns": frame.archive_session_monotonic_ns,
                    "decoded_session_seconds": None,
                }
                for frame in frames
            ],
            "usage": {
                "selection_basis": "zero-based decoded frame index",
                "worker_timestamp_binding": (
                    "resolve exact best_effort_timestamp_time after full stream decode"
                ),
                "archive_timestamp_status": "host-arrival-best-effort-diagnostic",
                "comparison_contract": (
                    "all E1 variants must use this exact ordered frame list and the same input "
                    "and calibration generations"
                ),
            },
        }
        write_json_atomic(staging / "manifest.json", manifest)
        os.chmod(staging / "manifest.json", 0o600)
        _fsync_directory(staging)
        os.replace(staging, final)
        _fsync_directory(root)
        published = True
    finally:
        if not published and staging.exists():
            shutil.rmtree(staging)

    return validate_recorded_qualification_slice(final, job_root=job.job_root)


def validate_recorded_qualification_slice(
    slice_root: Path,
    *,
    job_root: Path,
) -> RecordedQualificationSlice:
    root = slice_root.expanduser().resolve(strict=True)
    if not root.is_dir() or _SAFE_GENERATION.fullmatch(root.name) is None:
        raise RecordedQualificationSliceError("qualification slice root is invalid")
    manifest = _read_json_object(root / "manifest.json", root)
    identity = manifest.get("identity")
    identity_sha256 = manifest.get("identity_sha256")
    if (
        manifest.get("schema_version") != QUALIFICATION_SLICE_SCHEMA
        or manifest.get("generation_id") != root.name
        or not isinstance(identity, dict)
        or identity.get("schema_version") != QUALIFICATION_SLICE_IDENTITY_SCHEMA
        or not isinstance(identity_sha256, str)
        or _SHA256.fullmatch(identity_sha256) is None
        or root.name != f"qualification-slice-{identity_sha256}"
        or hashlib.sha256(_canonical_json(identity)).hexdigest() != identity_sha256
    ):
        raise RecordedQualificationSliceError("qualification slice identity is inconsistent")
    job = validate_camera_compute_job(job_root)
    if (
        identity.get("job_id") != job.job_id
        or identity.get("input_sha256") != job.input_sha256
        or identity.get("session_id") != job.session_id
        or identity.get("source_id") != job.source_id
        or identity.get("codec_epoch") != job.codec_epoch
        or identity.get("source_frame_count") != job.segment_count
        or identity.get("policy") != QUALIFICATION_POLICY
    ):
        raise RecordedQualificationSliceError("qualification slice is not bound to the job")
    rows = manifest.get("frames")
    selected = identity.get("selected_frames")
    admitted_count = identity.get("admitted_sample_count")
    requested_count = identity.get("requested_sample_count")
    if (
        not isinstance(rows, list)
        or not isinstance(selected, list)
        or rows == []
        or not isinstance(admitted_count, int)
        or isinstance(admitted_count, bool)
        or admitted_count != len(rows)
        or len(selected) != len(rows)
        or not isinstance(requested_count, int)
        or isinstance(requested_count, bool)
        or not 1 <= requested_count <= MAX_QUALIFICATION_FRAME_COUNT
        or admitted_count != min(requested_count, job.segment_count)
    ):
        raise RecordedQualificationSliceError("qualification slice frame count is invalid")
    archive_rows = _read_archive_index(job)
    frames: list[QualificationFrame] = []
    previous_index = -1
    for row, selected_row in zip(rows, selected, strict=True):
        if not isinstance(row, dict) or not isinstance(selected_row, dict):
            raise RecordedQualificationSliceError("qualification frame descriptor is invalid")
        frame_index = row.get("frame_index")
        if (
            not isinstance(frame_index, int)
            or isinstance(frame_index, bool)
            or not previous_index < frame_index < job.segment_count
        ):
            raise RecordedQualificationSliceError("qualification frame order is invalid")
        expected = _frame_from_index(archive_rows[frame_index], frame_index)
        expected_identity = {
            "frame_index": expected.frame_index,
            "sequence": expected.sequence,
            "segment_sha256": expected.segment_sha256,
        }
        expected_row = {
            **expected_identity,
            "archive_host_epoch_ns": expected.host_epoch_ns,
            "archive_host_monotonic_ns": expected.host_monotonic_ns,
            "archive_session_monotonic_ns": expected.archive_session_monotonic_ns,
            "decoded_session_seconds": None,
        }
        if selected_row != expected_identity or row != expected_row:
            raise RecordedQualificationSliceError("qualification frame binding changed")
        frames.append(expected)
        previous_index = frame_index
    if tuple(frame.frame_index for frame in frames) != _uniform_indices(
        job.segment_count,
        admitted_count,
    ):
        raise RecordedQualificationSliceError("qualification selection policy changed")
    return RecordedQualificationSlice(
        generation_id=root.name,
        root=root,
        manifest_path=root / "manifest.json",
        job_id=job.job_id,
        input_sha256=job.input_sha256,
        policy=QUALIFICATION_POLICY,
        source_frame_count=job.segment_count,
        frames=tuple(frames),
    )


def _uniform_indices(source_count: int, sample_count: int) -> tuple[int, ...]:
    if source_count < 1 or not 1 <= sample_count <= source_count:
        raise RecordedQualificationSliceError("uniform selection bounds are invalid")
    if sample_count == 1:
        return (source_count // 2,)
    denominator = sample_count - 1
    indices = tuple(
        (position * (source_count - 1) + denominator // 2) // denominator
        for position in range(sample_count)
    )
    if len(set(indices)) != sample_count or indices[0] != 0 or indices[-1] != source_count - 1:
        raise RecordedQualificationSliceError("uniform selection is not exact")
    return indices


def _read_archive_index(job: CameraComputeJob) -> list[dict[str, Any]]:
    index_path = (
        job.job_root
        / "input"
        / "camera"
        / job.source_id
        / f"epoch-{job.codec_epoch}"
        / "index.jsonl"
    )
    try:
        metadata = index_path.lstat()
        resolved = index_path.resolve(strict=True)
    except OSError as exc:
        raise RecordedQualificationSliceError("camera archive index 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(job.job_root)
    ):
        raise RecordedQualificationSliceError("camera archive index is not confined")
    rows: list[dict[str, Any]] = []
    try:
        with index_path.open("r", encoding="utf-8") as stream:
            for expected_index, line in enumerate(stream):
                if len(line.encode("utf-8")) > MAX_INDEX_LINE_BYTES:
                    raise RecordedQualificationSliceError("camera archive index row is too large")
                value = json.loads(line)
                if not isinstance(value, dict):
                    raise RecordedQualificationSliceError("camera archive index row is invalid")
                if value.get("sequence") != expected_index + 1 or value.get("kind") != "media":
                    raise RecordedQualificationSliceError("camera archive index order changed")
                rows.append(value)
    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise RecordedQualificationSliceError("camera archive index could not be read") from exc
    if len(rows) != job.segment_count:
        raise RecordedQualificationSliceError("camera archive index count changed")
    return rows


def _frame_from_index(row: dict[str, Any], frame_index: int) -> QualificationFrame:
    sequence = row.get("sequence")
    segment_sha256 = row.get("sha256")
    host_epoch_ns = row.get("host_epoch_ns")
    host_monotonic_ns = row.get("host_monotonic_ns")
    session_monotonic_ns = row.get("session_monotonic_ns")
    integers = (sequence, host_epoch_ns, host_monotonic_ns, session_monotonic_ns)
    if (
        not all(
            isinstance(value, int) and not isinstance(value, bool) and value >= 0
            for value in integers
        )
        or sequence != frame_index + 1
        or not isinstance(segment_sha256, str)
        or _SHA256.fullmatch(segment_sha256) is None
    ):
        raise RecordedQualificationSliceError("camera archive frame identity is invalid")
    return QualificationFrame(
        frame_index=frame_index,
        sequence=sequence,
        segment_sha256=segment_sha256,
        host_epoch_ns=cast(int, host_epoch_ns),
        host_monotonic_ns=cast(int, host_monotonic_ns),
        archive_session_monotonic_ns=cast(int, session_monotonic_ns),
    )


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 RecordedQualificationSliceError("qualification output root must be a real directory")
    root = candidate.resolve(strict=True)
    os.chmod(root, 0o700)
    return root


def _read_json_object(path: Path, root: Path) -> dict[str, Any]:
    try:
        metadata = path.lstat()
        resolved = path.resolve(strict=True)
    except OSError as exc:
        raise RecordedQualificationSliceError("qualification manifest 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(root)
        or not 0 < metadata.st_size <= MAX_SLICE_MANIFEST_BYTES
    ):
        raise RecordedQualificationSliceError("qualification manifest is outside bounds")
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise RecordedQualificationSliceError("qualification manifest could not be read") from exc
    if not isinstance(value, dict):
        raise RecordedQualificationSliceError("qualification manifest is not an object")
    return value


def _canonical_json(value: object) -> bytes:
    try:
        return json.dumps(
            value,
            ensure_ascii=False,
            sort_keys=True,
            separators=(",", ":"),
            allow_nan=False,
        ).encode("utf-8")
    except (TypeError, ValueError) as exc:
        raise SessionIntegrityError("qualification identity is not canonical JSON") from exc


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