#!/usr/bin/env python3
"""Mesh Optimizer agent CLI.

Usage:
    mesh-optimizer start [--config FILE] [--port PORT]
    mesh-optimizer stop
    mesh-optimizer status
    mesh-optimizer about
    mesh-optimizer settings
    mesh-optimizer settings set <key> <value>
    mesh-optimizer add
    mesh-optimizer nodes
    mesh-optimizer --config FILE          # (no subcommand = start)
"""
from __future__ import annotations

import argparse
import json
import logging
import os
import platform
import signal
import subprocess
import sys
from pathlib import Path

import uvicorn


def _default_config() -> str:
    """Find config in standard locations."""
    candidates = [
        Path.home() / ".mesh-optimizer" / "mesh_config.yaml",
        Path("mesh/mesh_config.yaml"),
        Path("mesh_config.yaml"),
    ]
    for p in candidates:
        if p.exists():
            return str(p)
    return str(candidates[0])


def _pid_file() -> Path:
    return Path.home() / ".mesh-optimizer" / "mesh-optimizer.pid"


def _is_running() -> tuple[bool, int | None]:
    """Check if agent is running. Returns (running, pid)."""
    pf = _pid_file()
    if not pf.exists():
        return False, None
    try:
        pid = int(pf.read_text().strip())
        os.kill(pid, 0)  # signal 0 = check existence
        return True, pid
    except (ValueError, ProcessLookupError, PermissionError):
        pf.unlink(missing_ok=True)
        return False, None


def _agent_port(config_path: str) -> int:
    """Read agent port from config."""
    try:
        import yaml
        with open(config_path) as f:
            cfg = yaml.safe_load(f)
        return cfg.get("node", {}).get("agent_port", 8400)
    except Exception:
        return 8400


def cmd_start(args):
    """Start the mesh agent."""
    running, pid = _is_running()
    if running:
        print(f"Agent already running (PID {pid})")
        sys.exit(0)

    project_root = Path(__file__).resolve().parents[2]
    sys.path.insert(0, str(project_root))

    from mesh.config import MeshConfig
    from mesh.agent.node_agent import NodeAgent
    from mesh.api.agent_api import app, wire_agent

    config = MeshConfig.from_yaml(args.config)
    if args.port:
        config.node.agent_port = args.port
    if args.log_level:
        config.log_level = args.log_level

    logging.basicConfig(
        level=getattr(logging, config.log_level.upper()),
        format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
    )
    logger = logging.getLogger("mesh.agent")
    logger.info("Starting agent on port %d", config.node.agent_port)

    # Write PID file
    pf = _pid_file()
    pf.parent.mkdir(parents=True, exist_ok=True)
    pf.write_text(str(os.getpid()))

    def _cleanup(*_a):
        pf.unlink(missing_ok=True)
        sys.exit(0)

    signal.signal(signal.SIGTERM, _cleanup)
    signal.signal(signal.SIGINT, _cleanup)

    try:
        agent = NodeAgent(config)
        wire_agent(node_agent=agent)
        uvicorn.run(app, host=config.node.host, port=config.node.agent_port,
                    log_level=config.log_level.lower())
    finally:
        pf.unlink(missing_ok=True)


def cmd_stop(args):
    """Stop the mesh agent."""
    running, pid = _is_running()
    if not running:
        print("Agent is not running.")
        sys.exit(0)

    print(f"Stopping agent (PID {pid})...")
    try:
        os.kill(pid, signal.SIGTERM)
        print("Agent stopped.")
    except ProcessLookupError:
        print("Agent already stopped.")
    _pid_file().unlink(missing_ok=True)


