from __future__ import annotations

import hashlib
import json
import math
import os
import pathlib
import statistics
import subprocess
import sys
import time
from dataclasses import asdict, dataclass
from functools import lru_cache
from typing import Any, Iterable, Mapping, Sequence


STAGE_NAMES = (
    "STRUCTURAL_PREFLIGHT",
    "KNOWN_FAILURE_SMOKE",
    "KNOWN_FAILURE_STABILITY",
    "MIXED_INTERFERENCE",
    "DETERMINISTIC_CANARY",
    "LINUX313_DEFAULT_COLD_FULL",
    "LINUX313_AUTHORITY_ROWS",
    "PLATFORM_VERSION_COVERAGE",
    "ROLLBACK_AND_SEAL",
    "EXACT_FINAL_REPLAY_24",
)

KNOWN_FAILURE_TEST_IDS = (
    "tests.test_r5_rg.R5RegressionTests.test_r5_rg_03",
    "tests.test_r6_rg.R6RegressionTests.test_r6_rg_05",
    "tests.test_r6_rg.R6RegressionTests.test_r6_rg_06",
    "tests.test_r6_rg.R6RegressionTests.test_r6_rg_07",
    "tests.test_r8_rg.R8RegressionTests.test_r8_rg_03",
    "tests.test_r8_rg.R8RegressionTests.test_r8_rg_06",
)

REPEAT_LADDER = (1, 3, 10, 20)
NON_REUSABLE_EVIDENCE = frozenset(
    {
        "test_result",
        "pid",
        "ready",
        "timeout",
        "stdout",
        "stderr",
        "cleanup",
        "challenge",
        "supervisor_record",
        "event_journal",
        "matrix_row",
        "rollback_result",
    }
)


def _p95(values: Sequence[float]) -> float:
    if not values:
        raise ValueError("p95 requires at least one observation")
    ordered = sorted(float(value) for value in values)
    index = max(0, math.ceil(len(ordered) * 0.95) - 1)
    return ordered[index]


@dataclass(frozen=True)
class DeadlineCalibration:
    startup_samples_seconds: tuple[float, ...]
    execution_samples_seconds: tuple[float, ...]
    drain_samples_seconds: tuple[float, ...]
    cleanup_samples_seconds: tuple[float, ...]
    startup_p95_seconds: float
    execution_p95_seconds: float
    drain_p95_seconds: float
    cleanup_p95_seconds: float
    finalization_reserve_seconds: float
    observed_at_epoch: float
    environment_sha256: str

    @classmethod
    def from_samples(
        cls,
        *,
        startup: Sequence[float],
        execution: Sequence[float],
        drain: Sequence[float],
        cleanup: Sequence[float],
        environment_sha256: str,
        observed_at_epoch: float | None = None,
    ) -> "DeadlineCalibration":
        groups = (startup, execution, drain, cleanup)
        if any(len(group) < 3 for group in groups):
            raise ValueError("each deadline phase requires at least three fresh observations")
        if any(value < 0 for group in groups for value in group):
            raise ValueError("deadline samples cannot be negative")
        startup_p95 = _p95(startup)
        execution_p95 = _p95(execution)
        drain_p95 = _p95(drain)
        cleanup_p95 = _p95(cleanup)
        resolution = time.get_clock_info("monotonic").resolution
        # Reserve several observable clock ticks without turning a coarse
        # Windows clock (15.625 ms on older interpreters) into a 15.625 s
        # deadline gap.  Drain and cleanup observations still dominate when
        # they are slower than the clock floor.
        finalization_reserve = max(resolution * 4.0, drain_p95 * 0.25, cleanup_p95 * 0.25)
        return cls(
            tuple(float(value) for value in startup),
            tuple(float(value) for value in execution),
            tuple(float(value) for value in drain),
            tuple(float(value) for value in cleanup),
            startup_p95,
            execution_p95,
            drain_p95,
            cleanup_p95,
            finalization_reserve,
            time.time() if observed_at_epoch is None else float(observed_at_epoch),
            environment_sha256,
        )

    @property
    def reserve_seconds(self) -> float:
        resolution = time.get_clock_info("monotonic").resolution
        return (
            max(self.drain_p95_seconds, resolution)
            + max(self.cleanup_p95_seconds, resolution)
            + self.finalization_reserve_seconds
        )

    def admission_cost(self, execution_cost_seconds: float | None = None) -> float:
        execution = self.execution_p95_seconds if execution_cost_seconds is None else float(execution_cost_seconds)
        return self.startup_p95_seconds + execution + self.reserve_seconds

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass(frozen=True)
class DeadlineHierarchy:
    reviewer_seconds: float
    supervisor_seconds: float
    controller_seconds: float
    worker_seconds: float
    reserve_seconds: float

    @property
    def passed(self) -> bool:
        tolerance = max(abs(self.reserve_seconds) * 1e-12, time.get_clock_info("monotonic").resolution)
        return (
            self.reviewer_seconds > self.supervisor_seconds
            > self.controller_seconds > self.worker_seconds > 0
            and self.reviewer_seconds - self.supervisor_seconds + tolerance >= self.reserve_seconds
            and self.supervisor_seconds - self.controller_seconds + tolerance >= self.reserve_seconds
            and self.controller_seconds - self.worker_seconds + tolerance >= self.reserve_seconds
        )

    def to_dict(self) -> dict[str, Any]:
        return {**asdict(self), "passed": self.passed}


