#!/usr/bin/env python3

import argparse
import os
from pathlib import Path
import re
import shutil
import subprocess
import sys
import time
import uuid

import yaml


TEMPLATE_RE = re.compile(r"{{\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*}}")


def render_template(value: str, variables: dict[str, str]) -> str:
    def replace(match: re.Match[str]) -> str:
        key = match.group(1)
        if key not in variables:
            raise KeyError(f"unknown template variable: {key}")
        return str(variables[key])

    return TEMPLATE_RE.sub(replace, value)


def resolve_mapping(raw: dict | None, variables: dict[str, str]) -> dict[str, str]:
    resolved: dict[str, str] = {}
    if not raw:
        return resolved
    pending = dict(raw)
    while pending:
        progressed = False
        for key, value in list(pending.items()):
            try:
                resolved[key] = render_template(str(value), {**variables, **resolved})
            except KeyError:
                continue
            del pending[key]
            progressed = True
        if not progressed:
            missing = ", ".join(sorted(pending))
            raise RuntimeError(f"could not resolve templated variables: {missing}")
    return resolved


def ensure_parent(path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)


def apply_mode(path: Path, mode: str | None) -> None:
    if not mode:
        return
    path.chmod(int(mode, 8))


def remove_path(path: Path) -> None:
    if not path.exists() and not path.is_symlink():
        return
    if path.is_dir() and not path.is_symlink():
        shutil.rmtree(path)
    else:
        path.unlink()


def run_command(command: str, cwd: Path) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        command,
        cwd=cwd,
        shell=True,
        executable="/bin/bash",
        text=True,
        capture_output=True,
    )


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="run_yaml_scenario.py")
    parser.add_argument("scenario", help="Path to the scenario YAML")
    parser.add_argument(
        "--step-by-step",
        action="store_true",
        help="Pause after each step so the sandbox can be inspected before continuing",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Print commands, stdout, stderr, and inspection hints for each step",
    )
    parser.add_argument(
        "--reset",
        action="store_true",
        help="Remove the scenario cleanup targets before starting",
    )
    return parser


def print_block(label: str, text: str) -> None:
    if not text:
        return
    print(f"{label}:")
    print(text)


