from __future__ import annotations

import hashlib
import json
import math
import os
import pathlib
import random
import re
import signal
import statistics
import subprocess
import sys
import tempfile
import time
import unittest
from dataclasses import asdict, dataclass
from typing import Any, Iterable, Mapping, Sequence

from .stage_control_r10r3r2 import cleanup_wait_seconds


SAFE_SHARED_WORKER = "SAFE_SHARED_WORKER"
ISOLATED_TEST_PROCESS = "ISOLATED_TEST_PROCESS"
ISOLATED_BATCH_CONTROLLER = "ISOLATED_BATCH_CONTROLLER"
EXTERNAL_PLATFORM_ONLY = "EXTERNAL_PLATFORM_ONLY"
ALLOWED_POLICIES = {
    SAFE_SHARED_WORKER,
    ISOLATED_TEST_PROCESS,
    ISOLATED_BATCH_CONTROLLER,
    EXTERNAL_PLATFORM_ONLY,
}
REQUIRED_ENTRY_FIELDS = {
    "test_id",
    "policy",
    "reason",
    "timeout_profile",
    "creates_process",
    "creates_thread",
    "uses_signal",
    "uses_global_cache",
    "mutates_environment",
    "required_cleanup_assertions",
    "requirement_ids",
}
TIMEOUT_PROFILE_SECONDS = {
    "shared_safe": 0.25,
    "adaptive_process": 6.0,
    "isolated_batch": 12.0,
    "isolated_heavy": 45.0,
}


@dataclass(frozen=True)
class StartupCalibration:
    samples_seconds: tuple[float, ...]
    sample_pids: tuple[int, ...]
    sample_process_identities: tuple[str, ...]
    median_seconds: float
    p95_seconds: float
    bootstrap_budget_seconds: float
    ready_budget_seconds: float
    execution_budget_seconds: float
    cleanup_reserve_seconds: float
    hard_max_bootstrap_seconds: float
    passed: bool

    def to_dict(self) -> dict[str, Any]:
        value = asdict(self)
        value["samples_seconds"] = list(self.samples_seconds)
        value["sample_pids"] = list(self.sample_pids)
        value["sample_process_identities"] = list(self.sample_process_identities)
        return value


def _flatten(suite: unittest.TestSuite) -> Iterable[unittest.TestCase]:
    for item in suite:
        if isinstance(item, unittest.TestSuite):
            yield from _flatten(item)
        else:
            yield item


def discover_test_ids(root: pathlib.Path) -> list[str]:
    suite = unittest.defaultTestLoader.discover(str(root / "tests"), top_level_dir=str(root))
    return sorted(test.id() for test in _flatten(suite))


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


def load_policy_manifest(path: pathlib.Path) -> dict[str, Any]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(payload, dict) or not isinstance(payload.get("tests"), list):
        raise ValueError("test execution manifest must contain a tests list")
    return payload


def validate_policy_manifest(payload: dict[str, Any], discovered_ids: Iterable[str]) -> dict[str, Any]:
    discovered = list(discovered_ids)
    rows = payload.get("tests") if isinstance(payload, dict) else None
    rows = rows if isinstance(rows, list) else []
    row_ids = [row.get("test_id") for row in rows if isinstance(row, dict)]
    duplicates = sorted({test_id for test_id in row_ids if isinstance(test_id, str) and row_ids.count(test_id) > 1})
    manifest_ids = {test_id for test_id in row_ids if isinstance(test_id, str)}
    discovered_set = set(discovered)
    missing = sorted(discovered_set - manifest_ids)
    unknown = sorted(manifest_ids - discovered_set)
    malformed: list[dict[str, Any]] = []
    unsafe_sensitive: list[str] = []
    for index, row in enumerate(rows):
        if not isinstance(row, dict):
            malformed.append({"index": index, "reason": "not_object"})
            continue
        absent = sorted(REQUIRED_ENTRY_FIELDS - set(row))
        if absent or row.get("policy") not in ALLOWED_POLICIES:
            malformed.append({"index": index, "test_id": row.get("test_id"), "missing": absent, "policy": row.get("policy")})
        sensitive = any(bool(row.get(field)) for field in ("creates_process", "creates_thread", "uses_signal", "uses_global_cache", "mutates_environment"))
        if sensitive and row.get("policy") == SAFE_SHARED_WORKER:
            unsafe_sensitive.append(str(row.get("test_id")))
        cleanup = row.get("required_cleanup_assertions")
        if not isinstance(cleanup, list) or not cleanup:
            malformed.append({"index": index, "test_id": row.get("test_id"), "reason": "cleanup_assertions_required"})
    checks = {
        "manifest_version": payload.get("version") in {"R9", "R10_R2", "R10_R3_R2"},
        "coverage_exact": not missing and not unknown,
        "duplicates_zero": not duplicates,
        "malformed_zero": not malformed,
        "unsafe_sensitive_zero": not unsafe_sensitive,
        "count_match": len(rows) == len(discovered),
    }
    return {
        "passed": all(checks.values()),
        "checks": checks,
        "discovered_count": len(discovered),
        "manifest_count": len(rows),
        "coverage_percent": round(100.0 * len(discovered_set & manifest_ids) / max(1, len(discovered_set)), 3),
        "missing_test_ids": missing,
        "unknown_test_ids": unknown,
        "duplicate_test_ids": duplicates,
        "malformed_entries": malformed,
        "unsafe_sensitive_test_ids": sorted(unsafe_sensitive),
    }


