import json
import logging
import os
import shutil
import time
import uuid
from types import SimpleNamespace
from typing import Any, Literal

from fastmcp import Context, FastMCP
from fastmcp.exceptions import ToolError
from fastmcp.server.lifespan import lifespan
from fastmcp.server.middleware import Middleware

from claude_teams import acpx_client, messaging, tasks, teams
from claude_teams.acpx_client import AcpxError
from claude_teams.models import (
    COLOR_PALETTE,
    InboxMessage,
    SendMessageResult,
    ShutdownApproved,
    SpawnResult,
    TeammateMember,
)
from claude_teams.spawner import (
    discover_harness_binary,
    spawn_teammate,
)

logger = logging.getLogger(__name__)

# NOTE(victor): Mutated by both app_lifespan and HarnessDetectionMiddleware.
# Safe under stdio (single session). Racy under SSE/streamable HTTP.
_lifespan_state: dict[str, Any] = {}
_spawn_tool: Any = None
_check_teammate_tool: Any = None
_read_inbox_tool: Any = None


_SPAWN_TOOL_DESCRIPTION = (
    "Spawn a new teammate in a new acpx session. "
    "Names must be unique within the team. "
    "cwd must be an absolute path to the teammate's working directory. "
    "Available backends: 'claude' (default, models: sonnet, opus, haiku). "
    "subagent_type selects the agent runtime: 'kiro' (default), 'claude', 'codex', 'gemini', 'opencode', 'pi'. "
    "IMPORTANT: spawn only initiates the session — it does NOT guarantee the agent is ready. "
    "After spawning, poll check_teammate every 5-10 seconds until agent_status is not "
    "'no-session' before assigning tasks or sending messages. "
    "If agent_status remains 'no-session' after ~60 seconds, consider the spawn failed."
)

_CHECK_TEAMMATE_BASE_DESCRIPTION = (
    "Check a single teammate's status: alive/dead, unread messages from them, "
    "their unread count. Always non-blocking. "
    "Use parallel calls to check multiple teammates. "
    "Push notifications are NOT available in this session "
    "(not supported by the current harness). "
    "Do NOT pass notify_after_minutes. "
    "NOTE: agent_status reflects the acpx session state. Any status other than "
    "'no-session' means the session exists and the agent can receive messages. "
    "'dead' means idle (no active task) — normal after spawn or task completion. "
    "'running'/'alive' means actively processing. "
    "Only 'no-session' means the session was never created or has been closed."
)


def _build_check_teammate_description(push_available: bool) -> str:
    if push_available:
        return (
            "Check a single teammate's status: alive/dead, unread messages from them, "
            "their unread count. Always non-blocking. "
            "Use parallel calls to check multiple teammates. "
            "Push notifications are available in this session. "
            "Use notify_after_minutes to schedule a deferred reminder. "
            "NOTE: agent_status reflects the acpx session state. Any status other than "
            "'no-session' means the session exists and the agent can receive messages. "
            "'dead' means idle (no active task) — normal after spawn or task completion. "
            "'running'/'alive' means actively processing. "
            "Only 'no-session' means the session was never created or has been closed."
        )
    return _CHECK_TEAMMATE_BASE_DESCRIPTION


_READ_INBOX_BASE_DESCRIPTION = (
    "Read messages from an agent's inbox. Returns unread messages by default "
    "and marks them as read."
)


def _build_read_inbox_description(is_lead_session: bool) -> str:
    if is_lead_session:
        return (
            _READ_INBOX_BASE_DESCRIPTION
            + " NOTE: As team-lead, prefer check_teammate to read messages"
            " from a specific teammate. check_teammate filters by sender"
            " and provides richer status."
        )
    return _READ_INBOX_BASE_DESCRIPTION


