#!/usr/bin/env python3
"""Orze CLI — GPU experiment orchestrator.

Usage:
    python -m orze.cli                          # all GPUs, continuous
    python -m orze.cli -c orze.yaml --gpus 0,1  # with project config
    python -m orze.cli --once                   # one cycle then exit
    python -m orze.cli --report-only            # regenerate report
    python -m orze.cli --init                   # initialize new project
    python -m orze.cli --admin                  # launch admin panel
"""

import argparse
import logging
import os
import sys
from pathlib import Path

from orze import __version__
from orze.core.config import load_project_config
from orze.hardware.gpu import detect_all_gpus

logger = logging.getLogger("orze")

# ---------------------------------------------------------------------------
# Template for --init scaffolding
# ---------------------------------------------------------------------------

BASELINE_TRAIN_PY = '''\
#!/usr/bin/env python3
"""Baseline training script generated by orze --init.

Replace this with your real training script. Orze will call it as:
    python train.py --idea_config <path_to_yaml>
"""
import argparse
import json
import sys
import time
from pathlib import Path

import yaml


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--idea_config", required=True, help="Path to idea config YAML")
    args = parser.parse_args()

    cfg = yaml.safe_load(Path(args.idea_config).read_text())
    print(f"Training with config: {cfg}")

    # --- Replace with your training logic ---
    time.sleep(2)
    accuracy = 0.5
    loss = 1.0
    # -----------------------------------------

    # Write metrics (orze reads this to track results)
    metrics = {
        "status": "COMPLETED",
        "test_accuracy": accuracy,
        "test_loss": loss,
        "training_time": 2.0,
    }
    idea_dir = Path(args.idea_config).parent
    (idea_dir / "metrics.json").write_text(json.dumps(metrics, indent=2))
    print(f"Done. Metrics: {metrics}")


if __name__ == "__main__":
    main()
'''


# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------

def setup_logging(verbose: bool = False):
    """Configure logging with timestamps."""
    level = logging.DEBUG if verbose else logging.INFO
    fmt = "%(asctime)s [%(levelname)s] %(message)s"
    datefmt = "%Y-%m-%d %H:%M:%S"
    logging.basicConfig(level=level, format=fmt, datefmt=datefmt)


# ---------------------------------------------------------------------------
# Uninstall
# ---------------------------------------------------------------------------

# Files inside each idea dir to keep (research results)
_IDEA_KEEP = {
    "metrics.json",
    "eval_report.json",
    "eval_output.log",
}

# Top-level results_dir files to keep
_RESULTS_KEEP = {"report.md", "status.json"}


