"""CLI handler for `axi agents` — agent service lifecycle management.

Usage:
    axi agents status              Show running/stopped state for all agents
    axi agents start [name]        Start one or all always-on agents
    axi agents stop [name]         Stop one or all always-on agents
    axi agents logs [name]         Tail service log output
    axi agents register            Register all daemon agents as system services
    axi agents unregister [name]   Remove service registrations
"""

from __future__ import annotations

import argparse
import sys
from dataclasses import dataclass

from axiom.extensions.contracts import Extension


@dataclass
class RegistrationResult:
    """Outcome of registering a single daemon agent as an OS service."""

    agent_name: str
    registered: bool  # install() returned True
    started: bool  # start() returned True
    provider: str = ""
    error: str = ""

    @property
    def ok(self) -> bool:
        return self.registered and self.started


def register_all_daemon_agents() -> list[RegistrationResult]:
    """Register every always-on agent as an OS service (idempotent).

    Used by the setup wizard, the installer finalize phase, and the CLI
    self-heal hook — one code path so drift stays impossible.

    Never raises; failures are reported in the returned results so callers
    can decide whether to soft-fail or hard-fail.
    """
    results: list[RegistrationResult] = []
    try:
        agents = _discover_agent_extensions()
    except Exception as exc:  # pragma: no cover — discovery should not crash
        return [
            RegistrationResult(
                agent_name="<discovery>", registered=False, started=False, error=str(exc)
            )
        ]

    daemons = [e for e in agents if e.agent and e.agent.is_always_on]
    for ext in daemons:
        # Refuse to register an agent that has no heartbeat_command. This
        # is the guard that fixes the v0.9.0–v0.10.2 crash-loop where we
        # happily emitted `axi <noun> heartbeat` as ExecStart even though
        # no such subcommand existed on any agent.
        if not ext.agent.is_registrable:
            results.append(
                RegistrationResult(
                    agent_name=ext.name,
                    registered=False,
                    started=False,
                    error=(
                        "no heartbeat_command declared in [agent]; "
                        "add one or change startup from 'daemon' to 'lazy'"
                    ),
                )
            )
            continue
        try:
            mgr = _make_service_manager(ext)
            installed = mgr.install()
            started = mgr.start() if installed else False
            results.append(
                RegistrationResult(
                    agent_name=ext.name,
                    registered=bool(installed),
                    started=bool(started),
                    provider=mgr.provider_name,
                    error="" if installed and started else "install or start returned False",
                )
            )
        except Exception as exc:
            results.append(
                RegistrationResult(
                    agent_name=ext.name,
                    registered=False,
                    started=False,
                    error=str(exc),
                )
            )
    return results


def missing_daemon_agents() -> list[str]:
    """Return names of always-on agents whose OS service is not installed.

    Used by the CLI self-heal hook as a cheap drift detector. Never raises.
    """
    try:
        agents = _discover_agent_extensions()
    except Exception:
        return []
    missing: list[str] = []
    for ext in agents:
        if not (ext.agent and ext.agent.is_always_on):
            continue
        try:
            mgr = _make_service_manager(ext)
            info = mgr.status()
            if info.status == "not_installed":
                missing.append(ext.name)
        except Exception:
            # Treat errors as "unknown" — don't add to missing list to avoid noise.
            pass
    return missing


# ---------------------------------------------------------------------------
# Discovery helpers
# ---------------------------------------------------------------------------


def _discover_agent_extensions() -> list[Extension]:
    """Find all extensions with kind='agent' and an [agent] section."""
    from axiom.extensions.discovery import discover_extensions

    return [ext for ext in discover_extensions() if ext.kind == "agent" and ext.agent is not None]


def _make_service_manager(ext: Extension):
    """Create a ServiceManager for an agent extension."""
    from axiom.infra.branding import get_branding
    from axiom.infra.services import ServiceManager

    cli = get_branding().cli_name
    return ServiceManager(
        name=f"{ext.name}-agent",
        binary=cli,
        args=_agent_service_args(ext),
        interval_secs=ext.agent.heartbeat_interval if ext.agent else 0,
    )


