#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from propose_self_improvement import ROOT, load_yaml, run_codex_proposal

DEFAULT_CONFIG = ROOT / "self_improve.yaml"
DEFAULT_SCHEMA = ROOT / "schemas" / "self_improve_proposal_v0.json"


def utc_stamp() -> str:
    return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")


def run_command(command: str, *, cwd: Path, stdout_path: Path, stderr_path: Path) -> int:
    completed = subprocess.run(
        command,
        cwd=cwd,
        shell=True,
        text=True,
        stdout=stdout_path.open("w", encoding="utf-8"),
        stderr=stderr_path.open("w", encoding="utf-8"),
        check=False,
    )
    return completed.returncode


def main() -> int:
    parser = argparse.ArgumentParser(description="Propose and optionally apply one bounded self-improvement run.")
    parser.add_argument("--goal", required=True)
    parser.add_argument("--apply", action="store_true")
    parser.add_argument("--stage-on-pass", action="store_true")
    parser.add_argument("--config", default=str(DEFAULT_CONFIG))
    parser.add_argument("--schema", default=str(DEFAULT_SCHEMA))
    parser.add_argument("--output-dir", default="")
    args = parser.parse_args()

    config_path = Path(args.config)
    schema_path = Path(args.schema)
    config = load_yaml(config_path)
    run_dir = Path(args.output_dir) if args.output_dir else (ROOT / "runs" / "self_improve" / utc_stamp())
    run_dir.mkdir(parents=True, exist_ok=True)

    system_context_path = run_dir / "system_context.json"
    subprocess.run(
        [str(ROOT / "api" / "build_system_context.sh"), str(system_context_path)],
        cwd=ROOT,
        check=True,
        stdout=subprocess.DEVNULL,
    )

    proposal_path = run_dir / "proposal.json"
    manifest_path = run_dir / "manifest.json"
    brief_path = run_dir / "brief.md"
    proposal = run_codex_proposal(
        goal=args.goal,
        config_path=config_path,
        system_context_path=system_context_path,
        output_path=proposal_path,
        schema_path=schema_path,
    )
    manifest_path.write_text(json.dumps(proposal["manifest"], indent=2, sort_keys=True) + "\n", encoding="utf-8")
    brief_lines = [
        "# Self-Improve Brief",
        "",
        f"Clean product one-liner: {proposal['one_liner']}",
        "",
        f"- Goal: `{proposal['goal']}`",
        f"- Decision: {proposal['decision_brief']}",
        f"- Target files: `{', '.join(proposal['target_files'])}`",
        f"- Benchmark: `{proposal['benchmark']['command']}`",
        "",
        "## Change Summary",
        "",
    ]
    brief_lines.extend([f"- {item}" for item in proposal["change_summary"]])
    brief_path.write_text("\n".join(brief_lines) + "\n", encoding="utf-8")

    receipt_path = None
    runtime_summary_path = None
    benchmark_status = None
    benchmark_stdout = run_dir / "benchmark.stdout"
    benchmark_stderr = run_dir / "benchmark.stderr"
    staged_files: list[str] = []

    if args.apply:
        runtime = subprocess.run(
            [str(ROOT / "runtime" / "execute_manifest.sh"), str(manifest_path)],
            cwd=ROOT,
            check=True,
            text=True,
            capture_output=True,
        )
        runtime_lines = [line for line in runtime.stdout.splitlines() if line.strip()]
        if len(runtime_lines) >= 2:
            receipt_path = runtime_lines[0]
            runtime_summary_path = runtime_lines[1]
        benchmark_status = run_command(
            proposal["benchmark"]["command"],
            cwd=ROOT,
            stdout_path=benchmark_stdout,
            stderr_path=benchmark_stderr,
        )
        if benchmark_status == 0 and args.stage_on_pass:
            subprocess.run(["git", "add", *proposal["target_files"]], cwd=ROOT, check=True)
            staged_files = proposal["target_files"]

    summary = {
        "version": "self_improve_run_v0",
        "generated_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "goal": args.goal,
        "config_path": str(config_path),
        "proposal_path": str(proposal_path),
        "manifest_path": str(manifest_path),
        "brief_path": str(brief_path),
        "applied": args.apply,
        "target_files": proposal["target_files"],
        "receipt_path": receipt_path,
        "runtime_summary_path": runtime_summary_path,
        "benchmark_command": proposal["benchmark"]["command"],
        "benchmark_status": benchmark_status,
        "benchmark_stdout": str(benchmark_stdout) if args.apply else None,
        "benchmark_stderr": str(benchmark_stderr) if args.apply else None,
        "staged_files": staged_files,
        "proposal": proposal,
    }
    summary_path = run_dir / "summary.json"
    summary_path.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    json.dump(summary, sys.stdout, indent=2)
    sys.stdout.write("\n")
    if args.apply and config.get("proposal_contract", {}).get("fail_on_benchmark_error", False):
        if benchmark_status is None:
            return 1
        if benchmark_status != 0:
            return benchmark_status
    return 0


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