"""Phase 79: MCP Microservice — Swarm & Agent Management.

Swarm formation, local LLM agent spawning/killing, backend discovery,
quick chat, compute grid, and file watcher tools.
"""

import asyncio
import json

from mcp.types import TextContent, Tool, ToolAnnotations

_SAFE_WRITE = ToolAnnotations(readOnlyHint=False, destructiveHint=False, idempotentHint=False, openWorldHint=False)
_SAFE_READ = ToolAnnotations(readOnlyHint=True, destructiveHint=False, idempotentHint=True, openWorldHint=False)

TOOLS = [
    Tool(
        name="broker_form_swarm",
        description="Form a swarm from submitted bids. Assigns a lead and workers.",
        inputSchema={
            "type": "object",
            "properties": {
                "swarm_id": {"type": "string", "description": "Swarm ID to form"},
            },
            "required": ["swarm_id"],
        },
        annotations=_SAFE_WRITE,
    ),
    Tool(
        name="broker_discover_backends",
        description="Discover available LLM backends (Ollama, LM Studio, etc) on this machine.",
        inputSchema={"type": "object", "properties": {}},
        annotations=_SAFE_READ,
    ),
    Tool(
        name="broker_spawn_local",
        description="Spawn a local LLM agent as a background process. It will monitor a channel and respond autonomously.",
        inputSchema={
            "type": "object",
            "properties": {
                "name": {"type": "string", "description": "Agent name (e.g. 'worker-1')"},
                "backend": {"type": "string", "description": "LLM backend", "default": "ollama"},
                "model": {"type": "string", "description": "Model name", "default": "llama3"},
                "channel": {"type": "string", "description": "Channel to monitor", "default": "general"},
            },
            "required": ["name"],
        },
        annotations=_SAFE_WRITE,
    ),
    Tool(
        name="broker_list_local_agents",
        description="List all running local LLM agents.",
        inputSchema={"type": "object", "properties": {}},
        annotations=_SAFE_READ,
    ),
    Tool(
        name="broker_kill_local",
        description="Kill a running local LLM agent by name.",
        inputSchema={
            "type": "object",
            "properties": {
                "name": {"type": "string", "description": "Agent name to kill"},
            },
            "required": ["name"],
        },
        annotations=_SAFE_WRITE,
    ),
    Tool(
        name="broker_quick_chat",
        description="Send a one-shot prompt to a local LLM and get a response. Auto-discovers available backends.",
        inputSchema={
            "type": "object",
            "properties": {
                "prompt": {"type": "string", "description": "The prompt to send"},
                "backend": {"type": "string", "description": "Backend to use (or 'auto')", "default": "auto"},
                "model": {"type": "string", "description": "Model to use (or 'auto')", "default": "auto"},
            },
            "required": ["prompt"],
        },
        annotations=_SAFE_READ,
    ),
    # -- Compute Grid --
    Tool(
        name="broker_grid_join",
        description="Join the compute grid as a worker node.",
        inputSchema={
            "type": "object",
            "properties": {
                "capabilities": {"type": "array", "items": {"type": "string"}, "description": "What this node can do"},
                "max_concurrent": {"type": "integer", "default": 2},
            },
            "required": ["capabilities"],
        },
        annotations=_SAFE_WRITE,
    ),
    Tool(
        name="broker_grid_stats",
        description="View compute grid statistics.",
        inputSchema={"type": "object", "properties": {}},
        annotations=_SAFE_READ,
    ),
    Tool(
        name="broker_grid_submit",
        description="Submit a job to the compute grid.",
        inputSchema={
            "type": "object",
            "properties": {
                "job_type": {"type": "string", "description": "Type of job (e.g., 'map', 'reduce')"},
                "payload": {"type": "object", "description": "Job payload data"},
                "priority": {"type": "string", "default": "normal"},
            },
            "required": ["job_type", "payload"],
        },
        annotations=_SAFE_WRITE,
    ),
    Tool(
        name="broker_grid_status",
        description="Check status of a grid job.",
        inputSchema={
            "type": "object",
            "properties": {
                "job_id": {"type": "string", "description": "Grid job ID to check"},
            },
            "required": ["job_id"],
        },
        annotations=_SAFE_READ,
    ),
    # -- File Watcher --
    Tool(
        name="broker_watch",
        description="Register a file watcher on a directory. When files matching the pattern change, a notification is posted to the specified channel.",
        inputSchema={
            "type": "object",
            "properties": {
                "directory": {"type": "string", "description": "Absolute path to directory to watch"},
                "pattern": {"type": "string", "description": "Glob pattern (e.g. '*.py', '*.json')", "default": "*"},
                "channel": {"type": "string", "description": "Channel for change notifications", "default": "file-events"},
                "recursive": {"type": "boolean", "description": "Watch subdirectories too", "default": True},
            },
            "required": ["directory"],
        },
        annotations=_SAFE_WRITE,
    ),
    Tool(
        name="broker_unwatch",
        description="Remove a file watcher by its watch ID.",
        inputSchema={
            "type": "object",
            "properties": {"watch_id": {"type": "string", "description": "The watch ID to remove"}},
            "required": ["watch_id"],
        },
        annotations=_SAFE_WRITE,
    ),
    Tool(
        name="broker_list_watches",
        description="List all active file watchers.",
        inputSchema={"type": "object", "properties": {}},
        annotations=_SAFE_READ,
    ),
    Tool(
        name="broker_check_watches",
        description="Check all file watchers for changes. Posts notifications to channels for any detected changes.",
        inputSchema={"type": "object", "properties": {}},
        annotations=_SAFE_READ,
    ),
]


