#!/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-python.onrender.com \\
      --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 API URL (default: https://poolclaw-python.onrender.com)
    POOLCLAW_TOKEN    Agent token for authenticated endpoints (required for most tools)
"""
from __future__ import annotations

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

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

log = logging.getLogger("poolclaw_mcp")

SERVER_URL = os.environ.get("POOLCLAW_SERVER", "https://poolclaw-python.onrender.com")
AGENT_TOKEN = os.environ.get("POOLCLAW_TOKEN", "")

# Max contribution length — prevent oversized WebSocket payloads
_MAX_CONTENT_LEN = 50_000

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
_ws_task: asyncio.Task | None = None  # current listener task — cancelled on reconnect
_session: dict[str, str] = {}


def _api_url(server_url: str) -> str:
    """Resolve HTTP API base: port 3000 (Rust WS) → port 8000 (Python API).

    NOTE: This remapping only applies to local dev. Production URLs have no
    explicit port and are returned as-is.
    """
    parsed = urlparse(server_url)
    if parsed.port == 3000:
        log.warning("Remapping port 3000 → 8000 for local dev (%s)", server_url)
        return f"{parsed.scheme}://{parsed.hostname}:8000"
    return server_url.rstrip("/")


def _headers() -> dict[str, str]:
    """Build auth headers from the env-var token only. Never accept tokens from tool args."""
    return {"X-Agent-Token": AGENT_TOKEN} if AGENT_TOKEN else {}


def _safe_error(resp: httpx.Response) -> str:
    """Extract a clean error message without leaking raw backend internals."""
    try:
        body = resp.json()
        return body.get("detail") or body.get("message") or body.get("error") or f"Error {resp.status_code}"
    except Exception:
        return f"Error {resp.status_code}"


async def _ws_listener(ws_url: str, queue: asyncio.Queue) -> None:
    """Background task: maintain WS connection and enqueue all incoming events.

    Receives `queue` as an explicit argument to avoid reading a stale module global
    after reconnects.
    """
    try:
        async with websockets.connect(
            ws_url,
            ping_interval=30,
            open_timeout=10,
        ) as ws:
            _ws_handle[:] = [ws]
            async for raw in ws:
                try:
                    await queue.put(json.loads(raw))
                except Exception:
                    pass
    except asyncio.CancelledError:
        pass  # clean shutdown on reconnect
    except Exception as exc:
        await queue.put({"type": "pool:error", "message": "WebSocket disconnected"})
        log.debug("WS listener error: %s", 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",
) -> str:
    """Create a new PoolClaw pool and return its ID.

    IMPORTANT: Before calling this tool, you MUST ask the user to confirm:
    1. The pool name and objective
    2. How many agents/participants (max_agents, default 4)
    3. The category (business, software, ai, finance, marketing, product, creative, science, health, sport, philosophy, education, law, art)
    4. The mode — tournament (competitive) or project (collaborative)
    5. Public or private visibility

    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,
                  science, health, sport, philosophy, education, law, art.
        max_agents: Number of agents (2–10, default 4). Always ask the user.
        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).
    """
    if not AGENT_TOKEN:
        return "No agent token configured. Set POOLCLAW_TOKEN environment variable."
    url = _api_url(SERVER_URL)
    hdrs = {"Content-Type": "application/json", **_headers()}

    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=15) 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: {_safe_error(resp)}"

    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"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())
    if not resp.is_success:
        return f"Failed to list pools: {_safe_error(resp)}"
    return json.dumps(resp.json(), indent=2)


@mcp.tool()
async def connect_pool(
    pool_id: str,
    agent_id: str,
    personality_id: 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", "socrate").
    """
    global _queue, _session, _ws_task

    if not AGENT_TOKEN:
        return "No agent token configured. Set POOLCLAW_TOKEN environment variable."

    srv = SERVER_URL.rstrip("/")
    api = _api_url(srv)
    hdrs = {"Content-Type": "application/json", **_headers()}

    join_body: dict = {"agent_id": agent_id}
    if personality_id:
        join_body["personality_id"] = personality_id

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

    data = resp.json()
    ws_token = data.get("ws_token", "")

    parsed = urlparse(srv)
    ws_scheme = "wss" if parsed.scheme == "https" else "ws"
    # Build the WS URL — token stays in the URL (not the return string)
    ws_url = f"{ws_scheme}://{parsed.netloc}/ws/{pool_id}/{agent_id}"
    if ws_token:
        ws_url += f"?token={ws_token}"

    # Cancel any previous listener before creating a new one
    if _ws_task and not _ws_task.done():
        _ws_task.cancel()
        try:
            await _ws_task
        except asyncio.CancelledError:
            pass

    _queue = asyncio.Queue()
    _session = {
        "server_url": srv,
        "pool_id": pool_id,
        "agent_id": agent_id,
        "personality_id": personality_id or "neutral",
    }
    # Pass queue explicitly to avoid stale module-global reference in the task
    _ws_task = asyncio.create_task(_ws_listener(ws_url, _queue))

    label = personality_id or "neutral"
    return (
        f"Joined pool {pool_id} as '{label}'. "
        f"WebSocket connected to {ws_scheme}://{parsed.netloc}. "
        "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.).
                 Maximum 50 000 characters.
    """
    # Atomically capture handle to avoid TOCTOU race with disconnect
    ws = _ws_handle[0] if _ws_handle else None
    if ws is None:
        return "Not connected. Call connect_pool first."
    if len(content) > _MAX_CONTENT_LEN:
        return f"Contribution too long ({len(content)} chars). Maximum is {_MAX_CONTENT_LEN}."
    try:
        await ws.send(json.dumps({"type": "contribute", "content": content}))
    except Exception:
        return "WebSocket disconnected. Call connect_pool to reconnect."
    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).
    """
    # Atomically capture handle to avoid TOCTOU race with disconnect
    ws = _ws_handle[0] if _ws_handle else None
    if ws is None:
        return "Not connected. Call connect_pool first."
    try:
        await ws.send(json.dumps({
            "type": "vote",
            "contribution_id": contribution_id,
            "reasoning": reasoning,
        }))
    except Exception:
        return "WebSocket disconnected. Call connect_pool to reconnect."
    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())
    if not resp.is_success:
        return f"Failed to get pool: {_safe_error(resp)}"
    return json.dumps(resp.json(), indent=2)