def format_duration(seconds: float) -> str:
    if seconds < 1:
        return f"{seconds * 1000:.0f}ms"
    if seconds < 60:
        return f"{seconds:.2f}s"
    minutes = int(seconds // 60)
    remainder = seconds - (minutes * 60)
    return f"{minutes}m {remainder:.2f}s"


def collect_inspect_paths(step: dict, variables: dict[str, str]) -> list[str]:
    paths: list[str] = []

    for entry in step.get("inspect", []):
        paths.append(render_template(str(entry), variables))

    for entry in step.get("mkdirs", []):
        paths.append(render_template(str(entry), variables))

    for item in step.get("copy", []):
        paths.append(render_template(str(item["to"]), variables))

    for item in step.get("render", []):
        paths.append(render_template(str(item["to"]), variables))

    for entry in step.get("assert_files", []):
        paths.append(render_template(str(entry), variables))

    deduped: list[str] = []
    seen: set[str] = set()
    for path in paths:
        if path in seen:
            continue
        seen.add(path)
        deduped.append(path)
    return deduped


def print_step_context(step: dict, variables: dict[str, str], inspect_paths: list[str]) -> None:
    if actor := step.get("actor"):
        print(f"actor: {render_template(str(actor), variables)}")
    if summary := step.get("summary"):
        print(f"what: {render_template(str(summary), variables)}")
    if inspect_paths:
        print("inspect:")
        for path in inspect_paths:
            print(f"  {path}")


def wait_for_step_confirmation(name: str) -> int:
    while True:
        try:
            response = input(f"inspect disk for '{name}', then press Enter to continue, 's' to skip pauses, or 'q' to quit: ")
        except EOFError:
            return 0
        normalized = response.strip().lower()
        if normalized == "":
            return 1
        if normalized == "s":
            return 2
        if normalized == "q":
            return 0
        print("enter, s, or q")


def print_timing_summary(step_timings: list[dict[str, float | str]]) -> None:
    if not step_timings:
        return
    print("== Timing Summary ==")
    total_seconds = sum(float(item["seconds"]) for item in step_timings)
    print(f"total: {format_duration(total_seconds)}")
    for item in sorted(step_timings, key=lambda value: float(value["seconds"]), reverse=True):
        status = str(item["status"])
        print(f"{format_duration(float(item['seconds']))} [{status}] {item['name']}")


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()

    scenario_path = Path(args.scenario).resolve()
    scenario_dir = scenario_path.parent
    repo_root = Path(__file__).resolve().parent.parent

    with scenario_path.open("r", encoding="utf-8") as handle:
        scenario = yaml.safe_load(handle)

    setup = scenario.get("setup", {})
    variables: dict[str, str] = {
        "repo_root": str(repo_root),
        "scenario_dir": str(scenario_dir),
        "assets_dir": str(scenario_dir / "assets"),
        "mode": str(
            os.environ.get("RHO_SCENARIO_MODE")
            or setup.get("default_mode", "default")
        ),
    }

    env_mapping = setup.get("env", {})
    for key, env_name in env_mapping.items():
        variables[key] = os.environ.get(str(env_name), "")

    generated = setup.get("generate", {})
    for key, generator in generated.items():
        override = os.environ.get(f"RHO_SCENARIO_GEN_{key.upper()}")
        if override:
            variables[key] = override
            continue
        if generator != "uuid":
            raise RuntimeError(f"unsupported generator for {key}: {generator}")
        variables[key] = str(uuid.uuid4()).lower()

    variables.update(resolve_mapping(setup.get("vars"), variables))
    variables.update(resolve_mapping(setup.get("paths"), variables))

    cleanup_targets = [
        Path(render_template(str(entry), variables)) for entry in setup.get("cleanup", [])
    ]
    cleanup_on_exit = bool(setup.get("cleanup_on_exit", True))
    step_by_step = args.step_by_step
    verbose = args.verbose or step_by_step

    if args.reset:
        for target in cleanup_targets:
            remove_path(target)

    if "sandbox_dir" in variables:
        Path(variables["sandbox_dir"]).mkdir(parents=True, exist_ok=True)
        print(f"sandbox: {variables['sandbox_dir']}")

    step_timings: list[dict[str, float | str]] = []

    try:
        for index, step in enumerate(scenario.get("steps", []), start=1):
            when = step.get("when")
            if when and str(when) != variables["mode"]:
                continue

            name = step.get("name", f"step-{index}")
            print(f"== {name} ==")
            inspect_paths = collect_inspect_paths(step, variables)
            print_step_context(step, variables, inspect_paths)
            step_started = time.perf_counter()

            for entry in step.get("mkdirs", []):
                Path(render_template(str(entry), variables)).mkdir(parents=True, exist_ok=True)

            for item in step.get("copy", []):
                src = Path(render_template(str(item["from"]), variables))
                dst = Path(render_template(str(item["to"]), variables))
                ensure_parent(dst)
                shutil.copy2(src, dst)
                apply_mode(dst, item.get("mode"))

            for item in step.get("render", []):
                src = Path(render_template(str(item["template"]), variables))
                dst = Path(render_template(str(item["to"]), variables))
                contents = src.read_text(encoding="utf-8")
                rendered = render_template(contents, variables)
                ensure_parent(dst)
                dst.write_text(rendered, encoding="utf-8")
                apply_mode(dst, item.get("mode"))

            if "run" in step:
                command = render_template(str(step["run"]), variables)
                if verbose:
                    print_block("command", command)
                result = run_command(command, repo_root)
                if result.returncode != 0:
                    duration = time.perf_counter() - step_started
                    step_timings.append({"name": name, "seconds": duration, "status": "failed"})
                    print(result.stdout, end="")
                    print(result.stderr, end="", file=sys.stderr)
                    print(f"duration: {format_duration(duration)}")
                    print_timing_summary(step_timings)
                    print(f"step failed with exit code {result.returncode}: {name}", file=sys.stderr)
                    return result.returncode

                stdout = result.stdout.rstrip("\n")
                stderr = result.stderr.rstrip("\n")
                if verbose:
                    print_block("stdout", stdout)
                    print_block("stderr", stderr)
                if "capture" in step:
                    variables[str(step["capture"])] = stdout
                if "expect_stdout" in step:
                    expected = render_template(str(step["expect_stdout"]), variables)
                    if stdout != expected:
                        duration = time.perf_counter() - step_started
                        step_timings.append({"name": name, "seconds": duration, "status": "failed"})
                        print(f"duration: {format_duration(duration)}")
                        print_timing_summary(step_timings)
                        print(
                            f"unexpected stdout for {name}: expected {expected!r}, got {stdout!r}",
                            file=sys.stderr,
                        )
                        return 1

            for entry in step.get("assert_files", []):
                path = Path(render_template(str(entry), variables))
                if not path.is_file():
                    duration = time.perf_counter() - step_started
                    step_timings.append({"name": name, "seconds": duration, "status": "failed"})
                    print(f"duration: {format_duration(duration)}")
                    print_timing_summary(step_timings)
                    print(f"expected file missing: {path}", file=sys.stderr)
                    return 1

            for check in step.get("assert_contains", []):
                path = Path(render_template(str(check["path"]), variables))
                pattern = render_template(str(check["pattern"]), variables)
                contents = path.read_text(encoding="utf-8")
                if pattern not in contents:
                    duration = time.perf_counter() - step_started
                    step_timings.append({"name": name, "seconds": duration, "status": "failed"})
                    print(f"duration: {format_duration(duration)}")
                    print_timing_summary(step_timings)
                    print(f"expected pattern {pattern!r} in {path}", file=sys.stderr)
                    return 1

            duration = time.perf_counter() - step_started
            step_timings.append({"name": name, "seconds": duration, "status": "ok"})
            print(f"duration: {format_duration(duration)}")

            if step_by_step:
                decision = wait_for_step_confirmation(name)
                if decision == 0:
                    print_timing_summary(step_timings)
                    print("stopped by operator", file=sys.stderr)
                    return 130
                if decision == 2:
                    step_by_step = False

        print("PASS: scenario completed")
        print_timing_summary(step_timings)
        return 0
    finally:
        if cleanup_on_exit:
            for target in cleanup_targets:
                remove_path(target)


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