def policy_index(payload: dict[str, Any]) -> dict[str, dict[str, Any]]:
    return {str(row["test_id"]): row for row in payload["tests"] if isinstance(row, dict) and isinstance(row.get("test_id"), str)}


def ordered_test_ids(test_ids: Iterable[str], policies: dict[str, dict[str, Any]], order_name: str, random_seed: int | None = None) -> list[str]:
    ordered = sorted(test_ids)
    if order_name == "default":
        return ordered
    if order_name == "reverse":
        return list(reversed(ordered))
    if order_name == "random":
        if random_seed is None:
            raise ValueError("random order requires a seed")
        random.Random(random_seed).shuffle(ordered)
        return ordered
    if order_name in {"sensitive-first", "sensitive-last"}:
        sensitive = [item for item in ordered if policies[item]["policy"] != SAFE_SHARED_WORKER]
        safe = [item for item in ordered if policies[item]["policy"] == SAFE_SHARED_WORKER]
        return sensitive + safe if order_name == "sensitive-first" else safe + sensitive
    raise ValueError(f"unsupported test order: {order_name}")


def derive_phase_budget(samples_seconds: Iterable[float], hard_max_bootstrap_seconds: float = 8.0) -> StartupCalibration:
    samples = tuple(round(float(value), 6) for value in samples_seconds)
    if not samples or any(value <= 0 for value in samples) or hard_max_bootstrap_seconds <= 0:
        raise ValueError("positive calibration samples and hard maximum are required")
    ordered = sorted(samples)
    p95 = ordered[min(len(ordered) - 1, max(0, math.ceil(len(ordered) * 0.95) - 1))]
    median = statistics.median(ordered)
    bootstrap = min(float(hard_max_bootstrap_seconds), max(1.0, p95 * 1.6 + 0.5))
    ready = min(12.0, max(2.0, bootstrap + 1.0))
    return StartupCalibration(
        samples_seconds=samples,
        sample_pids=(),
        sample_process_identities=(),
        median_seconds=round(median, 6),
        p95_seconds=round(p95, 6),
        bootstrap_budget_seconds=round(bootstrap, 6),
        ready_budget_seconds=round(ready, 6),
        execution_budget_seconds=120.0,
        cleanup_reserve_seconds=2.0,
        hard_max_bootstrap_seconds=float(hard_max_bootstrap_seconds),
        passed=bootstrap <= hard_max_bootstrap_seconds and bootstrap > p95,
    )