@lifespan
async def app_lifespan(server):
    global _spawn_tool, _check_teammate_tool, _read_inbox_tool

    # 找 acpx：优先从 nvm 各版本里找，确保 node 和 acpx 在同一目录
    acpx_binary = None
    _nvm_dir = os.path.expanduser("~/.nvm/versions/node")
    if os.path.isdir(_nvm_dir):
        for _v in sorted(os.listdir(_nvm_dir), reverse=True):
            _bin = os.path.join(_nvm_dir, _v, "bin")
            _acpx = os.path.join(_bin, "acpx")
            _node = os.path.join(_bin, "node")
            if os.path.isfile(_acpx) and os.path.isfile(_node):
                acpx_binary = _acpx
                # 把对应的 node bin 目录插到 PATH 最前面，确保 subprocess 用同一个 node
                _current_path = os.environ.get("PATH", "")
                if _bin not in _current_path:
                    os.environ["PATH"] = _bin + ":" + _current_path
                break
    if not acpx_binary:
        acpx_binary = shutil.which("acpx")
    if not acpx_binary:
        raise FileNotFoundError(
            "acpx binary not found. Install acpx via: npm install -g acpx"
        )
    logger.warning("[claude-teams] acpx_binary=%s", acpx_binary)

    tools = await mcp.list_tools()
    _check_teammate_tool = next((t for t in tools if t.name == "check_teammate"), None)
    _read_inbox_tool = next((t for t in tools if t.name == "read_inbox"), None)

    session_id = str(uuid.uuid4())
    _lifespan_state.clear()
    _lifespan_state.update(
        {
            "acpx_binary": acpx_binary,
            "session_id": session_id,
            "client_name": "unknown",
            "client_version": "unknown",
        }
    )
    yield _lifespan_state


class HarnessDetectionMiddleware(Middleware):
    async def on_initialize(self, context, call_next):
        _unknown = SimpleNamespace(name="unknown", version="unknown")
        client_info = context.message.params.clientInfo or _unknown
        client_name = client_info.name
        client_version = client_info.version

        result = await call_next(context)

        logger.info("MCP client connected: %s v%s", client_name, client_version)

        _lifespan_state["client_name"] = client_name
        _lifespan_state["client_version"] = client_version

        return result


mcp = FastMCP(
    name="claude-teams",
    instructions=(
        "MCP server for orchestrating Claude Code agent teams. "
        "Manages team creation, teammate spawning, messaging, and task tracking."
    ),
    lifespan=app_lifespan,
)
mcp.add_middleware(HarnessDetectionMiddleware())


def _get_lifespan(ctx: Context) -> dict[str, Any]:
    return ctx.lifespan_context


def _content_metadata(content: str, sender: str) -> str:
    """Append sender signature and reply reminder to outgoing message content."""
    return (
        f"{content}\n\n"
        f"<system_reminder>"
        f"This message was sent from {sender}. "
        f"Use your send_message tool to respond."
        f"</system_reminder>"
    )


@mcp.tool
def team_create(
    team_name: str,
    ctx: Context,
    description: str = "",
) -> dict:
    """Create a new agent team. Sets up team config and task directories under ~/.claude/.
    One team per server session. Team names must be filesystem-safe
    (letters, numbers, hyphens, underscores)."""
    ls = _get_lifespan(ctx)
    result = teams.create_team(
        name=team_name, session_id=ls["session_id"], description=description
    )
    return result.model_dump()


@mcp.tool
def team_delete(team_name: str) -> dict:
    """Delete a team and all its data. Fails if any teammates are still active.
    Removes both team config and task directories."""
    try:
        result = teams.delete_team(team_name)
    except (RuntimeError, FileNotFoundError) as e:
        raise ToolError(str(e))
    return result.model_dump()


@mcp.tool(name="spawn_teammate", description=_SPAWN_TOOL_DESCRIPTION)
async def spawn_teammate_tool(
    team_name: str,
    name: str,
    ctx: Context,
    cwd: str = "",
    model: str = "sonnet",
    subagent_type: str = "kiro",
    plan_mode_required: bool = False,
) -> dict:
    """Spawn a new teammate using acpx. Description is dynamically updated at startup.
    cwd: working directory for the agent. Defaults to the current project directory if not provided. Must be an absolute path to an existing directory."""
    import os.path

    if not cwd:
        cwd = _get_lifespan(ctx).get("client_cwd", os.getcwd())
    if not os.path.isabs(cwd):
        raise ToolError("cwd must be an absolute path.")
    if not os.path.isdir(cwd):
        raise ToolError(f"cwd does not exist: {cwd!r}. Please provide a valid directory.")
    ls = _get_lifespan(ctx)
    try:
        member = spawn_teammate(
            team_name=team_name,
            name=name,
            acpx_binary=ls.get("acpx_binary", "acpx"),
            lead_session_id=ls["session_id"],
            model=model,
            subagent_type=subagent_type,
            plan_mode_required=plan_mode_required,
            cwd=cwd,
        )
    except ValueError as e:
        raise ToolError(str(e))
    return SpawnResult(
        agent_id=member.agent_id,
        name=member.name,
        team_name=team_name,
    ).model_dump()


