#!/usr/bin/env python3
"""Immutable Transformer selector registry and failure-continuing driver."""
import argparse
import hashlib
import json
import os
import subprocess
import sys
import time


ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
PROTOCOL = os.path.join(ROOT, "TRANSFORMER_CROSSOVER.md")
RESULT_ROOT = os.path.join(ROOT, "results", "transformer_crossover")
METHODS = (
    "bp", "fa", "dfa", "pepita", "ff", "ep", "dualprop",
    "clean_kp", "sdil",
)
RATE_GRIDS = {
    "bp": (0.0003, 0.001, 0.003),
    "fa": (0.0003, 0.001, 0.003),
    "dfa": (0.0003, 0.001, 0.003),
    "pepita": (0.0001, 0.0003, 0.001),
    "ff": (0.0001, 0.0003, 0.001),
    "ep": (0.0001, 0.0003, 0.001),
    "dualprop": (0.0001, 0.0003, 0.001),
    "clean_kp": (0.0003, 0.001, 0.003),
    "sdil": (0.0003, 0.001, 0.003),
}


def sha256(path):
    digest = hashlib.sha256()
    with open(path, "rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def text_sha256(text):
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def git_output(*args):
    return subprocess.run(
        ["git", *args], cwd=ROOT, check=True, capture_output=True,
        text=True).stdout.strip()


def source_report():
    dirty = git_output(
        "status", "--porcelain", "--untracked-files=no")
    if dirty:
        raise RuntimeError(
            "Transformer crossover requires clean tracked source")
    paths = [
        PROTOCOL,
        os.path.abspath(__file__),
        os.path.join(
            ROOT, "experiments", "analyze_transformer_crossover_p1.py"),
        os.path.join(
            ROOT, "experiments", "transformer_crossover_native.py"),
        os.path.join(
            ROOT, "experiments", "transformer_feedback_smoke.py"),
        os.path.join(ROOT, "sdil", "transformer.py"),
    ]
    for path in paths:
        tracked = subprocess.run(
            ["git", "ls-files", "--error-unmatch",
             os.path.relpath(path, ROOT)],
            cwd=ROOT, capture_output=True).returncode == 0
        if not tracked:
            raise RuntimeError(f"untracked crossover source: {path}")
    freeze = subprocess.run(
        [sys.executable, "-m", "pip", "freeze"],
        check=True, capture_output=True, text=True).stdout
    return {
        "git_commit": git_output("rev-parse", "HEAD"),
        "protocol_path": PROTOCOL,
        "protocol_sha256": sha256(PROTOCOL),
        "tracked_files": {
            os.path.relpath(path, ROOT): sha256(path) for path in paths},
        "python": sys.executable,
        "environment_freeze_sha256": text_sha256(freeze),
        "environment_freeze": freeze.splitlines(),
    }


def rate_tag(rate):
    return f"{rate:g}".replace(".", "p")


def p1_command(method, rate):
    name = f"transformer-p1-{method}-d4-lr{rate_tag(rate)}"
    output = os.path.join(RESULT_ROOT, "p1", name + ".json")
    command = [
        sys.executable,
        "experiments/transformer_crossover_native.py",
        "--method", method,
        "--out", output,
        "--device", "cuda",
        "--depth", "4",
        "--width", "128",
        "--heads", "4",
        "--mlp_ratio", "4",
        "--context_length", "64",
        "--batch_size", "32",
        "--eval_batch_size", "32",
        "--train_steps", "1000",
        "--lr", str(rate),
        "--schedule", "constant",
        "--min_lr", "0.0001",
        "--warmup_steps", "100",
        "--weight_decay", "0.1",
        "--run_seed", "0",
        "--model_seed", "2027",
        "--loader_seed", "0",
        "--negative_seed", "5001",
        "--ep_sign_seed", "5002",
        "--traffic_ratio", "4.0",
        "--ff_threshold", "2.0",
        "--ff_score_from_layer", "1",
        "--ep_beta", "0.5",
        "--ep_dt", "0.5",
        "--ep_free_steps", "20",
        "--ep_nudge_steps", "4",
        "--dp_alpha", "0.0",
        "--dp_beta", "0.1",
        "--dp_inference_passes", "16",
        "--eval_every", "0",
        "--log_every", "50",
        "--max_val_batches", "0",
    ]
    return {
        "stage": "p1",
        "method": method,
        "architecture": "transformer4",
        "rate": rate,
        "experiment_name": name,
        "output": output,
        "timeout_seconds": 48 * 60 * 60,
        "command": command,
    }


def p1_jobs():
    # Interleave the expensive methods so two shards do not inherit all state
    # jobs from contiguous method groups.
    order = (
        "ep", "dualprop", "ff", "pepita", "dfa", "fa",
        "clean_kp", "sdil", "bp",
    )
    jobs = []
    for grid_index in range(3):
        for method in order:
            jobs.append(p1_command(
                method, RATE_GRIDS[method][grid_index]))
    return jobs


def registry_sha256(jobs):
    encoded = json.dumps(
        jobs, sort_keys=True, separators=(",", ":")).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()


def ensure_p1_launch(source, jobs):
    path = os.path.join(RESULT_ROOT, "p1_launch.json")
    expected = {
        "stage": "p1",
        "source": source,
        "registry_sha256": registry_sha256(jobs),
        "num_jobs": len(jobs),
        "allowed_physical_gpus": [5, 7],
    }
    if os.path.isfile(path):
        with open(path, encoding="utf-8") as handle:
            existing = json.load(handle)
        if existing != expected:
            raise RuntimeError("Transformer P1 launch lock drift")
        return path
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w", encoding="utf-8") as handle:
        json.dump(expected, handle, indent=2, sort_keys=True)
        handle.write("\n")
    return path


def assert_source_unchanged(source):
    if git_output("status", "--porcelain", "--untracked-files=no"):
        raise RuntimeError("tracked source changed after launch")
    if git_output("rev-parse", "HEAD") != source["git_commit"]:
        raise RuntimeError("source commit changed after launch")
    for relative, expected_hash in source["tracked_files"].items():
        path = os.path.join(ROOT, relative)
        if sha256(path) != expected_hash:
            raise RuntimeError(
                f"source file changed after launch: {relative}")


def physical_gpu_report(dry_run=False):
    visible = os.environ.get("CUDA_VISIBLE_DEVICES")
    if dry_run:
        return {
            "cuda_visible_devices": visible,
            "physical_gpu_index": None,
            "physical_gpu_uuid": None,
        }
    if visible not in {"5", "7"}:
        raise RuntimeError(
            "formal Transformer jobs require CUDA_VISIBLE_DEVICES=5 or 7")
    query = subprocess.run([
        "nvidia-smi",
        "--query-gpu=index,uuid,name",
        "--format=csv,noheader,nounits",
    ], check=True, capture_output=True, text=True).stdout.splitlines()
    rows = {}
    for line in query:
        index, uuid, name = [value.strip() for value in line.split(",", 2)]
        rows[index] = {"uuid": uuid, "name": name}
    if visible not in rows:
        raise RuntimeError(f"physical GPU {visible} not found")
    return {
        "cuda_visible_devices": visible,
        "physical_gpu_index": int(visible),
        "physical_gpu_uuid": rows[visible]["uuid"],
        "physical_gpu_name": rows[visible]["name"],
    }


def run_job(job, source, gpu, dry_run):
    if not dry_run:
        assert_source_unchanged(source)
    manifest_path = job["output"] + ".manifest.json"
    if os.path.exists(manifest_path):
        print(f"preserving {job['experiment_name']}", flush=True)
        return
    if os.path.exists(job["output"]):
        raise RuntimeError(
            f"orphaned output requires audit: {job['output']}")
    print("RUN", " ".join(job["command"]), flush=True)
    if dry_run:
        return
    os.makedirs(os.path.dirname(job["output"]), exist_ok=True)
    started = time.time()
    try:
        result = subprocess.run(
            job["command"], cwd=ROOT, timeout=job["timeout_seconds"])
        return_code = result.returncode
        status = "completed" if return_code == 0 else "nonzero_exit"
    except subprocess.TimeoutExpired:
        return_code = None
        status = "timeout"
    output_exists = os.path.isfile(job["output"])
    if status == "completed" and not output_exists:
        status = "missing_output"
    manifest = {
        **job,
        "source": source,
        "hardware_lock": gpu,
        "status": status,
        "return_code": return_code,
        "output_exists": output_exists,
        "output_sha256": sha256(job["output"]) if output_exists else None,
        "driver_wall_seconds": time.time() - started,
        "completed_unix_time": time.time(),
    }
    with open(manifest_path, "w", encoding="utf-8") as handle:
        json.dump(manifest, handle, indent=2, sort_keys=True)
        handle.write("\n")
    print(
        f"DONE status={status} {job['experiment_name']}", flush=True)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--stage", choices=("p1",), default="p1")
    parser.add_argument("--shard-index", type=int, default=0)
    parser.add_argument("--num-shards", type=int, default=1)
    parser.add_argument("--method", choices=METHODS)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()
    if not 0 <= args.shard_index < args.num_shards:
        raise ValueError("invalid shard")
    source = source_report()
    gpu = physical_gpu_report(args.dry_run)
    jobs = p1_jobs()
    if not args.dry_run:
        launch = ensure_p1_launch(source, jobs)
        print(f"P1 launch lock: {launch}", flush=True)
    if args.method:
        jobs = [job for job in jobs if job["method"] == args.method]
    jobs = [
        job for index, job in enumerate(jobs)
        if index % args.num_shards == args.shard_index]
    if not jobs:
        raise ValueError("no jobs selected")
    for job in jobs:
        run_job(job, source, gpu, args.dry_run)


if __name__ == "__main__":
    main()