def calibrate_python_startup(root: pathlib.Path, sample_count: int = 3, environment: dict[str, str] | None = None) -> StartupCalibration:
    if sample_count < 3:
        raise ValueError("at least three cold-start samples are required")
    launcher = root / "tests" / "fixtures" / "minimal_bootstrap_launcher.py"
    target = root / "tests" / "fixtures" / "noop_target.py"
    samples: list[float] = []
    pids: list[int] = []
    process_identities: list[str] = []
    with tempfile.TemporaryDirectory(prefix="r9-calibration-") as directory:
        base = pathlib.Path(directory)
        for index in range(sample_count):
            boot = base / f"boot-{index}.json"
            started = time.monotonic()
            process = subprocess.Popen(
                [sys.executable, str(launcher), "--bootstrapped-path", str(boot), "--target", str(target)],
                cwd=root,
                stdin=subprocess.DEVNULL,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
                start_new_session=os.name != "nt",
                creationflags=getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0) if os.name == "nt" else 0,
                env={**os.environ, "PYTHONDONTWRITEBYTECODE": "1", **(environment or {})},
            )
            deadline = started + 8.0
            while not boot.is_file() and process.poll() is None and time.monotonic() < deadline:
                time.sleep(0.005)
            elapsed = time.monotonic() - started
            try:
                process.wait(timeout=max(0.1, deadline - time.monotonic()))
            except subprocess.TimeoutExpired:
                process.kill()
                process.wait(timeout=cleanup_wait_seconds())
            if not boot.is_file() or process.returncode != 0:
                raise RuntimeError("cold-start calibration launcher failed")
            payload = json.loads(boot.read_text(encoding="utf-8"))
            if payload.get("status") != "BOOTSTRAPPED" or payload.get("pid") != process.pid:
                raise RuntimeError("cold-start calibration event identity mismatch")
            event_epoch = payload.get("event_epoch")
            if not isinstance(event_epoch, (int, float)) or event_epoch <= 0:
                raise RuntimeError("cold-start calibration event lacks a process-start identity")
            samples.append(elapsed)
            pids.append(process.pid)
            process_identities.append(f"{process.pid}:{event_epoch!r}")
    observed_max = max(samples)
    observed_jitter = observed_max - min(samples)
    measured_margin = max(
        observed_max * 0.25,
        observed_jitter * 2.0,
        time.get_clock_info("monotonic").resolution * 4.0,
    )
    derived = derive_phase_budget(
        samples,
        hard_max_bootstrap_seconds=max(8.0, observed_max + measured_margin),
    )
    return StartupCalibration(**{
        **asdict(derived),
        "sample_pids": tuple(pids),
        "sample_process_identities": tuple(process_identities),
        "passed": derived.passed and len(set(process_identities)) == sample_count,
    })


def worker_calibration_test_ids(root: pathlib.Path) -> tuple[str, ...]:
    """Select one deterministic SAFE test from every discovered test module."""
    manifest = load_policy_manifest(root / "test_execution_manifest.json")
    selected_by_module: dict[str, str] = {}
    for row in sorted(manifest["tests"], key=lambda value: str(value["test_id"])):
        if row["policy"] != SAFE_SHARED_WORKER:
            continue
        test_id = str(row["test_id"])
        module = test_id.rsplit(".", 2)[0]
        selected_by_module.setdefault(module, test_id)
    if not selected_by_module:
        raise RuntimeError("worker calibration requires at least one SAFE test module")
    return tuple(sorted(selected_by_module.values()))


def calibrate_worker_import_discovery(root: pathlib.Path, sample_count: int = 3) -> StartupCalibration:
    if sample_count < 3:
        raise ValueError("at least three fresh worker import/discovery samples are required")
    from .process_control import run_supervised_process

    calibration_test_ids = worker_calibration_test_ids(root)
    samples: list[float] = []
    pids: list[int] = []
    process_identities: list[str] = []
    if len(calibration_test_ids) >= sample_count:
        schedule = [
            (representative_index % sample_count, test_id)
            for representative_index, test_id in enumerate(calibration_test_ids)
        ]
    else:
        schedule = [
            (sample_index, test_id)
            for sample_index in range(sample_count)
            for test_id in calibration_test_ids
        ]
    observed_by_sample: list[list[tuple[float, int, str]]] = [[] for _ in range(sample_count)]
    for operation_index, (sample_index, test_id) in enumerate(schedule):
        environment = {
            **os.environ,
            "PYTHONDONTWRITEBYTECODE": "1",
            "WRLAB_TEST_WORKER": "1",
            "WRLAB_TEST_GROUP": f"import-calibration-{sample_index}-{operation_index}",
            "WRLAB_TEST_FILTER": json.dumps((test_id,)),
        }
        record = run_supervised_process(
            [sys.executable, "-B", str(root / "run_all.py"), "--test-worker"],
            root,
            environment=environment,
            total_timeout_seconds=30.0,
            max_output_bytes=1024 * 1024,
        )
        reports = [
            json.loads(line.split(" ", 1)[1])
            for line in str(record["stdout"]).splitlines()
            if line.startswith("TEST_WORKER_REPORT ")
        ]
        if (
            record["exit_code"] != 0
            or record["timed_out"] is not False
            or record["orphan_process_count"] != 0
            or record["zombie_process_count"] != 0
            or record["output_drain_threads_stopped"] is not True
            or record["stderr"]
            or not reports
            or reports[-1].get("tests_run") != 1
            or reports[-1].get("passed") is not True
        ):
            raise RuntimeError("worker import/discovery calibration failed")
        root_identity = record.get("root_process_identity")
        if (
            not isinstance(root_identity, Mapping)
            or int(root_identity.get("pid", -1)) != int(record["pid"])
            or not str(root_identity.get("started", "")).strip()
        ):
            raise RuntimeError("worker import/discovery calibration lacks a root process identity")
        process_identity = f"{int(root_identity['pid'])}:{root_identity['started']}"
        observed_by_sample[sample_index].append(
            (float(record["elapsed_ms"]) / 1000.0, int(record["pid"]), process_identity)
        )
    for observed_records in observed_by_sample:
        slowest_elapsed, slowest_pid, slowest_identity = max(observed_records, key=lambda value: value[0])
        samples.append(slowest_elapsed)
        pids.append(slowest_pid)
        process_identities.append(slowest_identity)
    derived = derive_phase_budget(samples)
    return StartupCalibration(**{
        **asdict(derived),
        "sample_pids": tuple(pids),
        "sample_process_identities": tuple(process_identities),
        "passed": derived.passed and len(set(process_identities)) == sample_count,
    })