def _find_teammate(team_name: str, name: str) -> TeammateMember | None:
    config = teams.read_config(team_name)
    for m in config.members:
        if isinstance(m, TeammateMember) and m.name == name:
            return m
    return None


@mcp.tool
def send_message(
    team_name: str,
    type: Literal[
        "message",
        "broadcast",
        "shutdown_request",
        "shutdown_response",
        "plan_approval_response",
    ],
    ctx: Context,
    recipient: str = "",
    content: str = "",
    summary: str = "",
    request_id: str = "",
    approve: bool | None = None,
    sender: str = "team-lead",
) -> dict:
    """Send a message to a teammate or respond to a protocol request.
    Type 'message' sends a direct message (requires recipient, content, summary). content is the message body the recipient reads; summary is a short label for team-lead's preview.
    Type 'broadcast' sends to all teammates (requires summary).
    Type 'shutdown_request' asks a teammate to shut down (requires recipient; content used as reason).
    Type 'shutdown_response' responds to a shutdown request (requires sender, request_id, approve).
    Type 'plan_approval_response' responds to a plan approval request (requires recipient, request_id, approve)."""

    try:
        teams.read_config(team_name)
    except FileNotFoundError:
        raise ToolError(f"Team {team_name!r} not found")

    if type == "message":
        if not summary:
            raise ToolError("Message summary must not be empty")
        if not recipient:
            raise ToolError("Message recipient must not be empty")
        config = teams.read_config(team_name)
        member_names = {m.name for m in config.members}
        if sender not in member_names:
            raise ToolError(f"Sender {sender!r} is not a member of team {team_name!r}")
        if recipient not in member_names:
            raise ToolError(
                f"Recipient {recipient!r} is not a member of team {team_name!r}"
            )
        if sender == recipient:
            raise ToolError("Cannot send a message to yourself")
        if sender != "team-lead" and recipient != "team-lead":
            raise ToolError("Teammates can only send direct messages to team-lead")
        target_color = None
        for m in config.members:
            if m.name == recipient and isinstance(m, TeammateMember):
                target_color = m.color
                break
        content = _content_metadata(content, sender)
        messaging.send_plain_message(
            team_name,
            sender,
            recipient,
            content,
            summary=summary,
            color=target_color,
        )
        # 如果收件人是队员，自动触发 acpx prompt 让它去读 inbox
        recipient_member = None
        for m in config.members:
            if m.name == recipient and isinstance(m, TeammateMember):
                recipient_member = m
                break
        if recipient_member:
            ls = _get_lifespan(ctx)
            acpx_bin = ls.get("acpx_binary", "acpx")
            try:
                acpx_client.send_prompt(
                    recipient_member.agent_type,
                    recipient_member.acpx_session_name,
                    "You have a new message. Please read your inbox and respond.",
                    name=recipient_member.name,
                    team_name=team_name,
                    cwd=recipient_member.cwd,
                    acpx_binary=acpx_bin,
                )
            except Exception as e:
                logger.warning("Failed to notify %r via acpx: %s", recipient, e)
        return SendMessageResult(
            success=True,
            message=f"Message sent to {recipient}",
            routing={
                "sender": sender,
                "target": recipient,
                "targetColor": target_color,
            },
        ).model_dump(exclude_none=True)

    elif type == "broadcast":
        if sender != "team-lead":
            raise ToolError("Only team-lead can send broadcasts")
        if not summary:
            raise ToolError("Broadcast summary must not be empty")
        config = teams.read_config(team_name)
        content = _content_metadata(content, sender)
        count = 0
        for m in config.members:
            if isinstance(m, TeammateMember):
                messaging.send_plain_message(
                    team_name,
                    "team-lead",
                    m.name,
                    content,
                    summary=summary,
                    color=None,
                )
                count += 1
        return SendMessageResult(
            success=True,
            message=f"Broadcast sent to {count} teammate(s)",
        ).model_dump(exclude_none=True)

    elif type == "shutdown_request":
        if not recipient:
            raise ToolError("Shutdown request recipient must not be empty")
        if recipient == "team-lead":
            raise ToolError("Cannot send shutdown request to team-lead")
        config = teams.read_config(team_name)
        member_names = {m.name for m in config.members}
        if recipient not in member_names:
            raise ToolError(
                f"Recipient {recipient!r} is not a member of team {team_name!r}"
            )
        req_id = messaging.send_shutdown_request(team_name, recipient, reason=content)
        # 触发队员去读 inbox 处理 shutdown_request
        for m in config.members:
            if m.name == recipient and isinstance(m, TeammateMember):
                ls = _get_lifespan(ctx)
                acpx_bin = ls.get("acpx_binary", "acpx")
                try:
                    acpx_client.send_prompt(
                        m.agent_type,
                        m.acpx_session_name,
                        "You have a new message. Please read your inbox and respond.",
                        name=m.name,
                        team_name=team_name,
                        cwd=m.cwd,
                        acpx_binary=acpx_bin,
                    )
                except Exception as e:
                    logger.warning("Failed to notify %r via acpx: %s", recipient, e)
                break
        return SendMessageResult(
            success=True,
            message=f"Shutdown request sent to {recipient}",
            request_id=req_id,
            target=recipient,
        ).model_dump(exclude_none=True)

    elif type == "shutdown_response":
        config = teams.read_config(team_name)
        member = None
        for m in config.members:
            if isinstance(m, TeammateMember) and m.name == sender:
                member = m
                break
        if member is None:
            raise ToolError(
                f"Sender {sender!r} is not a teammate in team {team_name!r}"
            )

        if approve:
            payload = ShutdownApproved(
                request_id=request_id,
                from_=sender,
                timestamp=messaging.now_iso(),
                acpx_session_name=member.acpx_session_name,
            )
            messaging.send_structured_message(team_name, sender, "team-lead", payload)
            return SendMessageResult(
                success=True,
                message=f"Shutdown approved for request {request_id}",
            ).model_dump(exclude_none=True)
        else:
            messaging.send_plain_message(
                team_name,
                sender,
                "team-lead",
                content or "Shutdown rejected",
                summary="shutdown_rejected",
            )
            return SendMessageResult(
                success=True,
                message=f"Shutdown rejected for request {request_id}",
            ).model_dump(exclude_none=True)

    elif type == "plan_approval_response":
        if not recipient:
            raise ToolError("Plan approval recipient must not be empty")
        config = teams.read_config(team_name)
        member_names = {m.name for m in config.members}
        if recipient not in member_names:
            raise ToolError(
                f"Recipient {recipient!r} is not a member of team {team_name!r}"
            )
        if approve:
            messaging.send_plain_message(
                team_name,
                sender,
                recipient,
                '{"type":"plan_approval","approved":true}',
                summary="plan_approved",
            )
        else:
            messaging.send_plain_message(
                team_name,
                sender,
                recipient,
                content or "Plan rejected",
                summary="plan_rejected",
            )
        return SendMessageResult(
            success=True,
            message=f"Plan {'approved' if approve else 'rejected'} for {recipient}",
        ).model_dump(exclude_none=True)

    raise ToolError(f"Unknown message type: {type}")