def _agent_service_args(ext: Extension) -> list[str]:
    """Build the argv the systemd timer fires on each heartbeat tick.

    Reads ``[agent] heartbeat_command`` from the extension manifest — a
    space-separated string like ``"mo health --json"`` that becomes a
    oneshot command. Callers MUST gate registration on
    ``ext.agent.is_registrable`` so agents without a declared
    heartbeat_command don't end up with a crash-looping invocation (the
    v0.9.0–v0.10.2 regression that motivated this design).
    """
    if ext.agent and ext.agent.heartbeat_command:
        return ext.agent.heartbeat_command.split()
    # Defensive: if caller failed to gate on is_registrable, emit a
    # command that will fail fast and obviously rather than loop.
    return ["--invalid-no-heartbeat-command-configured"]


def _get_service_status(ext: Extension) -> str:
    """Get the service status string for an agent."""

    mgr = _make_service_manager(ext)
    info = mgr.status()
    return info.status


def _find_agent(agents: list[Extension], name: str) -> Extension | None:
    """Find an agent by name (matches extension name or CLI noun)."""
    for ext in agents:
        if ext.name == name:
            return ext
        for cmd in ext.cli_commands:
            if cmd.noun == name:
                return ext
    return None


# ---------------------------------------------------------------------------
# Command handlers
# ---------------------------------------------------------------------------


def _cmd_status(args) -> int:
    agents = _discover_agent_extensions()
    if not agents:
        print("No agent extensions with lifecycle configuration found.")
        return 0

    print("Agent Services")
    print("=" * 60)
    print(f"  {'Agent':<20} {'Startup':<10} {'Interval':<10} {'Status':<15}")
    print(f"  {'-' * 20} {'-' * 10} {'-' * 10} {'-' * 15}")

    for ext in agents:
        status = _get_service_status(ext)
        interval = f"{ext.agent.heartbeat_interval}s"
        startup = ext.agent.startup

        # Color status
        color_map = {
            "running": "\033[32m",  # green
            "starting": "\033[33m",  # yellow — transient, expect running soon
            "stopped": "\033[33m",  # yellow — installed but not active
            "failed": "\033[31m",  # red — operator attention required
            "not_installed": "\033[90m",  # grey
            "unknown": "\033[90m",  # grey
        }
        color = color_map.get(status, "")
        label = "not installed" if status == "not_installed" else status
        status_str = f"{color}{label}\033[0m" if color else label

        print(f"  {ext.name:<20} {startup:<10} {interval:<10} {status_str:<15}")

    # Show watchers summary
    print()
    for ext in agents:
        if ext.agent.watchers:
            enabled = [w for w in ext.agent.watchers if w.enabled]
            if enabled:
                names = ", ".join(w.name for w in enabled)
                print(f"  {ext.name} watchers: {names}")

    print()
    return 0


def _cmd_start(args) -> int:
    agents = _discover_agent_extensions()
    name = getattr(args, "name", None)

    if name:
        ext = _find_agent(agents, name)
        if ext is None:
            print(f"Unknown agent: {name}")
            available = [e.name for e in agents]
            if available:
                print(f"  Available: {', '.join(available)}")
            return 1
        return _start_one(ext)

    # Start all daemon agents
    daemon_agents = [e for e in agents if e.agent.is_always_on]
    if not daemon_agents:
        print("No daemon agents configured.")
        return 0

    rc = 0
    for ext in daemon_agents:
        if _start_one(ext) != 0:
            rc = 1
    return rc


def _start_one(ext: Extension) -> int:
    mgr = _make_service_manager(ext)
    print(f"Starting {ext.name}...")
    if mgr.install() and mgr.start():
        info = mgr.status()
        print(f"  {ext.name}: {info.status} ({info.provider})")
        return 0
    print(f"  Failed to start {ext.name}")
    return 1