def calibrated_worker_process_floor(calibration: StartupCalibration) -> float:
    """Return a measured worker-lifecycle floor with a two-sided jitter envelope."""
    samples = tuple(float(value) for value in calibration.samples_seconds)
    if not calibration.passed or len(samples) < 3 or any(value <= 0 for value in samples):
        raise ValueError("worker process floor requires a passing fresh calibration")
    observed_jitter = max(samples) - min(samples)
    return float(calibration.p95_seconds) + observed_jitter * 2.0


def phase_budget_from_environment(root: pathlib.Path) -> StartupCalibration:
    encoded = os.environ.get("WRLAB_PHASE_BUDGET_JSON")
    if encoded:
        payload = json.loads(encoded)
        required = set(StartupCalibration.__dataclass_fields__)
        if set(payload) != required:
            raise ValueError("run-scoped phase budget is incomplete")
        payload["samples_seconds"] = tuple(payload["samples_seconds"])
        payload["sample_pids"] = tuple(payload["sample_pids"])
        payload["sample_process_identities"] = tuple(payload["sample_process_identities"])
        value = StartupCalibration(**payload)
        if not value.passed:
            raise ValueError("run-scoped phase budget did not pass calibration")
        return value
    return calibrate_python_startup(root)


def timeout_profile_budget(
    timeout_profiles: Iterable[str],
    *,
    shared_safe_observed_p95_seconds: float | None = None,
) -> float:
    """Return execution cost without treating a profile constant as observed runtime.

    A shared worker calibration is a fresh, supervised end-to-end lifecycle sample.
    Its p95 therefore provides a measured lower bound for each shared-safe test when
    the controller allocates a worker deadline.  The static profile remains only a
    fallback for callers that do not own a run-scoped calibration.
    """
    if shared_safe_observed_p95_seconds is not None:
        observed = float(shared_safe_observed_p95_seconds)
        if not math.isfinite(observed) or observed <= 0:
            raise ValueError("shared-safe observed p95 must be finite and positive")
    else:
        observed = TIMEOUT_PROFILE_SECONDS["shared_safe"]
    return sum(
        max(TIMEOUT_PROFILE_SECONDS["shared_safe"], observed)
        if str(profile) == "shared_safe"
        else TIMEOUT_PROFILE_SECONDS.get(str(profile), 8.0)
        for profile in timeout_profiles
    )


def partition_timeout_bounded_tests(
    test_ids: Sequence[str],
    timeout_profile_by_test_id: Mapping[str, str],
    *,
    maximum_execution_cost_seconds: float,
    worker_process_floor_seconds: float,
    shared_safe_observed_p95_seconds: float,
) -> tuple[tuple[str, ...], ...]:
    """Preserve test order while preventing one shared worker from owning the deadline.

    The ceiling is supplied by the run-scoped supervisor budget.  Static timeout
    profiles remain conservative execution estimates; this function only divides a
    contiguous shared-safe segment and never discounts, drops, or reorders a test.
    """
    ceiling = float(maximum_execution_cost_seconds)
    process_floor = float(worker_process_floor_seconds)
    if not math.isfinite(ceiling) or ceiling <= 0:
        raise ValueError("maximum execution cost must be finite and positive")
    if not math.isfinite(process_floor) or process_floor < 0 or process_floor >= ceiling:
        raise ValueError("worker process floor must be finite and below the ceiling")
    ordered = tuple(str(test_id) for test_id in test_ids)
    if len(set(ordered)) != len(ordered):
        raise ValueError("test ids must be unique within a safe segment")
    missing_profiles = [test_id for test_id in ordered if test_id not in timeout_profile_by_test_id]
    if missing_profiles:
        raise ValueError("every test requires a timeout profile")

    chunks: list[tuple[str, ...]] = []
    current: list[str] = []
    current_cost = process_floor
    for test_id in ordered:
        test_cost = timeout_profile_budget(
            (timeout_profile_by_test_id[test_id],),
            shared_safe_observed_p95_seconds=shared_safe_observed_p95_seconds,
        )
        if current and current_cost + test_cost > ceiling:
            chunks.append(tuple(current))
            current = []
            current_cost = process_floor
        current.append(test_id)
        current_cost += test_cost
    if current:
        chunks.append(tuple(current))
    if tuple(test_id for chunk in chunks for test_id in chunk) != ordered:
        raise RuntimeError("safe segment partition changed test identity or order")
    return tuple(chunks)


