#!/usr/bin/env python3
"""
Build a complete scenario directory (configuration.json + logic/*.py) from outputs/configuration.json.

Usage:
    python3 build_scenario.py --out outputs/scenario_run --overwrite

With process spec (uses LLM-generated physics instead of IR heuristics for HIL):
    python3 build_scenario.py --out outputs/scenario_run --process-spec outputs/process_spec.json --overwrite
"""

import argparse
import json
import shutil
import subprocess
import sys
from pathlib import Path
from typing import List, Set, Tuple


def get_logic_files_from_config(config_path: Path) -> Tuple[Set[str], Set[str]]:
    """
    Extract logic filenames referenced in configuration.json.

    Returns: (plc_logic_files, hil_logic_files)
    """
    config = json.loads(config_path.read_text(encoding="utf-8"))
    plc_files: Set[str] = set()
    hil_files: Set[str] = set()

    for plc in config.get("plcs", []):
        logic = plc.get("logic", "")
        if logic:
            plc_files.add(logic)

    for hil in config.get("hils", []):
        logic = hil.get("logic", "")
        if logic:
            hil_files.add(logic)

    return plc_files, hil_files


def verify_logic_files_exist(config_path: Path, logic_dir: Path) -> List[str]:
    """
    Verify all logic files referenced in config exist in logic_dir.

    Returns: list of missing file error messages (empty if all OK)
    """
    plc_files, hil_files = get_logic_files_from_config(config_path)
    all_files = plc_files | hil_files

    errors: List[str] = []
    for fname in sorted(all_files):
        fpath = logic_dir / fname
        if not fpath.exists():
            errors.append(f"Missing logic file: {fpath} (referenced in config)")

    return errors