def cmd_status(args):
    """Show agent status."""
    running, pid = _is_running()
    config_path = args.config

    print(f"Mesh Optimizer Agent")
    print(f"{'─' * 40}")

    if running:
        print(f"  Status:   \033[32mrunning\033[0m (PID {pid})")
    else:
        print(f"  Status:   \033[31mstopped\033[0m")

    cfg_path = Path(config_path)
    print(f"  Config:   {cfg_path}")
    print(f"  PID file: {_pid_file()}")

    if cfg_path.exists():
        try:
            import yaml
            with open(cfg_path) as f:
                cfg = yaml.safe_load(f)
            port = cfg.get("node", {}).get("agent_port", 8400)
            ctrl = cfg.get("controller", {}).get("url", "not set")
            nat = cfg.get("controller", {}).get("nat_mode", False)
            cluster = cfg.get("cluster_name", "default")
            print(f"  Port:     {port}")
            print(f"  Cluster:  {cluster}")
            print(f"  Controller: {ctrl}")
            print(f"  NAT mode: {nat}")
        except Exception:
            pass

    # If running, try to hit the local health endpoint
    if running:
        try:
            import urllib.request
            port = _agent_port(config_path)
            resp = urllib.request.urlopen(f"http://127.0.0.1:{port}/health", timeout=2)
            data = json.loads(resp.read())
            print(f"  Health:   {data.get('status', 'unknown')}")
        except Exception:
            print(f"  Health:   unreachable")

    log_dir = Path.home() / ".mesh-optimizer" / "logs"
    if log_dir.exists():
        logs = sorted(log_dir.glob("*.log"), key=lambda p: p.stat().st_mtime, reverse=True)
        if logs:
            print(f"  Logs:     {logs[0]}")


def _load_config(config_path: str) -> dict:
    """Load YAML config file, returning empty dict on failure."""
    try:
        import yaml
        with open(config_path) as f:
            return yaml.safe_load(f) or {}
    except Exception:
        return {}


def _mask_license(key: str) -> str:
    """Mask middle segments of a license key: MESH-XXXX-****-****-YYYY."""
    parts = key.split("-")
    if len(parts) >= 5:
        masked = parts[:2] + ["****"] * (len(parts) - 3) + [parts[-1]]
        return "-".join(masked)
    return key


def cmd_about(args):
    """Show product information."""
    try:
        from mesh._version import __version__
        version = __version__
    except Exception:
        version = "unknown"

    config_path = args.config
    py_version = platform.python_version()
    plat = f"{platform.system()} {platform.machine()}"

    print(f"Mesh Optimizer v{version}")
    print(f"\u2500" * 37)
    print(f"  Company:    Slento Systems")
    print(f"  Website:    https://slentosystems.com")
    print(f"  Docs:       https://docs.slentosystems.com")
    print(f"  Portal:     https://portal.slentosystems.com")
    print(f"  Python:     {py_version}")
    print(f"  Platform:   {plat}")
    print(f"  Config:     {config_path}")


def cmd_settings(args):
    """Show current settings from config file."""
    config_path = args.config
    cfg = _load_config(config_path)

    controller = cfg.get("controller_url", "not set")
    nat_mode = cfg.get("node", {}).get("nat_mode", False)
    agent_port = cfg.get("node", {}).get("agent_port", 8400)
    license_key = cfg.get("licensing", {}).get("license_key", "not set")
    cluster = cfg.get("cluster_name", "default")
    log_level = cfg.get("log_level", "info")

    masked = _mask_license(license_key) if license_key != "not set" else "not set"

    print("Mesh Optimizer Settings")
    print("\u2500" * 37)
    print(f"  Controller: {controller}")
    print(f"  NAT mode:   {str(nat_mode).lower()}")
    print(f"  Agent port: {agent_port}")
    print(f"  License:    {masked}")
    print(f"  Cluster:    {cluster}")
    print(f"  Log level:  {log_level.lower()}")