def suite_total_budget(calibration: StartupCalibration, isolated_count: int, safe_count: int, hard_max_seconds: float = 900.0, timeout_profiles: Iterable[str] | None = None) -> float:
    profiles = list(timeout_profiles or [])
    workload_reserve = timeout_profile_budget(profiles) if profiles else safe_count * .25 + isolated_count * 6.0
    estimate = 120.0 + workload_reserve + isolated_count * calibration.p95_seconds + calibration.cleanup_reserve_seconds
    return round(min(hard_max_seconds, max(360.0, estimate)), 3)


def environment_digest(environment: dict[str, str]) -> str:
    safe = {key: environment[key] for key in sorted(environment) if key.startswith(("PYTHON", "WRLAB_")) and "KEY" not in key and "TOKEN" not in key and "SECRET" not in key}
    return hashlib.sha256(json.dumps(safe, sort_keys=True).encode("utf-8")).hexdigest()


def deadline_allocation_valid(total_seconds: float, elapsed_seconds: float, cleanup_reserve_seconds: float) -> bool:
    return total_seconds > 0 and elapsed_seconds >= 0 and cleanup_reserve_seconds > 0 and elapsed_seconds <= total_seconds - cleanup_reserve_seconds


def required_order_matrix() -> list[dict[str, Any]]:
    return [
        {"order_name": "default", "random_seed": None, "minimum_repeats": 2},
        {"order_name": "reverse", "random_seed": None, "minimum_repeats": 2},
        {"order_name": "random", "random_seed": 901, "minimum_repeats": 2},
        {"order_name": "random", "random_seed": 902, "minimum_repeats": 2},
        {"order_name": "random", "random_seed": 903, "minimum_repeats": 2},
        {"order_name": "sensitive-first", "random_seed": None, "minimum_repeats": 1},
        {"order_name": "sensitive-last", "random_seed": None, "minimum_repeats": 1},
    ]


def parse_unittest_count(transcript: str) -> int | None:
    matches = re.findall(r"Ran\s+(\d+)\s+tests?\s+in", transcript)
    return int(matches[-1]) if matches else None


def validate_matrix_records(records: Iterable[dict[str, Any]], expected_test_count: int, source_hash: str, manifest_hash: str) -> dict[str, Any]:
    rows = list(records)
    invalid: list[dict[str, Any]] = []
    seen: set[str] = set()
    for index, row in enumerate(rows):
        checks = {
            "unique_run_id": isinstance(row.get("run_id"), str) and row.get("run_id") not in seen,
            "exit_zero": row.get("exit_code") == 0,
            "not_timed_out": row.get("timed_out") is False,
            "test_count": row.get("test_count") == expected_test_count,
            "cleanup": row.get("orphan_process_count") == 0 and row.get("zombie_process_count") == 0,
            "source_hash": row.get("source_tree_sha256") == source_hash,
            "manifest_hash": row.get("policy_manifest_sha256") == manifest_hash,
            "origin": row.get("source_origin") in {"active", "candidate", "final", "rollback"},
            "order": row.get("order_name") in {"default", "reverse", "random", "sensitive-first", "sensitive-last"},
        }
        if isinstance(row.get("run_id"), str):
            seen.add(row["run_id"])
        if not all(checks.values()):
            invalid.append({"index": index, "run_id": row.get("run_id"), "checks": checks})
    return {"passed": bool(rows) and not invalid, "record_count": len(rows), "invalid": invalid, "expected_test_count": expected_test_count}