# ── Handlers ──────────────────────────────────────────────────────────────


async def _handle_form_swarm(args: dict) -> list[TextContent]:
    from brynq import swarm
    success = await asyncio.to_thread(swarm.form_swarm, args["swarm_id"])
    if success:
        swarm_path = swarm.SWARM_DIR / f"{args['swarm_id']}.json"
        try:
            data = json.loads(swarm_path.read_text("utf-8"))
            members = data.get("members", {})
            lead = [k for k, v in members.items() if v == "lead"]
            lead_str = lead[0] if lead else "unknown"
            return [TextContent(type="text", text=f"Swarm formed! Lead: @{lead_str}. Members: {', '.join(members.keys())}")]
        except (FileNotFoundError, json.JSONDecodeError, OSError):
            return [TextContent(type="text", text="Swarm formed successfully.")]
    return [TextContent(type="text", text="Failed to form swarm. Either no bids or swarm not found.")]


async def _handle_discover_backends(args: dict) -> list[TextContent]:
    try:
        from brynq import llm_router
    except ImportError:
        return [TextContent(type="text", text="Error: llm_router module not installed. Install brynq[llm] extras.")]
    try:
        backends = await asyncio.to_thread(llm_router.discover_backends)
        if not backends:
            return [TextContent(type="text", text="No LLM backends discovered. Install Ollama, LM Studio, or set API keys for Claude/Gemini.")]
        lines = ["Discovered LLM Backends:"]
        for b in backends:
            name = b.get("name", "unknown")
            status = b.get("status", "unknown")
            models = b.get("models", [])
            model_str = ", ".join(models[:10]) if models else "none"
            if len(models) > 10:
                model_str += f" (+{len(models) - 10} more)"
            lines.append(f"  [{status.upper()}] {name}: {model_str}")
        return [TextContent(type="text", text="\n".join(lines))]
    except Exception as e:
        return [TextContent(type="text", text=f"Error discovering backends: {e}")]