def run_command(cmd: list[str], description: str) -> None:
    """Run a command and exit on failure."""
    print(f"\n{'='*60}")
    print(f"{description}")
    print(f"{'='*60}")
    print(f"$ {' '.join(cmd)}")
    result = subprocess.run(cmd)
    if result.returncode != 0:
        raise SystemExit(f"ERROR: {description} failed with code {result.returncode}")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Build scenario directory: config.json + IR + logic/*.py"
    )
    parser.add_argument(
        "--config",
        default="outputs/configuration.json",
        help="Input configuration.json (default: outputs/configuration.json)",
    )
    parser.add_argument(
        "--out",
        default="outputs/scenario_run",
        help="Output scenario directory (default: outputs/scenario_run)",
    )
    parser.add_argument(
        "--ir-file",
        default="outputs/ir/ir_v1.json",
        help="Intermediate IR file (default: outputs/ir/ir_v1.json)",
    )
    parser.add_argument(
        "--model",
        default="tank",
        choices=["tank"],
        help="Heuristic model for IR generation",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Overwrite existing files",
    )
    parser.add_argument(
        "--process-spec",
        default=None,
        help="Path to process_spec.json for HIL physics (optional, replaces IR-based HIL)",
    )
    parser.add_argument(
        "--skip-semantic",
        action="store_true",
        help="Skip semantic validation in config pipeline (for debugging)",
    )
    args = parser.parse_args()

    config_path = Path(args.config)
    out_dir = Path(args.out)
    ir_path = Path(args.ir_file)
    logic_dir = out_dir / "logic"
    process_spec_path = Path(args.process_spec) if args.process_spec else None

    # Validate input
    if not config_path.exists():
        raise SystemExit(f"ERROR: Configuration file not found: {config_path}")

    if process_spec_path and not process_spec_path.exists():
        raise SystemExit(f"ERROR: Process spec file not found: {process_spec_path}")

    print(f"\n{'#'*60}")
    print(f"# Building scenario: {out_dir}")
    print(f"# Using Python: {sys.executable}")
    print(f"{'#'*60}")

    # Step 0: Build and validate configuration (normalize -> enrich -> semantic validate)
    # Output enriched config to scenario output directory
    enriched_config_path = out_dir / "configuration.json"
    out_dir.mkdir(parents=True, exist_ok=True)
    cmd0 = [
        sys.executable,
        "-m",
        "tools.build_config",
        "--config",
        str(config_path),
        "--out-dir",
        str(out_dir),
        "--overwrite",
    ]
    if args.skip_semantic:
        cmd0.append("--skip-semantic")
    run_command(cmd0, "Step 0: Build and validate configuration")

    # Use enriched config for subsequent steps
    config_path = enriched_config_path

    # Step 1: Create IR from configuration.json
    ir_path.parent.mkdir(parents=True, exist_ok=True)
    cmd1 = [
        sys.executable,
        "-m",
        "tools.make_ir_from_config",
        "--config",
        str(config_path),
        "--out",
        str(ir_path),
        "--model",
        args.model,
    ]
    if args.overwrite:
        cmd1.append("--overwrite")
    run_command(cmd1, "Step 1: Generate IR from configuration.json")

    # Step 2: Compile IR to logic/*.py files
    logic_dir.mkdir(parents=True, exist_ok=True)
    cmd2 = [
        sys.executable,
        "-m",
        "tools.compile_ir",
        "--ir",
        str(ir_path),
        "--out-dir",
        str(logic_dir),
    ]
    if args.overwrite:
        cmd2.append("--overwrite")
    run_command(cmd2, "Step 2: Compile IR to logic/*.py files")

    # Step 2b (optional): Compile process_spec.json to HIL logic (replaces IR-generated HIL)
    if process_spec_path:
        # Get HIL logic filename from config
        _, hil_files = get_logic_files_from_config(config_path)
        if not hil_files:
            print("WARNING: No HIL logic files referenced in config, skipping process spec compilation")
        else:
            # Use first HIL logic filename (typically there's only one HIL)
            hil_logic_name = sorted(hil_files)[0]
            hil_logic_out = logic_dir / hil_logic_name

            cmd2b = [
                sys.executable,
                "-m",
                "tools.compile_process_spec",
                "--spec",
                str(process_spec_path),
                "--out",
                str(hil_logic_out),
                "--config",
                str(config_path),  # Pass config to initialize all HIL output keys
                "--overwrite",  # Always overwrite to replace IR-generated HIL
            ]
            run_command(cmd2b, f"Step 2b: Compile process_spec.json to {hil_logic_name}")

    # Step 3: Validate logic files
    cmd3 = [
        sys.executable,
        "-m",
        "tools.validate_logic",
        "--config",
        str(config_path),
        "--logic-dir",
        str(logic_dir),
        "--check-callbacks",
        "--check-hil-init",
    ]
    run_command(cmd3, "Step 3: Validate generated logic files")

    # Step 4: Verify all logic files referenced in config exist
    print(f"\n{'='*60}")
    print(f"Step 4: Verify all referenced logic files exist")
    print(f"{'='*60}")
    out_config = out_dir / "configuration.json"
    verify_errors = verify_logic_files_exist(out_config, logic_dir)
    if verify_errors:
        print("ERRORS:")
        for err in verify_errors:
            print(f"  - {err}")
        raise SystemExit("ERROR: Missing logic files. Scenario incomplete.")
    else:
        plc_files, hil_files = get_logic_files_from_config(out_config)
        print(f"  PLC logic files: {sorted(plc_files)}")
        print(f"  HIL logic files: {sorted(hil_files)}")
        print("  All logic files present: OK")

    # Summary
    print(f"\n{'#'*60}")
    print(f"# SUCCESS: Scenario built at {out_dir}")
    print(f"{'#'*60}")
    print(f"\nScenario contents:")
    print(f"  - {out_dir / 'configuration.json'}")
    print(f"  - {logic_dir}/")

    logic_files = sorted(logic_dir.glob("*.py"))
    for f in logic_files:
        print(f"      {f.name}")

    print(f"\nTo run with ICS-SimLab:")
    print(f"  cd ~/projects/ICS-SimLab-main/curtin-ics-simlab")
    print(f"  sudo ./start.sh {out_dir.absolute()}")


if __name__ == "__main__":
    main()