#!/usr/bin/env python3
"""FastMCP-based Claudy server for managing Claude agent sessions."""

import asyncio
import os
from contextlib import asynccontextmanager
from datetime import datetime
from pathlib import Path
from typing import Dict, Optional

from claude_agent_sdk import ClaudeSDKClient, ClaudeAgentOptions
from claude_agent_sdk.types import TextBlock
from fastmcp import FastMCP, Context

# Session management configuration
SESSION_IDLE_TIMEOUT = 7200  # 2 hours in seconds
SESSION_CLEANUP_INTERVAL = 300  # Check every 5 minutes


def get_current_claude_session_id() -> Optional[str]:
    """Get the current Claude Code session ID from the project's session files."""
    try:
        cwd = os.getcwd()
        project_name = cwd.replace('/', '-')
        sessions_dir = Path.home() / ".claude" / "projects" / project_name

        if sessions_dir.exists():
            jsonl_files = sorted(
                sessions_dir.glob("*.jsonl"),
                key=lambda x: x.stat().st_mtime,
                reverse=True
            )
            if jsonl_files:
                return jsonl_files[0].stem  # Filename without extension is the session ID
    except Exception:
        pass
    return None


# Global session storage (shared across all MCP connections)
# Key: session_name, Value: (client, metadata)
_global_sessions: Dict[str, tuple[ClaudeSDKClient, dict]] = {}


async def cleanup_idle_sessions():
    """Background task to cleanup idle sessions."""
    while True:
        try:
            await asyncio.sleep(SESSION_CLEANUP_INTERVAL)

            now = datetime.now()
            to_cleanup = []

            for name, (client, metadata) in _global_sessions.items():
                last_used = datetime.fromisoformat(metadata["last_used"])
                idle_seconds = (now - last_used).total_seconds()

                if idle_seconds > SESSION_IDLE_TIMEOUT:
                    to_cleanup.append(name)

            # Cleanup idle sessions
            for name in to_cleanup:
                if name in _global_sessions:
                    client, _ = _global_sessions[name]
                    try:
                        await client.disconnect()
                    except Exception:
                        pass  # Best effort cleanup
                    del _global_sessions[name]

        except Exception:
            pass  # Don't let background task crash


@asynccontextmanager
async def lifespan(server: FastMCP):
    """Lifespan context manager for background tasks."""
    # Start background cleanup task
    cleanup_task = asyncio.create_task(cleanup_idle_sessions())

    try:
        yield
    finally:
        # Cleanup on shutdown
        cleanup_task.cancel()
        try:
            await cleanup_task
        except asyncio.CancelledError:
            pass

        # Disconnect all sessions
        for client, _ in _global_sessions.values():
            try:
                await client.disconnect()
            except Exception:
                pass
        _global_sessions.clear()


# Initialize FastMCP server
mcp = FastMCP(
    name="claudy",
    instructions="Universal Claude agent manager for spawning and managing persistent Claude agent sessions.",
    lifespan=lifespan
)


async def get_or_create_session(
    name: str,
    parent_session_id: Optional[str] = None
) -> tuple[ClaudeSDKClient, dict]:
    """Get existing session or create new one with auto-spawn."""
    if name in _global_sessions:
        client, metadata = _global_sessions[name]
        # Update last_used
        metadata["last_used"] = datetime.now().isoformat()
        return client, metadata

    # Auto-spawn new session
    options = ClaudeAgentOptions()
    options.permission_mode = "bypassPermissions"

    if parent_session_id:
        options.resume = parent_session_id
        options.fork_session = True

    client = ClaudeSDKClient(options=options)
    await client.connect()

    metadata = {
        "created_at": datetime.now().isoformat(),
        "last_used": datetime.now().isoformat(),
        "message_count": 0,
        "session_id": None,
        "auto_created": True,
    }
    if parent_session_id:
        metadata["parent_session_id"] = parent_session_id

    _global_sessions[name] = (client, metadata)
    return client, metadata