def allocate_deadlines(reviewer_seconds: float, calibration: DeadlineCalibration) -> DeadlineHierarchy:
    reserve = calibration.reserve_seconds
    worker = reviewer_seconds - reserve * 3.0
    controller = worker + reserve
    supervisor = controller + reserve
    return DeadlineHierarchy(float(reviewer_seconds), supervisor, controller, worker, reserve)


def group_admission(
    remaining_seconds: float,
    calibration: DeadlineCalibration,
    execution_cost_seconds: float | None = None,
) -> dict[str, Any]:
    required = calibration.admission_cost(execution_cost_seconds)
    return {
        "remaining_seconds": float(remaining_seconds),
        "required_seconds": required,
        "startup_p95_seconds": calibration.startup_p95_seconds,
        "execution_cost_p95_seconds": calibration.execution_p95_seconds if execution_cost_seconds is None else float(execution_cost_seconds),
        "drain_p95_seconds": calibration.drain_p95_seconds,
        "cleanup_p95_seconds": calibration.cleanup_p95_seconds,
        "finalization_reserve_seconds": calibration.finalization_reserve_seconds,
        "admitted": float(remaining_seconds) >= required,
    }


def execution_cost_p95(calibration: DeadlineCalibration, profile_floor_seconds: float) -> float:
    if profile_floor_seconds < 0:
        raise ValueError("profile floor cannot be negative")
    return max(calibration.execution_p95_seconds, float(profile_floor_seconds))


def observed_execution_seconds(record: Mapping[str, Any]) -> float:
    if record.get("elapsed_seconds") is not None:
        value = float(record["elapsed_seconds"])
    elif record.get("elapsed_ms") is not None:
        value = float(record["elapsed_ms"]) / 1000.0
    else:
        raise ValueError("execution record has no elapsed duration")
    if value <= 0:
        raise ValueError("execution duration must be positive")
    return value


def environment_sha256(environment: Mapping[str, str] | None = None) -> str:
    source = os.environ if environment is None else environment
    safe = {
        key: source[key]
        for key in sorted(source)
        if key.startswith(("PYTHON", "WRLAB_", "R10_R3_"))
        and not any(token in key.upper() for token in ("KEY", "TOKEN", "SECRET", "PASSWORD"))
    }
    raw = json.dumps(safe, sort_keys=True, separators=(",", ":")).encode("utf-8")
    return hashlib.sha256(raw).hexdigest()


@lru_cache(maxsize=8)
def local_cleanup_calibration(python_executable: str = sys.executable) -> DeadlineCalibration:
    startup: list[float] = []
    execution: list[float] = []
    drain: list[float] = []
    cleanup: list[float] = []
    for _ in range(3):
        started = time.monotonic()
        process = subprocess.Popen(
            [python_executable, "-B", "-c", "import time;time.sleep(60)"],
            stdin=subprocess.DEVNULL,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            start_new_session=os.name != "nt",
            creationflags=getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0) if os.name == "nt" else 0,
        )
        startup.append(time.monotonic() - started)
        execution_started = time.monotonic()
        if process.poll() is not None:
            raise RuntimeError("cleanup calibration process exited before termination")
        execution.append(time.monotonic() - execution_started)
        cleanup_started = time.monotonic()
        process.terminate()
        stdout, stderr = process.communicate(timeout=10)
        cleanup.append(time.monotonic() - cleanup_started)
        drain_started = time.monotonic()
        _ = len(stdout) + len(stderr)
        drain.append(time.monotonic() - drain_started)
    return DeadlineCalibration.from_samples(
        startup=startup,
        execution=execution,
        drain=drain,
        cleanup=cleanup,
        environment_sha256=environment_sha256(),
    )


def load_deadline_calibration() -> DeadlineCalibration:
    encoded = os.environ.get("R10_R3_DEADLINE_CALIBRATION_JSON")
    if not encoded:
        return local_cleanup_calibration(sys.executable)
    payload = json.loads(encoded)
    tuple_fields = {
        "startup_samples_seconds",
        "execution_samples_seconds",
        "drain_samples_seconds",
        "cleanup_samples_seconds",
    }
    for field in tuple_fields:
        payload[field] = tuple(float(value) for value in payload[field])
    calibration = DeadlineCalibration(**payload)
    if any(len(getattr(calibration, field)) < 3 for field in tuple_fields):
        raise ValueError("deadline calibration payload has insufficient samples")
    return calibration


