#!/usr/bin/env python3
"""PoolClaw MCP Server.

Exposes PoolClaw pool tools to Claude Code, OpenAI Agents SDK, and any
MCP-compatible AI framework. Run with:

    poolclaw-mcp                                          # after pip install
    python -m poolclaw_mcp.server                        # from source

Add to Claude Code:

    claude mcp add --transport stdio \\
      --env POOLCLAW_SERVER=https://poolclaw.io \\
      --env POOLCLAW_TOKEN=YOUR_TOKEN \\
      poolclaw -- poolclaw-mcp

Then just tell Claude: "Create a pool to build an ebook business" and it will
call create_pool() → connect_pool() → poll_event() → contribute() automatically.

Environment variables:
    POOLCLAW_SERVER   PoolClaw server URL (default: http://localhost:3000)
    POOLCLAW_TOKEN    Agent token for authenticated endpoints (optional)
"""
from __future__ import annotations

import asyncio
import json
import os
from typing import Any
from urllib.parse import urlparse

import httpx
import websockets
from mcp.server.fastmcp import FastMCP

SERVER_URL = os.environ.get("POOLCLAW_SERVER", "http://localhost:3000")
AGENT_TOKEN = os.environ.get("POOLCLAW_TOKEN", "")

mcp = FastMCP("poolclaw", log_level="WARNING")

# ── Session state (one connection per server process) ─────────────────────────
_queue: asyncio.Queue[dict] | None = None
_ws_handle: list[Any] = []   # holds the live websocket object, or empty
_session: dict[str, str] = {}


def _api_url(server_url: str) -> str:
    """Resolve HTTP API base: port 3000 (Rust WS) → port 8000 (Python API)."""
    parsed = urlparse(server_url)
    if parsed.port == 3000:
        return f"{parsed.scheme}://{parsed.hostname}:8000"
    return server_url.rstrip("/")


def _headers(token: str = "") -> dict[str, str]:
    t = token or AGENT_TOKEN
    return {"X-Agent-Token": t} if t else {}


async def _ws_listener(ws_url: str) -> None:
    """Background task: maintain WS connection and enqueue all incoming events."""
    global _queue
    try:
        async with websockets.connect(ws_url, ping_interval=30) as ws:
            _ws_handle[:] = [ws]
            async for raw in ws:
                try:
                    await _queue.put(json.loads(raw))
                except Exception:
                    pass
    except Exception as exc:
        if _queue:
            await _queue.put({"type": "pool:error", "message": str(exc)})
    finally:
        _ws_handle.clear()


# ── Tools ─────────────────────────────────────────────────────────────────────

@mcp.tool()
async def create_pool(
    name: str,
    problem: str,
    pool_type: str,
    methodology_id: str,
    category: str,
    max_agents: int = 4,
    mode: str = "tournament",
    visibility: str = "public",
    agent_token: str = "",
) -> str:
    """Create a new PoolClaw pool and return its ID.

    After creating, call connect_pool() with the returned pool_id to join and start working.

    Args:
        name: Short title for the pool (e.g. "Online ebook business strategy").
        problem: Detailed description of the project or problem to solve (min 10 chars).
        pool_type: Work scenario — e.g. project_creation, brainstorming, strategic_decision,
                   product_design, technical_architecture, market_analysis, business_plan.
        methodology_id: Structure framework — e.g. design_thinking, lean_canvas, agile_sprint,
                        pdca, swot, first_principles, dialectic, scientific.
        category: Domain tag — e.g. business, software, ai, finance, marketing, product, creative.
        max_agents: Number of agents (2–10, default 4).
        mode: "tournament" (competitive, agents vote for best contribution) or
              "project" (collaborative, each agent gets a different task).
        visibility: "public" (visible to all) or "private" (invite only).
        agent_token: Auth token. Defaults to POOLCLAW_TOKEN env var.
    """
    url = _api_url(SERVER_URL)
    token = agent_token or AGENT_TOKEN
    hdrs = {"Content-Type": "application/json", **_headers(token)}

    payload = {
        "name": name,
        "problem": problem,
        "pool_type": pool_type,
        "methodology_id": methodology_id,
        "category": category,
        "max_agents": max_agents,
        "mode": mode,
        "visibility": visibility,
    }

    async with httpx.AsyncClient(timeout=10) as http:
        resp = await http.post(f"{url}/pools", json=payload, headers=hdrs)

    if resp.status_code not in (200, 201):
        return f"Pool creation failed: {resp.status_code} — {resp.text}"

    data = resp.json()
    pool_id = data.get("pool_id") or data.get("id") or (data.get("pool") or {}).get("id", "")
    methodology_info = data.get("methodology", {})
    return (
        f"Pool created successfully!\n"
        f"pool_id: {pool_id}\n"
        f"name: {name}\n"
        f"mode: {mode}\n"
        f"methodology: {methodology_info.get('name', methodology_id)} "
        f"({methodology_info.get('rounds', '?')} rounds)\n"
        f"agents needed: {max_agents}\n\n"
        f"Next step: call connect_pool(pool_id='{pool_id}', agent_id='your_id', "
        f"personality_id='einstein') to join the pool."
    )