def cmd_settings_set(args):
    """Update a config value."""
    import yaml

    key = args.settings_key
    value = args.settings_value
    config_path = args.config

    cfg_file = Path(config_path)
    if cfg_file.exists():
        with open(cfg_file) as f:
            cfg = yaml.safe_load(f) or {}
    else:
        cfg = {}

    supported_keys = {
        "controller-url": "Sets controller_url",
        "license-key": "Sets licensing.license_key",
        "nat-mode": "Sets node.nat_mode (true/false)",
        "agent-port": "Sets node.agent_port",
        "log-level": "Sets log_level",
    }

    if key not in supported_keys:
        print(f"Unknown setting: {key}")
        print(f"Supported keys: {', '.join(supported_keys.keys())}")
        sys.exit(1)

    if key == "controller-url":
        cfg["controller_url"] = value
    elif key == "license-key":
        cfg.setdefault("licensing", {})["license_key"] = value
    elif key == "nat-mode":
        if value.lower() not in ("true", "false"):
            print("Error: nat-mode must be 'true' or 'false'")
            sys.exit(1)
        cfg.setdefault("node", {})["nat_mode"] = value.lower() == "true"
    elif key == "agent-port":
        try:
            cfg.setdefault("node", {})["agent_port"] = int(value)
        except ValueError:
            print("Error: agent-port must be an integer")
            sys.exit(1)
    elif key == "log-level":
        cfg["log_level"] = value

    cfg_file.parent.mkdir(parents=True, exist_ok=True)
    with open(cfg_file, "w") as f:
        yaml.dump(cfg, f, default_flow_style=False, sort_keys=False)

    print(f"Updated '{key}' to '{value}' in {cfg_file}")


def cmd_add(args):
    """Print install commands for adding machines to the mesh."""
    config_path = args.config
    cfg = _load_config(config_path)

    license_key = cfg.get("licensing", {}).get("license_key", "")
    controller = cfg.get("controller_url", "https://mesh.slento.ai")
    nat_mode = cfg.get("node", {}).get("nat_mode", False)

    if not license_key:
        license_key = "YOUR-LICENSE-KEY"

    nat_flag_sh = " \\\n      --nat" if nat_mode else ""
    nat_flag_ps = " -NatMode" if nat_mode else ""

    print("Add machines to your mesh")
    print("\u2500" * 37)
    print()
    print("Linux:")
    print(f"  curl -fsSL https://mesh.slentosystems.com/install.sh | bash -s -- \\")
    print(f"    --license {license_key} \\")
    print(f"    --controller {controller}{nat_flag_sh}")
    print()
    print("macOS:")
    print(f"  curl -fsSL https://mesh.slentosystems.com/installers/macos/install.sh | bash -s -- \\")
    print(f"    --license {license_key} \\")
    print(f"    --controller {controller}{nat_flag_sh}")
    print()
    print("Windows (PowerShell):")
    print(f"  irm https://mesh.slentosystems.com/installers/windows/install.ps1 -OutFile install.ps1")
    print(f"  .\\install.ps1 -LicenseKey \"{license_key}\" -ControllerUrl \"{controller}\"{nat_flag_ps}")


def cmd_upgrade(args):
    """Upgrade mesh-optimizer to the latest version from PyPI."""
    import subprocess as sp

    try:
        from mesh._version import __version__
        current = __version__
    except Exception:
        current = "unknown"

    print(f"Mesh Optimizer Upgrade")
    print("\u2500" * 37)
    print(f"  Current version: {current}")
    print(f"  Checking PyPI for updates...")

    # Find the pip in the same venv as this script
    venv_pip = Path(sys.executable).parent / "pip"
    pip_cmd = str(venv_pip) if venv_pip.exists() else f"{sys.executable} -m pip"

    try:
        result = sp.run(
            [sys.executable, "-m", "pip", "install", "--upgrade",
             "--no-cache-dir", "slento-mesh-optimizer"],
            capture_output=True, text=True, timeout=120,
        )
        if result.returncode != 0:
            print(f"  \033[31mUpgrade failed:\033[0m {result.stderr.strip()[-200:]}")
            sys.exit(1)

        # Check new version
        result2 = sp.run(
            [sys.executable, "-m", "pip", "show", "slento-mesh-optimizer"],
            capture_output=True, text=True,
        )
        new_version = "unknown"
        for line in result2.stdout.splitlines():
            if line.startswith("Version:"):
                new_version = line.split(":", 1)[1].strip()
                break

        if new_version == current:
            print(f"  Already up to date (v{current})")
        else:
            print(f"  \033[32mUpgraded: v{current} -> v{new_version}\033[0m")
            print()
            running, pid = _is_running()
            if running:
                print(f"  Agent is running (PID {pid}). Restart to apply:")
                print(f"    mesh-optimizer stop && mesh-optimizer start")
            else:
                print(f"  Start the agent with: mesh-optimizer start")

    except sp.TimeoutExpired:
        print(f"  \033[31mUpgrade timed out\033[0m")
        sys.exit(1)
    except Exception as e:
        print(f"  \033[31mUpgrade error:\033[0m {e}")
        sys.exit(1)