@mcp.tool
def task_create(
    team_name: str,
    subject: str,
    description: str,
    active_form: str = "",
    metadata: dict | None = None,
) -> dict:
    """Create a new task for the team. Tasks are auto-assigned incrementing IDs.
    Optional metadata dict is stored alongside the task."""
    try:
        task = tasks.create_task(team_name, subject, description, active_form, metadata)
    except ValueError as e:
        raise ToolError(str(e))
    return {"id": task.id, "status": task.status}


@mcp.tool
def task_update(
    team_name: str,
    task_id: str,
    status: Literal["pending", "in_progress", "completed", "deleted"] | None = None,
    owner: str | None = None,
    subject: str | None = None,
    description: str | None = None,
    active_form: str | None = None,
    add_blocks: list[str] | None = None,
    add_blocked_by: list[str] | None = None,
    metadata: dict | None = None,
) -> dict:
    """Update a task's fields. Setting owner auto-notifies the assignee via
    inbox. Setting status to 'deleted' removes the task file from disk.
    Metadata keys are merged into existing metadata (set a key to null to delete it)."""
    if owner is not None:
        try:
            config = teams.read_config(team_name)
        except FileNotFoundError:
            raise ToolError(f"Team {team_name!r} not found")
        member_names = {m.name for m in config.members}
        if owner not in member_names:
            raise ToolError(f"Owner {owner!r} is not a member of team {team_name!r}")
    try:
        task = tasks.update_task(
            team_name,
            task_id,
            status=status,
            owner=owner,
            subject=subject,
            description=description,
            active_form=active_form,
            add_blocks=add_blocks,
            add_blocked_by=add_blocked_by,
            metadata=metadata,
        )
    except FileNotFoundError:
        raise ToolError(f"Task {task_id!r} not found in team {team_name!r}")
    except ValueError as e:
        raise ToolError(str(e))
    if owner is not None and task.owner is not None and task.status != "deleted":
        messaging.send_task_assignment(team_name, task, assigned_by="team-lead")
    return {"id": task.id, "status": task.status}