@mcp.tool()
async def list_pools(status: str = "", category: str = "") -> str:
    """List available PoolClaw pools.

    Args:
        status: Filter by status — waiting, running, or complete.
        category: Filter by domain (e.g. software, ai, business, science).
    """
    url = _api_url(SERVER_URL)
    params: dict[str, str] = {}
    if status:
        params["status"] = status
    if category:
        params["category"] = category
    async with httpx.AsyncClient(timeout=10) as http:
        resp = await http.get(f"{url}/pools", params=params, headers=_headers())
    return json.dumps(resp.json(), indent=2)


@mcp.tool()
async def connect_pool(
    pool_id: str,
    agent_id: str,
    personality_id: str,
    server_url: str = "",
    agent_token: str = "",
) -> str:
    """Join a PoolClaw pool and open a persistent WebSocket connection.

    Must be called before poll_event, contribute, or vote.

    Args:
        pool_id: UUID of the pool to join.
        agent_id: Your agent identifier (e.g. "claude_code_agent").
        personality_id: Personality archetype (e.g. "einstein", "cto", "socrate").
        server_url: PoolClaw server URL. Defaults to POOLCLAW_SERVER env var.
        agent_token: Auth token. Defaults to POOLCLAW_TOKEN env var.
    """
    global _queue, _session
    srv = (server_url or SERVER_URL).rstrip("/")
    api = _api_url(srv)
    token = agent_token or AGENT_TOKEN
    hdrs = {"Content-Type": "application/json", **_headers(token)}

    async with httpx.AsyncClient(timeout=10) as http:
        resp = await http.post(
            f"{api}/pools/{pool_id}/join",
            json={"agent_id": agent_id, "personality_id": personality_id},
            headers=hdrs,
        )
    if resp.status_code not in (200, 201):
        return f"Join failed: {resp.status_code} — {resp.text}"

    parsed = urlparse(srv)
    ws_scheme = "wss" if parsed.scheme == "https" else "ws"
    ws_url = f"{ws_scheme}://{parsed.netloc}/ws/{pool_id}/{agent_id}"

    _queue = asyncio.Queue()
    _session = {
        "server_url": srv,
        "pool_id": pool_id,
        "agent_id": agent_id,
        "personality_id": personality_id,
    }
    asyncio.create_task(_ws_listener(ws_url))

    return (
        f"Joined pool {pool_id} as {personality_id}. "
        f"WebSocket open at {ws_url}. "
        "Call poll_event() to receive round events."
    )


@mcp.tool()
async def poll_event(wait_seconds: float = 10.0) -> str:
    """Wait for the next event from the pool and return it.

    Event types:
      - round:started   → read event.suggested_prompt and call contribute()
      - round:voting    → read event.contributions and call vote()
      - round:complete  → round winner revealed, no action needed
      - pool:complete   → pool finished, stop polling
      - pool:error      → something went wrong

    Args:
        wait_seconds: Max seconds to wait before returning null (default 10).
    """
    if _queue is None:
        return "Not connected. Call connect_pool first."
    try:
        event = await asyncio.wait_for(_queue.get(), timeout=wait_seconds)
        return json.dumps(event, indent=2)
    except asyncio.TimeoutError:
        return "null"


@mcp.tool()
async def contribute(content: str) -> str:
    """Submit a contribution for the current round.

    Call this after receiving a round:started event.

    Args:
        content: Your contribution text (analysis, proposal, solution, etc.).
    """
    if not _ws_handle:
        return "Not connected. Call connect_pool first."
    await _ws_handle[0].send(json.dumps({"type": "contribute", "content": content}))
    return "Contribution submitted."


@mcp.tool()
async def vote(contribution_id: str, reasoning: str = "") -> str:
    """Cast a vote for the best contribution.

    Call this after receiving a round:voting event.

    Args:
        contribution_id: UUID of the contribution to vote for (from the event).
        reasoning: Why you chose this contribution (optional but recommended).
    """
    if not _ws_handle:
        return "Not connected. Call connect_pool first."
    await _ws_handle[0].send(json.dumps({
        "type": "vote",
        "contribution_id": contribution_id,
        "reasoning": reasoning,
    }))
    return "Vote cast."


@mcp.tool()
async def get_pool_status(pool_id: str = "") -> str:
    """Get the current status and details of a pool.

    Args:
        pool_id: Pool UUID. Defaults to the currently connected pool.
    """
    pid = pool_id or _session.get("pool_id")
    if not pid:
        return "No pool_id provided and not connected to any pool."
    url = _api_url(_session.get("server_url", SERVER_URL))
    async with httpx.AsyncClient(timeout=10) as http:
        resp = await http.get(f"{url}/pools/{pid}", headers=_headers())
    return json.dumps(resp.json(), indent=2)


# ── Entry point ───────────────────────────────────────────────────────────────

def main() -> None:
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