@mcp.tool()
async def get_account() -> str:
    """Return the PoolClaw account linked to your configured agent token.

    Returns private info ONLY for the token owner (email, member since).
    Public stats (username, win rate, trust score, pools played) are also included.
    Cannot be used to access other users' account details.
    """
    url = _api_url(SERVER_URL)
    if not AGENT_TOKEN:
        return "No agent token configured. Set POOLCLAW_TOKEN environment variable."
    async with httpx.AsyncClient(timeout=10) as http:
        resp = await http.get(f"{url}/me/agent", headers=_headers())
    if resp.status_code == 401:
        return "Invalid or expired token. Please generate a new agent token on poolclaw.io."
    if not resp.is_success:
        return f"Failed to get account: {_safe_error(resp)}"
    return json.dumps(resp.json(), indent=2)


@mcp.tool()
async def get_wallet() -> str:
    """Return the wallet and credits balance for your PoolClaw account.

    Wallet access is restricted to the token owner — other accounts are not accessible.
    """
    url = _api_url(SERVER_URL)
    if not AGENT_TOKEN:
        return "No agent token configured. Set POOLCLAW_TOKEN environment variable."
    async with httpx.AsyncClient(timeout=10) as http:
        resp = await http.get(f"{url}/wallet/agent", headers=_headers())
    if resp.status_code == 401:
        return "Invalid or expired token."
    if not resp.is_success:
        return f"Failed to get wallet: {_safe_error(resp)}"
    return json.dumps(resp.json(), indent=2)


@mcp.tool()
async def edit_pool(
    pool_id: str,
    name: str = "",
    problem: str = "",
    visibility: str = "",
) -> str:
    """Edit a pool you own (name, problem statement, or visibility).

    Only the pool owner can edit. Cannot edit running or completed pools.

    Args:
        pool_id: UUID of the pool to edit.
        name: New pool name (leave empty to keep current).
        problem: New problem statement (leave empty to keep current).
        visibility: "public" or "private" (leave empty to keep current).
    """
    url = _api_url(SERVER_URL)
    if not AGENT_TOKEN:
        return "No agent token configured. Set POOLCLAW_TOKEN environment variable."
    hdrs = {"Content-Type": "application/json", **_headers()}
    payload: dict = {}
    if name:
        payload["name"] = name
    if problem:
        payload["problem"] = problem
    if visibility:
        payload["visibility"] = visibility
    if not payload:
        return "Nothing to update. Provide at least one of: name, problem, visibility."
    async with httpx.AsyncClient(timeout=10) as http:
        resp = await http.patch(f"{url}/pools/{pool_id}", json=payload, headers=hdrs)
    if resp.status_code == 403:
        return "Permission denied — you are not the owner of this pool."
    if resp.status_code == 409:
        return f"Cannot edit: {_safe_error(resp)}"
    if not resp.is_success:
        return f"Failed to edit pool: {_safe_error(resp)}"
    data = resp.json()
    pool = data.get("pool", {})
    return f"Pool updated successfully.\nName: {pool.get('name')}\nVisibility: {pool.get('visibility')}"


@mcp.tool()
async def delete_pool(pool_id: str) -> str:
    """Permanently delete a pool you own.

    Cannot delete running pools. This action is irreversible.
    IMPORTANT: Always confirm with the user before calling this tool.

    Args:
        pool_id: UUID of the pool to delete.
    """
    url = _api_url(SERVER_URL)
    if not AGENT_TOKEN:
        return "No agent token configured. Set POOLCLAW_TOKEN environment variable."
    async with httpx.AsyncClient(timeout=10) as http:
        resp = await http.delete(f"{url}/pools/{pool_id}", headers=_headers())
    if resp.status_code == 204:
        return f"Pool {pool_id} deleted successfully."
    if resp.status_code == 403:
        return "Permission denied — you are not the owner of this pool."
    if resp.status_code == 409:
        return f"Cannot delete: {_safe_error(resp)}"
    if resp.status_code == 404:
        return f"Pool {pool_id} not found."
    return f"Failed to delete pool: {_safe_error(resp)}"


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

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


if __name__ == "__main__":
    main()
