from __future__ import annotations

import argparse
import hashlib
import json
import os
import pathlib
import secrets
import sys
import time
from typing import Any, Iterable


ROOT = pathlib.Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

from reliability_core.archive_r10r2 import source_tree_manifest
from reliability_core.process_control import run_supervised_process
from reliability_core.stage_control_r10r3r2 import (
    KNOWN_FAILURE_TEST_IDS,
    DeadlineCalibration,
    canary_selection,
    checkpoint_payload,
    environment_sha256,
    execution_cost_p95,
    group_admission,
    local_cleanup_calibration,
    observed_execution_seconds,
)
from reliability_core.suite_isolation import (
    calibrate_python_startup,
    discover_test_ids,
    load_policy_manifest,
    policy_index,
    validate_policy_manifest,
)
from reliability_core.supervisor_r10r2 import (
    EventJournal,
    _start_load_controller,
    canonical_hash,
    load_controller_cleanup_budget,
    load_controller_startup_budget,
    load_json_lines,
    stop_load_controller,
)


PROFILE_COST_SECONDS = {
    "shared_safe": 0.25,
    "adaptive_process": 6.0,
    "isolated_batch": 12.0,
    "isolated_heavy": 45.0,
}


def sha256_file(path: pathlib.Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def atomic_json(path: pathlib.Path, value: object) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8", newline="\n")
    os.replace(temporary, path)


def classify_failure(test_id: str, record: dict[str, Any]) -> str:
    if record.get("orphan_process_count") or record.get("zombie_process_count"):
        return "CLEANUP"
    if record.get("timed_out"):
        return "PHASE_STALL"
    mapping = {
        "test_r5_rg_03": "PROCESS_TREE_CAPTURE",
        "test_r6_rg_05": "OUTPUT_DRAIN",
        "test_r6_rg_06": "CLEANUP",
        "test_r6_rg_07": "PROCESS_TREE_CAPTURE",
        "test_r8_rg_03": "READY_IDENTITY_RACE",
        "test_r8_rg_06": "COLD_START_BUDGET",
    }
    return next((category for suffix, category in mapping.items() if test_id.endswith(suffix)), "ORDER_DEPENDENCY")


def provisional_calibration() -> DeadlineCalibration:
    startup = calibrate_python_startup(ROOT)
    cleanup = local_cleanup_calibration(sys.executable)
    execution = tuple(PROFILE_COST_SECONDS["isolated_heavy"] for _ in range(3))
    return DeadlineCalibration.from_samples(
        startup=startup.samples_seconds,
        execution=execution,
        drain=cleanup.drain_samples_seconds,
        cleanup=cleanup.cleanup_samples_seconds,
        environment_sha256=environment_sha256(),
    )


def load_calibration(path: pathlib.Path | None) -> DeadlineCalibration:
    if path is None:
        return provisional_calibration()
    payload = json.loads(path.read_text(encoding="utf-8"))
    for key in (
        "startup_samples_seconds",
        "execution_samples_seconds",
        "drain_samples_seconds",
        "cleanup_samples_seconds",
    ):
        payload[key] = tuple(payload[key])
    value = DeadlineCalibration(**payload)
    if value.environment_sha256 != environment_sha256():
        raise ValueError("calibration environment digest mismatch")
    return value


def selected_tests(stage: str, manifest: dict[str, Any]) -> list[str]:
    discovered = discover_test_ids(ROOT)
    if stage in {"smoke", "repeat", "mixed"}:
        return list(KNOWN_FAILURE_TEST_IDS)
    if stage == "canary":
        return canary_selection(discovered, policy_index(manifest))
    if stage == "full":
        return discovered
    raise ValueError(stage)


def test_plan(test_ids: Iterable[str], repeat: int, order: str, seed: int | None) -> list[str]:
    rows = list(test_ids)
    if order == "reverse":
        rows.reverse()
    elif order == "random":
        import random
        random.Random(seed).shuffle(rows)
    return [test_id for _ in range(repeat) for test_id in rows]


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--stage", choices=("smoke", "repeat", "mixed", "canary", "full"), required=True)
    parser.add_argument("--repeat", type=int, default=1)
    parser.add_argument("--order", choices=("default", "reverse", "random"), default="default")
    parser.add_argument("--seed", type=int)
    parser.add_argument("--load-profile", choices=("normal", "cpu-medium"), default="normal")
    parser.add_argument("--calibration", type=pathlib.Path)
    parser.add_argument("--output-dir", type=pathlib.Path, required=True)
    parser.add_argument("--workspec-sha256", required=True)
    parser.add_argument("--matrix-contract-sha256", required=True)
    args = parser.parse_args()
    if args.repeat < 1:
        raise SystemExit("repeat must be positive")
    if args.stage == "smoke" and args.repeat != 1:
        raise SystemExit("smoke repeat must be one")
    if args.stage == "repeat" and args.repeat not in {3, 10, 20}:
        raise SystemExit("repeat stage must use 3, 10, or 20")
    if args.order == "random" and args.seed is None:
        raise SystemExit("random order requires seed")

    output = args.output_dir.resolve()
    output.mkdir(parents=True, exist_ok=False)
    manifest = load_policy_manifest(ROOT / "test_execution_manifest.json")
    discovered = discover_test_ids(ROOT)
    manifest_validation = validate_policy_manifest(manifest, discovered)
    if not manifest_validation["passed"]:
        raise SystemExit("policy manifest validation failed")
    policies = policy_index(manifest)
    calibration = load_calibration(args.calibration)
    phase = calibrate_python_startup(ROOT)
    phase_json = json.dumps(phase.to_dict(), sort_keys=True)
    calibration_json = json.dumps(calibration.to_dict(), sort_keys=True)
    source = source_tree_manifest(ROOT)
    chosen = selected_tests(args.stage, manifest)
    plan = test_plan(chosen, args.repeat, args.order, args.seed)
    plan_digest = canonical_hash(plan)
    run_id = f"r10r3r2-{args.stage}-{os.getpid()}-{time.time_ns()}"
    challenge = hashlib.sha256((secrets.token_hex(32) + run_id).encode("utf-8")).hexdigest()
    journal_path = output / "event_journal.jsonl"
    journal = EventJournal(journal_path, run_id, challenge)
    journal.append("STAGE_START", stage=args.stage, repeat=args.repeat, planned_count=len(plan), plan_digest=plan_digest)
    records: list[dict[str, Any]] = []
    failure_category: str | None = None
    started = time.monotonic()
    expected_cost = sum(
        calibration.admission_cost(execution_cost_p95(
            calibration,
            PROFILE_COST_SECONDS.get(str(policies[test_id]["timeout_profile"]), calibration.execution_p95_seconds),
        ))
        for test_id in plan
    ) + calibration.reserve_seconds
    deadline = started + expected_cost

    load_process = None
    load_record: dict[str, Any] = {
        "level": "normal", "passed": True, "cleanup_passed": True,
        "orphan_process_count": 0, "zombie_process_count": 0,
    }
    if args.load_profile == "cpu-medium":
        load_startup_budget = load_controller_startup_budget(
            historical_startup_p95_seconds=calibration.startup_p95_seconds,
            current_phase_bootstrap_budget_seconds=phase.bootstrap_budget_seconds,
        )
        load_process, load_record = _start_load_controller(
            ROOT,
            output / "load_controller",
            startup_timeout_seconds=load_startup_budget,
        )
        journal.append(
            "LOAD_CONTROLLER_START",
            load_pid=load_record.get("load_pid"),
            ready_handshake=load_record.get("ready_handshake"),
        )
    load_ready = load_record.get("passed") is True or (
        (load_record.get("ready_handshake") or {}).get("ready") is True
    )

    for index, test_id in enumerate(plan):
        if not load_ready:
            failure_category = "PHASE_STALL"
            journal.append("LOAD_CONTROLLER_REJECTED", index=index, record=load_record)
            break
        remaining = deadline - time.monotonic()
        execution_cost = execution_cost_p95(
            calibration,
            PROFILE_COST_SECONDS.get(str(policies[test_id]["timeout_profile"]), calibration.execution_p95_seconds),
        )
        admission = group_admission(remaining, calibration, execution_cost)
        if not admission["admitted"]:
            failure_category = "GROUP_ADMISSION"
            journal.append("GROUP_REJECTED", test_id=test_id, index=index, admission=admission)
            break
        raw_stem = f"{index:04d}-{test_id.rsplit('.', 1)[-1]}"
        stdout_path = output / f"{raw_stem}.stdout.txt"
        stderr_path = output / f"{raw_stem}.stderr.txt"
        journal.append("TEST_START", test_id=test_id, index=index)
        environment = {
            **os.environ,
            "PYTHONDONTWRITEBYTECODE": "1",
            "WRLAB_TEST_WORKER": "1",
            "WRLAB_TEST_GROUP": raw_stem,
            "WRLAB_TEST_FILTER": json.dumps([test_id]),
            "WRLAB_PHASE_BUDGET_JSON": phase_json,
            "R10_R3_DEADLINE_CALIBRATION_JSON": calibration_json,
            "WRLAB_ORDER_NAME": args.order,
            "WRLAB_RANDOM_SEED": "" if args.seed is None else str(args.seed),
            "WRLAB_LOAD_PROFILE": args.load_profile,
        }
        record = run_supervised_process(
            [sys.executable, str(ROOT / "run_all.py"), "--test-worker"],
            ROOT,
            environment=environment,
            total_timeout_seconds=admission["required_seconds"] - calibration.cleanup_p95_seconds - calibration.finalization_reserve_seconds,
            max_output_bytes=4 * 1024 * 1024,
        )
        stdout_path.write_text(str(record.pop("stdout")), encoding="utf-8", newline="\n")
        stderr_path.write_text(str(record.pop("stderr")), encoding="utf-8", newline="\n")
        worker_lines = [
            json.loads(line.split(" ", 1)[1])
            for line in stdout_path.read_text(encoding="utf-8").splitlines()
            if line.startswith("TEST_WORKER_REPORT ")
        ]
        worker = worker_lines[-1] if worker_lines else {}
        passed = (
            record.get("exit_code") == 0
            and record.get("timed_out") is False
            and record.get("orphan_process_count") == 0
            and record.get("zombie_process_count") == 0
            and record.get("output_drain_threads_stopped") is True
            and worker.get("passed") is True
            and worker.get("tests_run") == 1
        )
        row = {
            **record,
            "index": index,
            "test_id": test_id,
            "challenge_sha256": hashlib.sha256((challenge + str(index)).encode("utf-8")).hexdigest(),
            "admission": admission,
            "stdout_file": stdout_path.name,
            "stderr_file": stderr_path.name,
            "stdout_sha256": sha256_file(stdout_path),
            "stderr_sha256": sha256_file(stderr_path),
            "worker_report": worker,
            "passed": passed,
        }
        records.append(row)
        journal.append("TEST_END", test_id=test_id, index=index, passed=passed, exit_code=record.get("exit_code"))
        if not passed:
            failure_category = classify_failure(test_id, row)
            journal.append("FAIL_FAST", test_id=test_id, index=index, failure_category=failure_category)
            break

    if load_process is not None:
        load_cleanup_seconds = load_controller_cleanup_budget(
            load_record,
            minimum_seconds=(
                calibration.cleanup_p95_seconds * 3.0
                + calibration.finalization_reserve_seconds
            ),
        )
        load_cleanup_deadline = time.monotonic() + load_cleanup_seconds
        load_record["cleanup_budget_seconds"] = load_cleanup_seconds
        load_record.update(stop_load_controller(load_process, load_record, load_cleanup_deadline))
        journal.append(
            "LOAD_CONTROLLER_END",
            passed=load_record.get("passed"),
            observed_busy_fraction=load_record.get("observed_busy_fraction"),
            orphan_process_count=load_record.get("orphan_process_count"),
            zombie_process_count=load_record.get("zombie_process_count"),
        )
        if load_record.get("passed") is not True and failure_category is None:
            failure_category = "CLEANUP"

    execution_samples = tuple(observed_execution_seconds(row) for row in records if row.get("passed"))
    if len(execution_samples) >= 3:
        measured = DeadlineCalibration.from_samples(
            startup=phase.samples_seconds,
            execution=execution_samples,
            drain=calibration.drain_samples_seconds,
            cleanup=calibration.cleanup_samples_seconds,
            environment_sha256=environment_sha256(),
        )
    else:
        measured = calibration
    completed = len(records)
    passed = completed == len(plan) and all(row["passed"] for row in records) and load_record.get("passed") is True
    journal.append("STAGE_END", passed=passed, completed_count=completed, planned_count=len(plan), failure_category=failure_category)
    journal_rows = load_json_lines(journal_path)
    raw_digest = canonical_hash([
        {"test_id": row["test_id"], "stdout_sha256": row["stdout_sha256"], "stderr_sha256": row["stderr_sha256"], "passed": row["passed"]}
        for row in records
    ])
    stage_id = {"smoke": 1, "repeat": 2, "mixed": 3, "canary": 4, "full": 5}[args.stage]
    checkpoint = checkpoint_payload(
        stage_id=stage_id,
        source_tree_sha256=source["tree_sha256"],
        result_zip_sha256=None,
        workspec_sha256=args.workspec_sha256,
        matrix_contract_sha256=args.matrix_contract_sha256,
        completed_static_gates=("stage0",),
        completed_final_rows=(),
        next_action="promote" if passed else "rerun-minimal-failed-stage",
        invalidated_scopes=() if passed else (args.stage,),
        last_failure_category=failure_category,
    )
    report = {
        "stage_id": stage_id,
        "stage": args.stage,
        "run_id": run_id,
        "source_tree_sha256": source["tree_sha256"],
        "challenge_sha256": challenge,
        "order": args.order,
        "seed": args.seed,
        "load_profile": args.load_profile,
        "repeat": args.repeat,
        "selected_test_count": len(chosen),
        "planned_execution_count": len(plan),
        "completed_execution_count": completed,
        "plan_digest": plan_digest,
        "raw_digest": raw_digest,
        "journal_sha256": canonical_hash(journal_rows),
        "failure_category": failure_category,
        "records": records,
        "load_controller_record": load_record,
        "load_controller_record_sha256": canonical_hash(load_record),
        "calibration": measured.to_dict(),
        "passed": passed,
    }
    atomic_json(output / "stage_report.json", report)
    atomic_json(output / "calibration.json", measured.to_dict())
    atomic_json(output / "checkpoint.json", checkpoint)
    print(json.dumps({
        "stage": args.stage,
        "passed": passed,
        "planned": len(plan),
        "completed": completed,
        "failure_category": failure_category,
        "source_tree_sha256": source["tree_sha256"],
        "raw_digest": raw_digest,
    }, ensure_ascii=False, sort_keys=True))
    return 0 if passed else 1


if __name__ == "__main__":
    raise SystemExit(main())