def cmd_nodes(args):
    """List nodes from the controller."""
    import urllib.request
    import urllib.error

    config_path = args.config
    cfg = _load_config(config_path)

    controller = cfg.get("controller_url", "https://mesh.slento.ai")
    url = f"{controller.rstrip('/')}/nodes"

    print(f"Mesh Nodes (via {controller})")
    print("\u2500" * 37)

    try:
        req = urllib.request.Request(url, headers={"Accept": "application/json"})
        resp = urllib.request.urlopen(req, timeout=10)
        data = json.loads(resp.read())
    except urllib.error.URLError as e:
        print(f"  Error: could not reach controller at {url}")
        print(f"         {e}")
        sys.exit(1)
    except Exception as e:
        print(f"  Error: {e}")
        sys.exit(1)

    # Handle both list-of-nodes and dict with "nodes" key
    nodes = data if isinstance(data, list) else data.get("nodes", [])

    if not nodes:
        print("  No nodes found.")
        return

    # Calculate column widths for alignment
    for node in nodes:
        name = node.get("name", node.get("hostname", "unknown"))
        os_name = node.get("os", node.get("platform", "unknown"))
        status = node.get("status", "unknown")
        gpu_count = node.get("gpu_count", node.get("gpus", 0))
        memory_gb = node.get("memory_gb", node.get("ram_gb", 0))

        gpu_label = f"{gpu_count} GPU" + ("s" if gpu_count != 1 else "")
        mem_label = f"{memory_gb} GB"

        # Color status
        if status == "online":
            status_str = f"\033[32m{status}\033[0m"
        elif status == "offline":
            status_str = f"\033[31m{status}\033[0m"
        else:
            status_str = status

        print(f"  {name:<20s} {os_name:<9s} {status_str:<18s} {gpu_label:<8s} {mem_label}")


def main():
    parser = argparse.ArgumentParser(
        prog="mesh-optimizer",
        description="Mesh Optimizer — distributed hardware optimization agent",
    )
    parser.add_argument("--config", default=_default_config(), help="Config file path")
    parser.add_argument("--port", type=int, default=None, help="Override agent port")
    parser.add_argument("--log-level", default=None, help="Log level")

    sub = parser.add_subparsers(dest="command")

    sub.add_parser("start", help="Start the agent")
    sub.add_parser("stop", help="Stop the agent")
    sub.add_parser("status", help="Show agent status")
    sub.add_parser("about", help="Show product info")
    sub.add_parser("add", help="Print install commands for adding machines")
    sub.add_parser("nodes", help="List nodes from the controller")
    sub.add_parser("upgrade", help="Upgrade mesh-optimizer to the latest version")
    sub.add_parser("update", help="Upgrade mesh-optimizer to the latest version")

    settings_parser = sub.add_parser("settings", help="Show or update settings")
    settings_sub = settings_parser.add_subparsers(dest="settings_command")
    set_parser = settings_sub.add_parser("set", help="Update a config value")
    set_parser.add_argument("settings_key", metavar="key",
                            help="Setting key (controller-url, license-key, nat-mode, agent-port, log-level)")
    set_parser.add_argument("settings_value", metavar="value", help="New value")

    args = parser.parse_args()

    # Default to start if no subcommand given
    if args.command is None or args.command == "start":
        cmd_start(args)
    elif args.command == "stop":
        cmd_stop(args)
    elif args.command == "status":
        cmd_status(args)
    elif args.command == "about":
        cmd_about(args)
    elif args.command == "settings":
        if hasattr(args, "settings_command") and args.settings_command == "set":
            cmd_settings_set(args)
        else:
            cmd_settings(args)
    elif args.command == "add":
        cmd_add(args)
    elif args.command == "nodes":
        cmd_nodes(args)
    elif args.command in ("upgrade", "update"):
        cmd_upgrade(args)


if __name__ == "__main__":
    main()