async def _handle_spawn_local(args: dict) -> list[TextContent]:
    try:
        from brynq import spawner
    except ImportError:
        return [TextContent(type="text", text="Error: spawner module not installed. Install brynq[llm] extras.")]
    from brynq.server import BROKER_DIR
    name = args.get("name", "")
    if not name:
        return [TextContent(type="text", text="Error: name is required")]
    backend = args.get("backend", "ollama")
    model = args.get("model", "llama3")
    channel = args.get("channel", "general")
    try:
        result = await asyncio.to_thread(
            spawner.spawn_agent,
            name=name,
            backend=backend,
            model=model,
            channel=channel,
            broker_dir=str(BROKER_DIR),
        )
        status = result.get("status", "unknown")
        if status == "running":
            pid = result.get("pid", "?")
            return [TextContent(type="text", text=(
                f"Agent '{name}' spawned successfully.\n"
                f"  Backend: {backend}\n"
                f"  Model: {model}\n"
                f"  Channel: #{channel}\n"
                f"  PID: {pid}"
            ))]
        else:
            error = result.get("error", "unknown error")
            return [TextContent(type="text", text=f"Failed to spawn agent '{name}': {error}")]
    except Exception as e:
        return [TextContent(type="text", text=f"Error spawning agent: {e}")]


async def _handle_list_local_agents(args: dict) -> list[TextContent]:
    try:
        from brynq import spawner
    except ImportError:
        return [TextContent(type="text", text="Error: spawner module not installed. Install brynq[llm] extras.")]
    try:
        agents = await asyncio.to_thread(spawner.list_agents)
        if not agents:
            return [TextContent(type="text", text="No local LLM agents running.")]
        lines = ["Running Local LLM Agents:"]
        for a in agents:
            name = a.get("name", "unknown")
            model = a.get("model", "?")
            backend = a.get("backend", "?")
            status = a.get("status", "?")
            uptime = a.get("uptime", "?")
            channel = a.get("channel", "?")
            pid = a.get("pid", "?")
            lines.append(f"  {name} [{status.upper()}] — {backend}/{model} on #{channel} (PID {pid}, uptime {uptime})")
        return [TextContent(type="text", text="\n".join(lines))]
    except Exception as e:
        return [TextContent(type="text", text=f"Error listing agents: {e}")]


async def _handle_kill_local(args: dict) -> list[TextContent]:
    try:
        from brynq import spawner
    except ImportError:
        return [TextContent(type="text", text="Error: spawner module not installed. Install brynq[llm] extras.")]
    name = args.get("name", "")
    if not name:
        return [TextContent(type="text", text="Error: name is required")]
    try:
        result = await asyncio.to_thread(spawner.kill_agent, name=name)
        if result.get("killed"):
            return [TextContent(type="text", text=f"Agent '{name}' killed successfully.")]
        else:
            reason = result.get("error", "agent not found")
            return [TextContent(type="text", text=f"Could not kill agent '{name}': {reason}")]
    except Exception as e:
        return [TextContent(type="text", text=f"Error killing agent: {e}")]


async def _handle_quick_chat(args: dict) -> list[TextContent]:
    try:
        from brynq import llm_router
    except ImportError:
        return [TextContent(type="text", text="Error: llm_router module not installed. Install brynq[llm] extras.")]
    prompt = args.get("prompt", "")
    if not prompt:
        return [TextContent(type="text", text="Error: prompt is required")]
    backend = args.get("backend", "auto")
    model = args.get("model", "auto")
    try:
        result = await asyncio.to_thread(
            llm_router.quick_chat,
            prompt=prompt,
            backend=backend,
            model=model,
        )
        used_backend = result.get("backend", "unknown")
        used_model = result.get("model", "unknown")
        response = result.get("response", "")
        return [TextContent(type="text", text=f"[{used_backend}/{used_model}]\n\n{response}")]
    except Exception as e:
        return [TextContent(type="text", text=f"Error in quick chat: {e}")]


async def _handle_grid_join(args: dict) -> list[TextContent]:
    from brynq.compute_grid import grid_join
    result = await asyncio.to_thread(
        grid_join,
        capabilities=args.get("capabilities", []),
        max_concurrent=args.get("max_concurrent", 2),
    )
    return [TextContent(type="text", text=json.dumps(result, indent=2, default=str))]