@mcp.tool
async def claudy_call(
    name: str,
    message: str,
    verbosity: str = "normal",
    fork: bool = False,
    fork_name: Optional[str] = None,
    parent_session_id: Optional[str] = None,
) -> str:
    """Send a message to an agent session.

    If the session doesn't exist, it will be automatically created inheriting the current
    Claude session's settings. Context is preserved across calls. Optionally fork the
    session to explore alternative paths.

    Args:
        name: Session name to call (will be auto-created if doesn't exist)
        message: Message to send to the agent
        verbosity: Output verbosity level (quiet, normal, verbose)
        fork: If true, fork the session before sending the message (creates a new branch)
        fork_name: Optional name for forked session (auto-generated if not provided)
        parent_session_id: Optional parent session ID to inherit from when auto-creating
    """
    # Auto-detect current Claude session ID if not provided
    if not parent_session_id:
        parent_session_id = get_current_claude_session_id()

    # Handle forking
    if fork:
        # Get parent session
        if name not in _global_sessions:
            return {
                "success": False,
                "error": f"Cannot fork non-existent session '{name}'",
            }

        _, parent_metadata = _global_sessions[name]
        parent_session_id = parent_metadata.get("session_id")

        if not parent_session_id:
            return {
                "success": False,
                "error": f"Cannot fork session '{name}': no session_id available. Send at least one message first.",
            }

        # Determine fork session name
        if not fork_name:
            fork_name = f"{name}_fork_{datetime.now().strftime('%Y%m%d_%H%M%S')}"

        # Check if fork name already exists
        if fork_name in _global_sessions:
            return {
                "success": False,
                "error": f"Fork session name '{fork_name}' already exists",
            }

        # Use fork_name as the target session
        name = fork_name

    # Get or create session
    client, metadata = await get_or_create_session(name, parent_session_id)

    # Send message and collect response
    await client.query(message)

    response_parts = []
    all_events = []
    session_id = None

    async for msg in client.receive_response():
        # Extract session_id from first message if available
        if session_id is None and hasattr(msg, "session_id"):
            session_id = msg.session_id
            if metadata["session_id"] is None:
                metadata["session_id"] = session_id

        # Collect all events for verbose mode
        if verbosity == "verbose":
            msg_dict = {"type": type(msg).__name__}

            if hasattr(msg, "content"):
                msg_dict["content"] = []
                for block in msg.content:
                    block_dict = {"type": getattr(block, "type", type(block).__name__)}

                    if isinstance(block, TextBlock):
                        block_dict["text"] = block.text
                    elif hasattr(block, "thinking"):
                        block_dict["thinking"] = block.thinking
                    elif hasattr(block, "name") and hasattr(block, "input"):
                        block_dict["name"] = block.name
                        block_dict["input"] = block.input
                    elif hasattr(block, "tool_use_id"):
                        block_dict["tool_use_id"] = block.tool_use_id
                        if hasattr(block, "content"):
                            block_dict["content"] = block.content

                    msg_dict["content"].append(block_dict)

            if hasattr(msg, "stop_reason"):
                msg_dict["stop_reason"] = msg.stop_reason
            if hasattr(msg, "role"):
                msg_dict["role"] = msg.role

            all_events.append(msg_dict)

        # Collect text for response
        if hasattr(msg, "content"):
            for block in msg.content:
                if isinstance(block, TextBlock):
                    response_parts.append(block.text)

    # Update metadata
    metadata["last_used"] = datetime.now().isoformat()
    metadata["message_count"] += 1

    # Build response based on verbosity
    result = {
        "success": True,
        "name": name,
        "response": "\n".join(response_parts),
    }

    if session_id:
        result["session_id"] = session_id

    if fork:
        result["forked"] = True

    if verbosity == "verbose":
        result["events"] = all_events

    import json
    return json.dumps(result, indent=2, ensure_ascii=False)


@mcp.tool
async def claudy_list() -> str:
    """List all active agent sessions with their metadata."""
    sessions_list = [
        {"name": name, **metadata}
        for name, (_, metadata) in _global_sessions.items()
    ]

    import json
    return json.dumps(
        {"success": True, "sessions": sessions_list},
        indent=2,
        ensure_ascii=False
    )


@mcp.tool
async def claudy_status(name: str) -> str:
    """Get detailed status of a specific agent session.

    Args:
        name: Session name to check status
    """
    if name not in _global_sessions:
        import json
        return json.dumps(
            {
                "success": False,
                "error": f"Session '{name}' not found",
                "available_sessions": list(_global_sessions.keys()),
            },
            indent=2,
            ensure_ascii=False
        )

    _, metadata = _global_sessions[name]

    import json
    return json.dumps(
        {"success": True, "name": name, **metadata},
        indent=2,
        ensure_ascii=False
    )


@mcp.tool
async def claudy_cleanup(name: Optional[str] = None, all: bool = False) -> str:
    """Cleanup one or all agent sessions. Use this when done with agents.

    Args:
        name: Session name to cleanup (optional if all=true)
        all: If true, cleanup all sessions
    """
    if all:
        # Cleanup all sessions
        count = len(_global_sessions)
        for client, _ in list(_global_sessions.values()):
            try:
                await client.disconnect()
            except Exception:
                pass
        _global_sessions.clear()

        import json
        return json.dumps(
            {"success": True, "message": f"Cleaned up {count} session(s)"},
            indent=2,
            ensure_ascii=False
        )
    else:
        if not name:
            import json
            return json.dumps(
                {"success": False, "error": "Session name is required"},
                indent=2,
                ensure_ascii=False
            )

        if name not in _global_sessions:
            import json
            return json.dumps(
                {
                    "success": False,
                    "error": f"Session '{name}' not found",
                    "available_sessions": list(_global_sessions.keys()),
                },
                indent=2,
                ensure_ascii=False
            )

        # Cleanup single session
        client, _ = _global_sessions[name]
        try:
            await client.disconnect()
        except Exception:
            pass

        del _global_sessions[name]

        import json
        return json.dumps(
            {
                "success": True,
                "name": name,
                "message": f"Session '{name}' cleaned up successfully",
            },
            indent=2,
            ensure_ascii=False
        )


if __name__ == "__main__":
    # Run with FastMCP CLI
    mcp.run()