@mcp.tool
def task_list(team_name: str) -> list[dict]:
    """List all tasks for a team with their current status and assignments."""
    try:
        result = tasks.list_tasks(team_name)
    except ValueError as e:
        raise ToolError(str(e))
    return [t.model_dump(by_alias=True, exclude_none=True) for t in result]


@mcp.tool
def task_get(team_name: str, task_id: str) -> dict:
    """Get full details of a specific task by ID."""
    try:
        task = tasks.get_task(team_name, task_id)
    except FileNotFoundError:
        raise ToolError(f"Task {task_id!r} not found in team {team_name!r}")
    return task.model_dump(by_alias=True, exclude_none=True)


@mcp.tool(description=_READ_INBOX_BASE_DESCRIPTION)
def read_inbox(
    team_name: str,
    agent_name: str,
    unread_only: bool = True,
    mark_as_read: bool = True,
) -> list[dict]:
    """Read inbox messages. Description is dynamically updated at startup."""
    try:
        config = teams.read_config(team_name)
    except FileNotFoundError:
        raise ToolError(f"Team {team_name!r} not found")
    member_names = {m.name for m in config.members}
    if agent_name not in member_names:
        raise ToolError(f"Agent {agent_name!r} is not a member of team {team_name!r}")
    msgs = messaging.read_inbox(
        team_name, agent_name, unread_only=unread_only, mark_as_read=mark_as_read
    )
    return [m.model_dump(by_alias=True, exclude_none=True) for m in msgs]


@mcp.tool
def read_config(team_name: str) -> dict:
    """Read the current team configuration including all members."""
    try:
        config = teams.read_config(team_name)
    except FileNotFoundError:
        raise ToolError(f"Team {team_name!r} not found")
    data = config.model_dump(by_alias=True)
    return data


@mcp.tool
def team_status(team_name: str, ctx: Context) -> dict:
    """Check the live acpx session status of all teammates in a team.
    Returns each member's alive state and agent_status (running/idle/no-session)."""
    try:
        config = teams.read_config(team_name)
    except FileNotFoundError:
        raise ToolError(f"Team {team_name!r} not found")

    ls = _get_lifespan(ctx)
    acpx_bin = ls.get("acpx_binary", "acpx")
    members_status = []
    for m in config.members:
        if not isinstance(m, TeammateMember):
            continue
        try:
            status = acpx_client.get_session_status(
                m.agent_type, m.acpx_session_name, m.cwd, acpx_binary=acpx_bin
            )
            alive = status in ("alive", "running")
        except AcpxError as e:
            status = str(e)
            alive = False
        members_status.append({
            "name": m.name,
            "alive": alive,
            "agent_status": status,
            "acpx_session": m.acpx_session_name,
        })
    return {"team": team_name, "members": members_status}