def _do_uninstall(cfg: dict):
    """Full uninstall: stop orze, strip runtime files, pip uninstall.

    Preserves research results:
        - results/{idea-*}/metrics.json  (core experiment data)
        - results/{idea-*}/eval_report.json, eval_output.log
        - results/report.md, status.json
        - ideas.md, idea_lake.db
    """
    import shutil
    import subprocess

    results_dir = Path(cfg["results_dir"])
    print("\n\033[1mOrze — Uninstall\033[0m")
    print("-----------------")
    print(f"Results dir : {results_dir.resolve()}")
    print(f"Config      : {cfg.get('_config_path', 'orze.yaml')}")
    print()

    # --- 1. Stop running instances -----------------------------------
    print("[1/4] Stopping running orze instances...")
    stop_path = results_dir / ".orze_stop_all"
    if results_dir.exists():
        stop_path.write_text("uninstall", encoding="utf-8")
        # Also send SIGTERM to any PID files we find
        for pid_file in results_dir.glob(".orze.pid*"):
            try:
                pid = int(pid_file.read_text(encoding="utf-8").strip())
                os.kill(pid, 15)  # SIGTERM
                print(f"  Sent SIGTERM to PID {pid}")
            except (ValueError, ProcessLookupError, PermissionError, OSError):
                pass
    import time
    time.sleep(2)  # give processes a moment to exit

    # --- 2. Clean results directory ----------------------------------
    print("[2/4] Cleaning runtime files from results/...")
    if results_dir.exists():
        # Fast path: rename the whole dir, create a fresh one, move only
        # keeper files back (O(1) renames), then bulk-delete the old tree.
        stale = results_dir.with_name(results_dir.name + "._orze_uninstall_tmp")
        if stale.exists():
            shutil.rmtree(stale, ignore_errors=True)
        results_dir.rename(stale)
        results_dir.mkdir()

        # Restore top-level keeper files (report.md, status.json)
        for name in _RESULTS_KEEP:
            src = stale / name
            if src.exists():
                src.rename(results_dir / name)

        # Restore keeper files and subdirs from idea directories
        for p in stale.iterdir():
            if not p.is_dir() or p.name.startswith((".","_")):
                continue
            keep_files = [c for c in p.iterdir()
                          if c.is_file() and c.name in _IDEA_KEEP]
            keep_dirs = [c for c in p.iterdir() if c.is_dir()]
            if keep_files or keep_dirs:
                dest = results_dir / p.name
                dest.mkdir()
                for f in keep_files:
                    f.rename(dest / f.name)
                for d in keep_dirs:
                    d.rename(dest / d.name)

        # Bulk-delete the old tree (fast, handled in C)
        shutil.rmtree(stale, ignore_errors=True)

    print("  Cleaned runtime files")

    # --- 3. Remove orze config file ----------------------------------
    print("[3/4] Removing orze config...")
    config_path = Path(cfg.get("_config_path", "orze.yaml"))
    if config_path.exists():
        config_path.unlink()
        print(f"  Removed {config_path}")
    else:
        print("  No config file found")

    # --- 4. Pip uninstall --------------------------------------------
    print("[4/4] Uninstalling orze package...")
    try:
        subprocess.run(
            [sys.executable, "-m", "pip", "uninstall", "orze", "-y"],
            check=True,
        )
        print("  orze package removed")
    except subprocess.CalledProcessError:
        print("  WARNING: pip uninstall failed — you may need to run manually:")
        print("    pip uninstall orze")

    # --- Summary -----------------------------------------------------
    print()
    print("\033[1mUninstall complete.\033[0m")
    print("Preserved research results:")
    if results_dir.exists():
        kept = list(results_dir.rglob("metrics.json"))
        print(f"  {len(kept)} experiment result(s) in {results_dir}/")
    ideas_file = Path(cfg.get("ideas_file", "ideas.md"))
    if ideas_file.exists():
        print(f"  {ideas_file} (experiment definitions)")
    lake_path = ideas_file.parent / "idea_lake.db"
    if lake_path.exists():
        print(f"  {lake_path} (idea archive database)")


def _stop_running_instance(results_dir: Path) -> bool:
    """Stop a running orze instance via PID file. Returns True if one was stopped."""
    import time as _time
    pid_file = results_dir / ".orze.pid"
    if not pid_file.exists():
        return False
    try:
        old_pid = int(pid_file.read_text(encoding="utf-8").strip())
    except (ValueError, OSError):
        return False
    print(f"Stopping orze (PID {old_pid})...")
    try:
        os.kill(old_pid, 15)  # SIGTERM
    except ProcessLookupError:
        print("  Process already exited.")
        return False
    for _ in range(30):  # 15s
        try:
            os.kill(old_pid, 0)
            _time.sleep(0.5)
        except ProcessLookupError:
            break
    else:
        print("  Still running after 15s, sending SIGKILL...")
        try:
            os.kill(old_pid, 9)
            _time.sleep(1)
        except ProcessLookupError:
            pass
    print("  Stopped.")
    return True