def _cmd_stop(args) -> int:
    agents = _discover_agent_extensions()
    name = getattr(args, "name", None)

    if name:
        ext = _find_agent(agents, name)
        if ext is None:
            print(f"Unknown agent: {name}")
            return 1
        return _stop_one(ext)

    # Stop all
    rc = 0
    for ext in agents:
        if _stop_one(ext) != 0:
            rc = 1
    return rc


def _stop_one(ext: Extension) -> int:
    mgr = _make_service_manager(ext)
    if mgr.stop():
        print(f"  {ext.name}: stopped")
        return 0
    print(f"  {ext.name}: failed to stop (may not be running)")
    return 1


def _cmd_register(args) -> int:
    """Register all daemon agents as system services."""
    results = register_all_daemon_agents()
    if not results:
        print("No daemon agents to register.")
        return 0

    print("Registering agent services...")
    rc = 0
    for r in results:
        if r.ok:
            print(f"  Registered: {r.agent_name} ({r.provider})")
        else:
            detail = r.error or "install or start failed"
            print(f"  Failed: {r.agent_name} — {detail}")
            rc = 1
    return rc


def _cmd_unregister(args) -> int:
    agents = _discover_agent_extensions()
    name = getattr(args, "name", None)

    if name:
        ext = _find_agent(agents, name)
        if ext is None:
            print(f"Unknown agent: {name}")
            return 1
        mgr = _make_service_manager(ext)
        if mgr.uninstall():
            print(f"  Unregistered: {ext.name}")
            return 0
        print(f"  Failed to unregister {ext.name}")
        return 1

    # Unregister all
    rc = 0
    for ext in agents:
        mgr = _make_service_manager(ext)
        if mgr.uninstall():
            print(f"  Unregistered: {ext.name}")
        else:
            print(f"  Failed: {ext.name}")
            rc = 1
    return rc


def _cmd_logs(args) -> int:
    from axiom.infra.paths import get_user_state_dir

    agents = _discover_agent_extensions()
    name = getattr(args, "name", None)

    if name:
        ext = _find_agent(agents, name)
        if ext is None:
            print(f"Unknown agent: {name}")
            return 1
        agents = [ext]

    services_dir = get_user_state_dir() / "services"
    for ext in agents:
        stdout_log = services_dir / f"{ext.name}-agent.stdout.log"
        stderr_log = services_dir / f"{ext.name}-agent.stderr.log"

        print(f"--- {ext.name} ---")
        for log_path in (stdout_log, stderr_log):
            if log_path.exists():
                lines = log_path.read_text(encoding="utf-8").splitlines()
                for line in lines[-20:]:
                    print(f"  {line}")
            else:
                print(f"  No log: {log_path.name}")
        print()

    return 0


# ---------------------------------------------------------------------------
# CLI parser
# ---------------------------------------------------------------------------


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="axi agents",
        description="Manage always-on agent services",
    )
    sub = parser.add_subparsers(dest="action")

    sub.add_parser("status", help="Show agent service status")

    start_p = sub.add_parser("start", help="Start one or all agents")
    start_p.add_argument("name", nargs="?", help="Agent name (omit for all)")

    stop_p = sub.add_parser("stop", help="Stop one or all agents")
    stop_p.add_argument("name", nargs="?", help="Agent name (omit for all)")

    sub.add_parser("register", help="Register all daemon agents as system services")

    unreg_p = sub.add_parser("unregister", help="Remove service registrations")
    unreg_p.add_argument("name", nargs="?", help="Agent name (omit for all)")

    logs_p = sub.add_parser("logs", help="Tail agent service logs")
    logs_p.add_argument("name", nargs="?", help="Agent name (omit for all)")

    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    if not args.action:
        args.action = "status"

    handlers = {
        "status": _cmd_status,
        "start": _cmd_start,
        "stop": _cmd_stop,
        "register": _cmd_register,
        "unregister": _cmd_unregister,
        "logs": _cmd_logs,
    }

    handler = handlers.get(args.action)
    if handler is None:
        parser.print_help()
        return 1

    return handler(args)


if __name__ == "__main__":
    sys.exit(main())