async def _handle_grid_stats(args: dict) -> list[TextContent]:
    from brynq.compute_grid import grid_stats
    result = await asyncio.to_thread(grid_stats)
    return [TextContent(type="text", text=json.dumps(result, indent=2, default=str))]


async def _handle_grid_submit(args: dict) -> list[TextContent]:
    from brynq.compute_grid import grid_submit
    result = await asyncio.to_thread(
        grid_submit,
        job_type=args.get("job_type", ""),
        payload=args.get("payload", {}),
        priority=args.get("priority", "normal"),
    )
    return [TextContent(type="text", text=json.dumps(result, indent=2, default=str))]


async def _handle_grid_status(args: dict) -> list[TextContent]:
    from brynq.compute_grid import grid_status
    result = await asyncio.to_thread(grid_status, job_id=args.get("job_id", ""))
    return [TextContent(type="text", text=json.dumps(result, indent=2, default=str))]


async def _handle_watch(args: dict) -> list[TextContent]:
    from brynq.file_watcher import register_watch
    directory = args.get("directory", "")
    if not directory:
        return [TextContent(type="text", text="Error: directory is required")]
    watch = await asyncio.to_thread(
        register_watch,
        directory=directory,
        pattern=args.get("pattern", "*"),
        channel=args.get("channel", "file-events"),
        recursive=args.get("recursive", True),
    )
    return [TextContent(type="text", text=f"Watch registered: {watch['watch_id']}\n  Directory: {watch['directory']}\n  Pattern: {watch['pattern']}\n  Channel: #{watch['channel']}\n  Recursive: {watch['recursive']}")]


async def _handle_unwatch(args: dict) -> list[TextContent]:
    from brynq.file_watcher import remove_watch
    watch_id = args.get("watch_id", "")
    if not watch_id:
        return [TextContent(type="text", text="Error: watch_id is required")]
    removed = await asyncio.to_thread(remove_watch, watch_id)
    if removed:
        return [TextContent(type="text", text=f"Watch {watch_id} removed")]
    return [TextContent(type="text", text=f"Watch {watch_id} not found")]


async def _handle_list_watches(args: dict) -> list[TextContent]:
    from brynq.file_watcher import list_watches
    watches = await asyncio.to_thread(list_watches)
    if not watches:
        return [TextContent(type="text", text="No active file watches")]
    lines = ["Active file watches:"]
    for w in watches:
        lines.append(f"  [{w['watch_id']}] {w['directory']} ({w['pattern']}) -> #{w['channel']}")
    return [TextContent(type="text", text="\n".join(lines))]


async def _handle_check_watches(args: dict) -> list[TextContent]:
    from brynq.file_watcher import check_all_watches
    results = await asyncio.to_thread(check_all_watches)
    if not results:
        return [TextContent(type="text", text="No file watches configured")]
    lines = ["File watch check results:"]
    for r in results:
        total = r["total_changes"]
        if total > 0:
            lines.append(f"  [{r['watch_id']}] {r['directory']}: {total} changes (+{len(r['created'])} ~{len(r['modified'])} -{len(r['deleted'])})")
        else:
            lines.append(f"  [{r['watch_id']}] {r['directory']}: no changes")
    return [TextContent(type="text", text="\n".join(lines))]


HANDLERS = {
    "broker_form_swarm": _handle_form_swarm,
    "broker_discover_backends": _handle_discover_backends,
    "broker_spawn_local": _handle_spawn_local,
    "broker_list_local_agents": _handle_list_local_agents,
    "broker_kill_local": _handle_kill_local,
    "broker_quick_chat": _handle_quick_chat,
    "broker_grid_join": _handle_grid_join,
    "broker_grid_stats": _handle_grid_stats,
    "broker_grid_submit": _handle_grid_submit,
    "broker_grid_status": _handle_grid_status,
    "broker_watch": _handle_watch,
    "broker_unwatch": _handle_unwatch,
    "broker_list_watches": _handle_list_watches,
    "broker_check_watches": _handle_check_watches,
}
