#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""
Unified swarm command -- THE one way to start workers.

Usage:
  pulse swarm              Interactive picker (provider -> model -> count)
  pulse swarm gemini       Skip provider, pick model + count
  pulse swarm mistral:7b   Skip provider+model, pick count
  pulse swarm stop         Kill all workers
  pulse swarm status       Show running workers
"""
from __future__ import annotations

import sys

from pulse.tasks.task_ops import detect_tier, ensure_orchestrator
from pulse.cli.swarm_providers import (
    MODELS, _PROVIDER_META, _ollama_models,
    _detect_providers, _find_provider_for_model,
)
from pulse.cli.swarm_workers import (
    _is_headless_provider, launch_swarm_worker,
    _swarm_stop, _swarm_status, _swarm_watch,
)

BACK = -1  # sentinel for "go back"

def _pick(prompt: str, options: list[tuple[str, str]], allow_back: bool = False) -> int | None:
    """Show numbered menu. Returns index, BACK (-1), or None (quit)."""
    print(f"\n  {prompt}")
    for i, (label, detail) in enumerate(options, 1):
        print(f"  [{i}] {label}" + (f" ({detail})" if detail else ""))
    hints = "b=back " if allow_back else ""
    print(f"  ({hints}q=quit)")
    while True:
        try:
            raw = input("  > ").strip()
        except (EOFError, KeyboardInterrupt):
            return None
        if raw.lower() in ("q", "quit", "exit"):
            return None
        if allow_back and raw.lower() in ("b", "back"):
            return BACK
        try:
            idx = int(raw)
            if 1 <= idx <= len(options):
                return idx - 1
        except ValueError:
            pass
        print(f"  Enter 1-{len(options)}"
              + (", 'b' to go back" if allow_back else "") + ", or 'q' to quit.")


def _ask_count(default: int = 1) -> int | None:
    try:
        raw = input(f"\n  How many agents? [{default}] (b=back): ").strip()
    except (EOFError, KeyboardInterrupt):
        return None
    if raw.lower() in ("b", "back"):
        return None
    if not raw:
        return default
    try:
        return max(1, min(int(raw), 10))
    except ValueError:
        return default


def _show_crew(crew: list[tuple[str, str, str, int]]):
    if not crew:
        print("\n  Crew: (empty)")
        return
    print(f"\n  Crew ({sum(c for _, _, _, c in crew)} agents):")
    for provider_key, model, tier, count in crew:
        mode = "CLI" if not _is_headless_provider(provider_key) else "API"
        print(f"    {count}x {model} [{tier}] ({mode})")


def _crew_builder(providers: list[tuple[str, str]], skills_used: bool = False):
    crew: list[tuple[str, str, str, int]] = []
    while True:
        _show_crew(crew)
        if crew:
            print(f"\n  [+] Add more agents   [go] Launch crew   [c] Clear")
            try:
                raw = input("  > ").strip().lower()
            except (EOFError, KeyboardInterrupt):
                return
            if raw in ("go", "g", ""):
                # D3: Crew risk analysis before launch
                try:
                    from pulse.daemons.synthesis.swarm.crew_composer import (
                        score_crew, get_pending_tasks, format_crew_analysis,
                    )
                    tasks = get_pending_tasks()
                    crew_dicts = [{"model": m, "tier": t, "count": c, "provider": p} for p, m, t, c in crew]
                    analysis = score_crew(crew_dicts, tasks)
                    if analysis.get("warnings") or analysis.get("suggestions"):
                        print(format_crew_analysis(analysis))
                        print(f"\n  [go] Launch anyway   [+] Add more   [c] Clear")
                        try:
                            raw2 = input("  > ").strip().lower()
                        except (EOFError, KeyboardInterrupt):
                            return
                        if raw2 in ("c", "clear"):
                            crew.clear()
                            continue
                        if raw2 in ("+", "add", "a"):
                            continue
                        # Any other input (go, g, enter) → proceed to launch
                except Exception:
                    pass  # Crew analysis is advisory — never block launch
                break
            if raw in ("c", "clear"):
                crew.clear()
                continue
            if raw not in ("+", "add", "a"):
                print("  Enter '+' to add, 'go' to launch, 'c' to clear.")
                continue
        # Step 1: Pick provider
        idx = _pick("Select Provider:",
                     [(label, "") for _, label in providers], allow_back=bool(crew))
        if idx is None:
            return
        if idx == BACK:
            if crew:
                crew.pop()
            continue
        provider_key = providers[idx][0]
        # Step 2: Pick model
        models = MODELS.get(provider_key, [])
        if provider_key == "ollama":
            models = _ollama_models()
            if not models:
                print("  [PULSE] No Ollama models. Pull one first: ollama pull mistral:7b")
                continue
        if not models:
            print(f"  [PULSE] No models for {provider_key}.")
            continue

        midx = _pick("Select Model:",
                      [(name, tier) for name, tier in models], allow_back=True)
        if midx is None:
            return
        if midx == BACK:
            continue
        model, tier = models[midx]
        # Step 3: Pick count
        count = _ask_count()
        if count is None:
            continue
        crew.append((provider_key, model, tier, count))
    # Launch the whole crew
    if not crew:
        return
    total = sum(c for _, _, _, c in crew)
    print(f"\n  Launching {total} agent(s)...")
    for provider_key, model, tier, count in crew:
        launch_swarm_worker(provider_key, model, tier, count, skills_used=skills_used)


def _pick_model_and_launch(provider_key: str, skills_used: bool = False):
    models = MODELS.get(provider_key, [])
    if provider_key == "ollama":
        models = _ollama_models()
        if not models:
            print("[PULSE] No Ollama models found. Pull one first: ollama pull mistral:7b")
            return
    if not models:
        print(f"[PULSE] No models configured for {provider_key}.")
        return
    idx = _pick("Select Model:", [(name, tier) for name, tier in models], allow_back=True)
    if idx is None or idx == BACK:
        return
    model, tier = models[idx]
    count = _ask_count()
    if count is None or count == 0:
        return
    print(f"\n  Starting {count}x {model} workers ({tier} tier)...")
    launch_swarm_worker(provider_key, model, tier, count, skills_used=skills_used)

def _swarm_by_tier(tier: str, count: int | None, skills_used: bool = False):
    providers = _detect_providers()
    matches = []
    for pk, _ in providers:
        models = MODELS.get(pk, [])
        if pk == "ollama":
            models = _ollama_models()
        for name, t in models:
            if t == tier:
                label = _PROVIDER_META.get(pk, {}).get("label", pk).split("(")[0].strip()
                matches.append((pk, name, label))
    if not matches:
        print(f"[PULSE] No {tier} models available.")
        return
    idx = _pick(f"{tier.title()} Models (all providers):",
                [(name, label) for _, name, label in matches])
    if idx is None:
        return
    pk, model, _ = matches[idx]
    if count is None:
        count = _ask_count()
    if not count:
        return
    count = max(1, min(count, 10))
    print(f"\n  Launching {count}x {model} [{tier}]...")
    launch_swarm_worker(pk, model, tier, count, skills_used=skills_used)

def cmd_swarm(args: list[str]):
    if args and args[0] == "stop":
        return _swarm_stop()
    if args and args[0] == "status":
        return _swarm_status()
    if args and args[0] == "watch":
        return _swarm_watch()
    skills_used = "--no-skills" not in args  # D10: always-on, opt-OUT via --no-skills
    if "--no-skills" in args:
        args.remove("--no-skills")
    ensure_orchestrator()
    # Tier shortcut: pulse swarm [count] fast|standard|premium
    TIER_NAMES = {"fast", "standard", "premium"}
    tier_arg = None
    count_arg = None
    remaining = []
    for a in args:
        if a in TIER_NAMES:
            tier_arg = a
        elif a.isdigit():
            count_arg = int(a)
        else:
            remaining.append(a)
    if tier_arg:
        return _swarm_by_tier(tier_arg, count_arg, skills_used)
    # Shortcut: model name given directly (e.g. "pulse swarm mistral:7b 3")
    if remaining:
        hint = remaining[0]
        provider_key = _find_provider_for_model(hint)
        if provider_key:
            tier = detect_tier(hint)
            count = count_arg or _ask_count()
            if not count:
                return
            print(f"\n  Starting {count}x {hint} workers ({tier} tier)...")
            launch_swarm_worker(provider_key, hint, tier, count, skills_used=skills_used)
            return
        # Check if it's a provider name shortcut
        provider_aliases = {"codex": "openai", "anthropic": "claude", "google": "gemini"}
        resolved_hint = provider_aliases.get(hint, hint)
        if resolved_hint in _PROVIDER_META:
            providers = _detect_providers()
            if any(pk == resolved_hint for pk, _ in providers):
                return _pick_model_and_launch(resolved_hint, skills_used=skills_used)
            print(f"[PULSE] Provider '{hint}' not available. Check API keys or CLI installation.")
            return
        print(f"[PULSE] Unknown argument: {hint}")
        return
    # Full interactive flow
    providers = _detect_providers()
    if not providers:
        print("[PULSE] No providers available.")
        print("  Set GEMINI_API_KEY, ANTHROPIC_API_KEY, or install agent CLIs.")
        return
    print("\n  PULSE Swarm -- Build Your Crew")
    _crew_builder(providers, skills_used=skills_used)

if __name__ == "__main__":
    import argparse
    from pulse.cli.autopilot.loop import _run_autopilot_loop
    parser = argparse.ArgumentParser(description="PULSE Swarm Autopilot (background)")
    parser.add_argument("--autopilot", action="store_true", required=True)
    parser.add_argument("--agent", required=True, choices=["gemini", "claude", "codex", "ollama-cloud"])
    parser.add_argument("--model", required=True)
    parser.add_argument("--tier", default="fast", choices=["fast", "standard", "premium"])
    parser.add_argument("--id", default=None)
    parser.add_argument("--skills-used", action="store_true")
    ap = parser.parse_args()
    ensure_orchestrator()
    sys.exit(_run_autopilot_loop(
        ap.agent, ap.model, ap.tier, worker_id=ap.id, skills_used=ap.skills_used) or 0)