@mcp.tool
def force_kill_teammate(team_name: str, agent_name: str, ctx: Context) -> dict:
    """Forcibly kill a teammate's acpx session. Use when graceful shutdown via
    send_message(type='shutdown_request') is not possible or not responding.
    Closes the acpx session, removes member from config, and resets their tasks."""
    config = teams.read_config(team_name)
    member = None
    for m in config.members:
        if isinstance(m, TeammateMember) and m.name == agent_name:
            member = m
            break
    if member is None:
        raise ToolError(f"Teammate {agent_name!r} not found in team {team_name!r}")
    ls = _get_lifespan(ctx)
    acpx_bin = ls.get("acpx_binary", "acpx")
    try:
        acpx_client.close_session(
            member.agent_type, member.acpx_session_name, member.cwd, acpx_binary=acpx_bin
        )
    except AcpxError as e:
        logger.warning("Failed to close acpx session %r: %s", member.acpx_session_name, e)
    teams.remove_member(team_name, agent_name)
    tasks.reset_owner_tasks(team_name, agent_name)
    return {"success": True, "message": f"{agent_name} has been stopped."}


@mcp.tool
def process_shutdown_approved(team_name: str, agent_name: str, ctx: Context) -> dict:
    """Process a teammate's shutdown by removing them from config and resetting
    their tasks. Call this after confirming shutdown_approved in the lead inbox."""
    if agent_name == "team-lead":
        raise ToolError("Cannot process shutdown for team-lead")
    member = _find_teammate(team_name, agent_name)
    if member is None:
        raise ToolError(f"Teammate {agent_name!r} not found in team {team_name!r}")
    ls = _get_lifespan(ctx)
    acpx_bin = ls.get("acpx_binary", "acpx")
    try:
        acpx_client.close_session(
            member.agent_type, member.acpx_session_name, member.cwd, acpx_binary=acpx_bin
        )
    except AcpxError as e:
        logger.warning("Failed to close acpx session %r: %s", member.acpx_session_name, e)
    teams.remove_member(team_name, agent_name)
    tasks.reset_owner_tasks(team_name, agent_name)
    return {"success": True, "message": f"{agent_name} removed from team."}


@mcp.tool(description=_CHECK_TEAMMATE_BASE_DESCRIPTION)
async def check_teammate(
    team_name: str,
    agent_name: str,
    ctx: Context,
    include_messages: bool = True,
    max_messages: int = 5,
) -> dict:
    """Check a single teammate's status. Description is dynamically updated
    at startup with push notification availability."""
    max_messages = max(1, min(max_messages, 20))

    try:
        config = teams.read_config(team_name)
    except FileNotFoundError:
        raise ToolError(f"Team {team_name!r} not found")

    member = None
    for m in config.members:
        if isinstance(m, TeammateMember) and m.name == agent_name:
            member = m
            break
    if member is None:
        raise ToolError(f"Teammate {agent_name!r} not found in team {team_name!r}")

    # 1. Read lead's inbox for unread messages FROM this teammate
    pending_from: list[dict] = []
    if include_messages:
        msgs = messaging.read_inbox_filtered(
            team_name=team_name,
            agent_name="team-lead",
            sender_filter=agent_name,
            unread_only=True,
            mark_as_read=True,
            limit=max_messages,
        )
        pending_from = [m.model_dump(by_alias=True, exclude_none=True) for m in msgs]

    # 2. Check teammate's unread count (messages they haven't read)
    try:
        their_unread = messaging.read_inbox(
            team_name, agent_name, unread_only=True, mark_as_read=False
        )
        their_unread_count = len(their_unread)
    except (FileNotFoundError, json.JSONDecodeError):
        their_unread_count = 0

    # 3. alive status via acpx_client.get_session_status
    ls = _get_lifespan(ctx)
    acpx_bin = ls.get("acpx_binary", "acpx")
    alive = False
    agent_status = None
    error = None
    try:
        status = acpx_client.get_session_status(
            member.agent_type, member.acpx_session_name, member.cwd, acpx_binary=acpx_bin
        )
        alive = status in ("alive", "running")
        agent_status = status
        if not alive:
            error = f"session status: {status}"
    except AcpxError as e:
        alive = False
        error = str(e)

    result: dict = {
        "name": agent_name,
        "alive": alive,
        "agent_status": agent_status,
        "pending_from": pending_from,
        "their_unread_count": their_unread_count,
        "push_available": False,
        "notification_scheduled": False,
    }
    if error:
        result["error"] = error
    return result


def main():
    logging.basicConfig(
        level=logging.INFO, format="%(levelname)s %(name)s: %(message)s"
    )
    mcp.run()


if __name__ == "__main__":
    main()