def _do_upgrade(cfg: dict):
    """Upgrade orze to the latest version from PyPI, then restart if running."""
    import subprocess

    print(f"\n\033[1mOrze — Upgrade\033[0m")
    print("---------------")
    print(f"Current version: {__version__}")

    # 1. Stop running instance first (so new code loads on restart)
    results_dir = Path(cfg["results_dir"])
    was_running = _stop_running_instance(results_dir)

    # 2. Upgrade
    print("Upgrading via pip...\n")
    try:
        subprocess.run(
            [sys.executable, "-m", "pip", "install", "--upgrade", "orze"],
            check=True,
        )
    except subprocess.CalledProcessError:
        print("\n\033[31mUpgrade failed.\033[0m Try manually:")
        print("  pip install --upgrade orze")
        return

    # Report the new version
    result = subprocess.run(
        [sys.executable, "-c", "from orze import __version__; print(__version__)"],
        capture_output=True, text=True,
    )
    new_ver = result.stdout.strip() if result.returncode == 0 else "unknown"
    if new_ver == __version__:
        print(f"\nAlready up-to-date (v{__version__}).")
    else:
        print(f"\n\033[32mUpgraded: v{__version__} -> v{new_ver}\033[0m")

    # 3. Restart if it was running
    if was_running:
        print("\nRestarting orze with new version...")
        os.execv(sys.executable, [sys.executable, "-m", "orze.cli",
                                  "-c", str(cfg.get("_config_path", "orze.yaml"))])


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="orze: GPU experiment orchestrator",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python -m orze.cli                             # all GPUs, continuous
  python -m orze.cli -c orze.yaml --gpus 0,1     # with project config
  python -m orze.cli --once                      # one cycle then exit
  python -m orze.cli --report-only               # regenerate report
        """,
    )
    parser.add_argument("-V", "--version", action="version",
                        version=f"orze {__version__}")
    parser.add_argument("-c", "--config-file", type=str, default=None,
                        help="Path to orze.yaml project config")
    parser.add_argument("--gpus", type=str, default=None,
                        help="Comma-separated GPU IDs (default: auto-detect)")
    parser.add_argument("--timeout", type=int, default=None,
                        help="Max training time in seconds")
    parser.add_argument("--poll", type=int, default=None,
                        help="Seconds between iterations")
    parser.add_argument("--once", action="store_true",
                        help="Run one cycle and exit")
    parser.add_argument("--stop", action="store_true",
                        help="Gracefully stop a running orze instance")
    parser.add_argument("--restart", action="store_true",
                        help="Stop the running instance and start a new one")
    parser.add_argument("--disable", action="store_true",
                        help="Stop and persistently disable Orze (survives restarts)")
    parser.add_argument("--enable", action="store_true",
                        help="Remove persistent disable flag to allow Orze to run")
    parser.add_argument("--report-only", action="store_true",
                        help="Only regenerate report")
    parser.add_argument("--role-only", type=str, default=None, metavar="NAME",
                        help="Run a single agent role once and exit")
    parser.add_argument("--research-only", action="store_true",
                        help="Alias for --role-only research")
    parser.add_argument("--ideas-md", type=str, default=None,
                        help="Path to ideas markdown file")
    parser.add_argument("--base-config", type=str, default=None,
                        help="Path to base config YAML")
    parser.add_argument("--results-dir", type=str, default=None,
                        help="Directory for results")
    parser.add_argument("--train-script", type=str, default=None,
                        help="Training script to run per idea")
    parser.add_argument("--init", action="store_true",
                        help="Initialize a new orze project in the current directory")
    parser.add_argument("--admin", action="store_true",
                        help="Launch admin panel instead of farm loop")
    parser.add_argument("--upgrade", action="store_true",
                        help="Upgrade orze to the latest version from PyPI")
    parser.add_argument("--uninstall", action="store_true",
                        help="Full uninstall: stop orze, remove runtime files, "
                             "pip uninstall — keeps only research results")
    parser.add_argument("-v", "--verbose", action="store_true",
                        help="Debug logging")
    args = parser.parse_args()

    setup_logging(args.verbose)

    # Load project config, then apply CLI overrides
    cfg = load_project_config(args.config_file)
    cfg["_config_path"] = args.config_file or "orze.yaml"  # stored for mode: research

    # --admin: launch web panel
    if args.admin:
        from orze.admin.server import run_admin
        run_admin(cfg)
        return

    # --upgrade: upgrade orze from PyPI (stops + restarts if running)
    if args.upgrade:
        _do_upgrade(cfg)
        return

    # --uninstall: full cleanup, keep only research results
    if args.uninstall:
        _do_uninstall(cfg)
        return

    # --init: initialize a new project in the current directory
    if args.init:
        print("\n\033[1mOrze — Initialization\033[0m")
        print("---------------------")

        # 1. Find/create training script (depth ≤ 2)
        _skip = {"venv", ".venv", "env", ".env", "node_modules", "__pycache__",
                 ".git", ".tox", ".nox", "results", "wandb", ".eggs", "dist", "build"}
        train_scripts = []
        for entry in Path(".").iterdir():
            if entry.is_file() and entry.suffix == ".py" and entry.name.startswith("train") and "orze" not in str(entry):
                train_scripts.append(str(entry))
            elif entry.is_dir() and entry.name not in _skip:
                for child in entry.iterdir():
                    if child.is_file() and child.suffix == ".py" and child.name.startswith("train") and "orze" not in str(child):
                        train_scripts.append(str(child))

        train_script = None
        if train_scripts:
            train_script = train_scripts[0]
            print(f"Detected training script: {train_script}")
        else:
            train_script = "train.py"
            if not Path(train_script).exists():
                print(f"No training script found. Creating {train_script}...")
                Path(train_script).write_text(BASELINE_TRAIN_PY.strip() + "\n", encoding="utf-8")
            else:
                print(f"Using existing {train_script}")

        # 3. Create orze.yaml
        if not Path("orze.yaml").exists():
            print("Creating orze.yaml...")
            backend = ("anthropic" if os.environ.get("ANTHROPIC_API_KEY") else
                       "gemini" if os.environ.get("GEMINI_API_KEY") else
                       "openai" if os.environ.get("OPENAI_API_KEY") else "ollama")

            yaml_content = f"""# orze.yaml — Project configuration
