"""Generate smooth wormhole feeds for a suite of seeds.

This utility replays the canonical smooth configuration for the requested seeds
and writes per-seed artefacts (metrics, markdown summary, and wireframe feed)
under the supplied output directory.
"""

from __future__ import annotations

import argparse
import json
from dataclasses import asdict
from pathlib import Path
import sys
from typing import Dict, List

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from run_smooth import build_wireframe_payload, format_markdown  # type: ignore
from wormhole_proof.core.wormhole_smooth import SmoothConfig, run_smooth_experiment

CONFIG_OVERRIDES: Dict[str, float | int] = {
    "epsilon": 0.02,
    "base_step_scale": 0.06,
    "lambda_h": 180.0,
    "wormhole_step_scale": 0.25,
    "verify_tolerance": 0.08,
    "traversable_mu": 0.03,
    "shock_energy": 0.04,
    "epsilon_final": 0.001,
    "base_step_scale_final": 0.04,
    "lambda_h_final": 260.0,
    "wormhole_step_scale_final": 0.18,
    "verify_tolerance_final": 0.045,
    "traversable_mu_final": 0.05,
    "shock_energy_final": 0.06,
    "warmup_steps": 200,
    "steps": 500,
    "horizon_steps": 40,
    "soft_rank": 8,
    "wormhole_budget_fraction": 0.10,
    "lambda_c": 12.0,
    "debt_repay_rate": 0.03,
    "max_packets": 72,
    "traversable_window": 4,
    "hp_message_bits": 3,
    "hp_message_magnitude": 0.26,
    "hp_scramble_wait": 5,
    "hp_decode_packets": 6,
    "budget_total": 1.0,
}


def _instantiate_config() -> SmoothConfig:
    return SmoothConfig(**CONFIG_OVERRIDES)  # type: ignore[arg-type]


def _run_seed(seed: int, output_dir: Path, overwrite: bool) -> dict:
    seed_dir = output_dir / f"seed_{seed}"
    seed_dir.mkdir(parents=True, exist_ok=True)
    feed_path = seed_dir / "wireframe.json"

    if feed_path.exists() and not overwrite:
        metrics_path = seed_dir / "smooth_metrics.json"
        if metrics_path.exists():
            return json.loads(metrics_path.read_text(encoding="utf-8"))

    cfg = _instantiate_config()
    result = run_smooth_experiment(cfg, seed)

    metrics = {
        "seed": seed,
        "metrics": asdict(result.metrics),
    }

    metrics_path = seed_dir / "smooth_metrics.json"
    metrics_path.write_text(json.dumps(metrics, indent=2), encoding="utf-8")

    summary_path = seed_dir / "smooth_summary.md"
    summary = format_markdown([result])
    summary_path.write_text(summary, encoding="utf-8")

    payload = build_wireframe_payload([result])
    feed_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")

    config_path = seed_dir / "config.json"
    config_path.write_text(json.dumps(asdict(cfg), indent=2), encoding="utf-8")

    return metrics


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate smooth wormhole feeds for multiple seeds")
    parser.add_argument(
        "--output",
        type=Path,
        default=ROOT / "artifacts" / "gr_suite",
        help="Directory where per-seed artefacts will be written",
    )
    parser.add_argument(
        "--seeds",
        type=int,
        nargs="*",
        default=[7, 11, 13, 17, 19],
        help="Seeds to simulate",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Re-run simulations even if artefacts already exist",
    )
    args = parser.parse_args()

    args.output.mkdir(parents=True, exist_ok=True)

    suite_metrics: List[dict] = []
    for seed in args.seeds:
        payload = _run_seed(seed, args.output, overwrite=args.overwrite)
        suite_metrics.append(payload)
        print(f"Seed {seed} written to {args.output / f'seed_{seed}'}")

    suite_summary = {
        "config_overrides": CONFIG_OVERRIDES,
        "seeds": suite_metrics,
    }
    summary_path = args.output / "smooth_suite_summary.json"
    summary_path.write_text(json.dumps(suite_summary, indent=2), encoding="utf-8")
    print(f"Suite summary written to {summary_path}")


if __name__ == "__main__":
    main()