def cleanup_wait_seconds(hard_deadline: float | None = None) -> float:
    calibration = load_deadline_calibration()
    local = local_cleanup_calibration(sys.executable)
    budget = max(
        calibration.cleanup_p95_seconds + calibration.finalization_reserve_seconds,
        local.cleanup_p95_seconds + local.finalization_reserve_seconds,
    )
    if hard_deadline is not None:
        budget = min(budget, max(0.0, hard_deadline - time.monotonic()))
    return max(time.get_clock_info("monotonic").resolution, budget)


def canary_selection(
    discovered_ids: Sequence[str],
    policy_rows: Mapping[str, Mapping[str, Any]],
    minimum_fraction: float = 0.20,
) -> list[str]:
    if not 0 < minimum_fraction <= 1:
        raise ValueError("minimum_fraction must be in (0, 1]")
    required_count = math.ceil(len(discovered_ids) * minimum_fraction)
    sensitive = [
        test_id for test_id in discovered_ids
        if policy_rows[test_id].get("policy") != "SAFE_SHARED_WORKER"
        or any(bool(policy_rows[test_id].get(field)) for field in ("creates_process", "creates_thread", "uses_signal", "mutates_environment"))
        or any(token in test_id for token in ("test_r5_", "test_r6_", "test_r7_", "test_r8_", "test_r9_", "test_r10"))
    ]
    selected = list(dict.fromkeys(sensitive))
    for test_id in discovered_ids:
        if len(selected) >= required_count:
            break
        if test_id not in selected:
            selected.append(test_id)
    return selected


def checkpoint_payload(
    *,
    stage_id: int,
    source_tree_sha256: str,
    result_zip_sha256: str | None,
    workspec_sha256: str,
    matrix_contract_sha256: str,
    completed_static_gates: Iterable[str],
    completed_final_rows: Iterable[Mapping[str, Any]],
    next_action: str,
    invalidated_scopes: Iterable[str],
    last_failure_category: str | None,
) -> dict[str, Any]:
    if not 0 <= stage_id < len(STAGE_NAMES):
        raise ValueError("invalid stage")
    return {
        "stage_id": stage_id,
        "source_tree_sha256": source_tree_sha256,
        "result_zip_sha256": result_zip_sha256,
        "workspec_sha256": workspec_sha256,
        "matrix_contract_sha256": matrix_contract_sha256,
        "completed_static_gates": list(completed_static_gates),
        "completed_final_rows": [dict(row) for row in completed_final_rows],
        "next_action": next_action,
        "invalidated_scopes": list(invalidated_scopes),
        "last_failure_category": last_failure_category,
        "checkpoint_is_evidence": False,
    }


def validate_checkpoint(
    checkpoint: Mapping[str, Any],
    *,
    source_tree_sha256: str,
    result_zip_sha256: str | None,
    workspec_sha256: str,
    matrix_contract_sha256: str,
    raw_root: pathlib.Path,
) -> dict[str, Any]:
    hashes_match = {
        "source": checkpoint.get("source_tree_sha256") == source_tree_sha256,
        "result": checkpoint.get("result_zip_sha256") == result_zip_sha256,
        "workspec": checkpoint.get("workspec_sha256") == workspec_sha256,
        "contract": checkpoint.get("matrix_contract_sha256") == matrix_contract_sha256,
    }
    row_checks: list[dict[str, Any]] = []
    for row in checkpoint.get("completed_final_rows", []):
        relative = pathlib.PurePosixPath(str(row.get("raw_path", "")))
        safe_path = raw_root / pathlib.Path(*relative.parts)
        exists = relative.parts and ".." not in relative.parts and safe_path.is_file()
        digest = hashlib.sha256(safe_path.read_bytes()).hexdigest() if exists else None
        row_checks.append({"run_id": row.get("run_id"), "exists": exists, "digest_match": digest == row.get("raw_sha256")})
    source_changed = not hashes_match["source"]
    result_changed = not hashes_match["result"]
    invalidated = []
    if source_changed:
        invalidated.extend(["static_gates", "candidate_results", "sealed_result", "exact_final_rows"])
    elif result_changed:
        invalidated.append("exact_final_rows")
    return {
        "hashes_match": hashes_match,
        "row_checks": row_checks,
        "invalidated_scopes": invalidated,
        "passed": all(hashes_match.values()) and all(row["exists"] and row["digest_match"] for row in row_checks),
    }