train_script: {train_script}
ideas_file: ideas.md
results_dir: results
python: {sys.executable}

# Research agent settings (optional)
# roles:
#   research:
#     mode: research
#     backend: {backend}
#     cooldown: 600
#     timeout: 300
"""
            Path("orze.yaml").write_text(yaml_content, encoding="utf-8")
        else:
            print("orze.yaml already exists — skipping")

        # 4. Create ideas.md
        if not Path("ideas.md").exists():
            print("Creating ideas.md...")
            Path("ideas.md").write_text(
                "# Ideas\n\n"
                "## idea-0001: Baseline\n"
                "- **Priority**: high\n\n"
                "```yaml\n"
                "learning_rate: 0.001\n"
                "epochs: 10\n"
                "```\n",
                encoding="utf-8",
            )
        else:
            print("ideas.md already exists — skipping")

        # 5. Create results directory
        Path("results").mkdir(exist_ok=True)

        # 6. Copy AGENT.md and RULES.md into the project directory
        pkg_dir = Path(__file__).resolve().parent
        _copy_map = {"AGENT.md": "ORZE-AGENT.md", "RULES.md": "ORZE-RULES.md"}
        for src_name, dst_name in _copy_map.items():
            src = pkg_dir / src_name
            dst = Path(dst_name)
            if src.exists() and not dst.exists():
                dst.write_text(src.read_text(encoding="utf-8"), encoding="utf-8")
                print(f"Copied {dst_name}")
            elif dst.exists():
                print(f"{dst_name} already exists — skipping")

        print()
        print("\033[32m✔ Initialization complete.\033[0m")

        # --- API key guidance ---
        _detected = []
        if os.environ.get("ANTHROPIC_API_KEY"):
            _detected.append("Anthropic")
        if os.environ.get("GEMINI_API_KEY"):
            _detected.append("Gemini")
        if os.environ.get("OPENAI_API_KEY"):
            _detected.append("OpenAI")

        print()
        print("\033[1mResearch agent setup:\033[0m")
        if _detected:
            print(f"  \033[32m✔\033[0m Detected API key(s): {', '.join(_detected)}")
            print(f"    Orze will use \033[36m{_detected[0].lower()}\033[0m for idea generation.")
        else:
            print("  \033[33m⚠ No LLM API key detected.\033[0m")
            print("  Orze needs an API key to auto-generate experiment ideas.")
            print("  Set one of these in your shell profile (~/.bashrc or ~/.zshrc):")
            print()
            print("    \033[36mexport ANTHROPIC_API_KEY=\"sk-ant-...\"  \033[0m# Claude (recommended)")
            print("    \033[36mexport GEMINI_API_KEY=\"AI...\"          \033[0m# Gemini")
            print("    \033[36mexport OPENAI_API_KEY=\"sk-...\"         \033[0m# OpenAI")
            print()
            print("  Then uncomment the \033[36mroles:\033[0m section in orze.yaml.")
            print("  Without a key, orze will only run ideas you add manually to ideas.md.")

        print()
        admin_port = os.environ.get("ORZE_ADMIN_PORT", "8787")
        print("\033[1mWhen orze is running:\033[0m")
        print(f"  Admin UI  →  \033[36mhttp://localhost:{admin_port}\033[0m")
        print()
        print("\033[1mTelegram notifications (optional):\033[0m")
        print("  Add to your \033[36morze.yaml\033[0m:")
        print("    notifications:")
        print("      enabled: true")
        print("      on: [completed, failed, new_best]")
        print("      channels:")
        print("        - type: telegram")
        print("          bot_token: \"YOUR_BOT_TOKEN\"")
        print("          chat_id: \"YOUR_CHAT_ID\"")
        print("  See README for setup instructions.")
        print()
        print("\033[1mNext steps:\033[0m")
        print("  If in Claude/Gemini/Codex CLI  →  do \033[36m@ORZE-AGENT.md\033[0m")
        print("  If not                         →  \033[36m$ orze\033[0m")
        return

    # Apply CLI overrides
    if args.timeout is not None:
        cfg["timeout"] = args.timeout
    if args.poll is not None:
        cfg["poll"] = args.poll
    if args.ideas_md:
        cfg["ideas_file"] = args.ideas_md
    if args.base_config:
        cfg["base_config"] = args.base_config
    if args.results_dir:
        cfg["results_dir"] = args.results_dir
    if args.train_script:
        cfg["train_script"] = args.train_script

    # --stop
    if args.stop:
        import datetime
        from orze.core.fs import atomic_write
        stop_path = Path(cfg["results_dir"]) / ".orze_stop_all"
        atomic_write(stop_path, datetime.datetime.now().isoformat())
        print(f"Stop sentinel written to {stop_path}")
        return

    # --restart: stop running instance, then continue to start a new one
    if args.restart:
        _stop_running_instance(Path(cfg["results_dir"]))
        print("Starting new orze instance...")

    # --disable
    if args.disable:
        import datetime
        from orze.core.fs import atomic_write
        disable_path = Path(cfg["results_dir"]) / ".orze_disabled"
        atomic_write(disable_path, f"Disabled at {datetime.datetime.now().isoformat()}")
        print(f"Orze disabled. Remove {disable_path} to re-enable.")
        return

    # --enable
    if args.enable:
        disable_path = Path(cfg["results_dir"]) / ".orze_disabled"
        if disable_path.exists():
            disable_path.unlink()
            print("Orze re-enabled.")
        else:
            print("Orze was not disabled.")
        return

    # --report-only
    if args.report_only:
        from orze.core.ideas import parse_ideas
        from orze.reporting.leaderboard import update_report
        ideas = parse_ideas(cfg["ideas_file"])
        results_dir = Path(cfg["results_dir"])
        lake = None
        lake_path = Path(cfg["ideas_file"]).parent / "idea_lake.db"
        if lake_path.exists():
            from orze.idea_lake import IdeaLake
            lake = IdeaLake(str(lake_path))
        update_report(results_dir, ideas, cfg, lake=lake)
        print("Report updated.")
        return

    # --research-only is an alias for --role-only research
    if args.research_only:
        args.role_only = "research"

    # Detect GPUs
    if args.gpus:
        gpu_ids = [int(g.strip()) for g in args.gpus.split(",")]
    else:
        gpu_ids = detect_all_gpus()

    if not gpu_ids:
        logger.error("No GPUs detected. Use --gpus to specify manually.")
        sys.exit(1)

    # Start admin panel in background thread (unless --role-only or --admin-off)
    if not args.role_only and not getattr(args, 'no_admin', False):
        try:
            import threading
            from orze.admin.server import run_admin as _run_admin_server
            admin_port = int(os.environ.get("ORZE_ADMIN_PORT", "8787"))

            def _admin_thread():
                try:
                    _run_admin_server(cfg, port=admin_port)
                except Exception as e:
                    logger.warning("Admin panel failed to start: %s", e)

            t = threading.Thread(target=_admin_thread, daemon=True)
            t.start()
            logger.info("Admin panel starting on http://0.0.0.0:%d", admin_port)
        except Exception as e:
            logger.warning("Could not start admin panel: %s", e)

    # Launch orchestrator
    from orze.engine.orchestrator import Orze
    orze = Orze(gpu_ids, cfg, once=args.once)

    if args.role_only:
        orze._run_role_once(args.role_only)
    else:
        orze.run()


if __name__ == "__main__":
    main()
