"""
PenBot MCP Server — expose PenBot tools to AI agents via the
Model Context Protocol.

Usage:
    penbot mcp                     # Start MCP server (stdio transport)
    penbot mcp --transport http    # Start MCP server (HTTP/SSE transport)
    penbot mcp --port 9100         # Custom port for HTTP transport

The server exposes the following tools:
    - penbot_doctor       — environment health check
    - penbot_agents       — list available agents
    - penbot_agent_detail — describe a specific agent
    - penbot_init         — create a target config
    - penbot_validate     — validate config + connectivity
    - penbot_test         — run a security test campaign
    - penbot_sessions     — list past sessions
    - penbot_session_view — view a session's data
    - penbot_patterns     — list attack pattern libraries
    - penbot_pattern_search — search patterns by keyword
"""

import asyncio
import json
import re
from pathlib import Path

from mcp.server.fastmcp import FastMCP

mcp = FastMCP(
    name="penbot",
    instructions=(
        "PenBot is an AI chatbot security testing framework. "
        "Use these tools to run adversarial tests, check environment health, "
        "manage sessions, and browse attack patterns. "
        "Start with penbot_doctor to verify the environment, then penbot_agents "
        "to see available agents, and penbot_test to run a campaign."
    ),
)


# ---------------------------------------------------------------------------
# Tool: doctor
# ---------------------------------------------------------------------------
@mcp.tool()
def penbot_doctor() -> dict:
    """Check PenBot environment health (Python, packages, API keys, browsers).

    Returns a dict with 'checks' array, 'summary' counts, and 'healthy' bool.
    """
    from src.cli.doctor import (
        FAIL,
        PASS,
        WARN,
        _check_api_keys,
        _check_config_files,
        _check_cookie_dir,
        _check_core_packages,
        _check_env_file,
        _check_playwright_browsers,
        _check_python_version,
        _check_sessions_db,
    )

    results = []
    results.append(_check_python_version())
    results.extend(_check_core_packages())
    results.extend(_check_api_keys())
    results.append(_check_playwright_browsers())
    results.append(_check_env_file())
    results.append(_check_config_files())
    results.append(_check_cookie_dir())
    results.append(_check_sessions_db())

    status_map = {PASS: "pass", WARN: "warn", FAIL: "fail"}
    checks = [
        {"check": name, "status": status_map.get(status, "unknown"), "detail": detail}
        for name, status, detail in results
    ]
    passes = sum(1 for c in checks if c["status"] == "pass")
    warns = sum(1 for c in checks if c["status"] == "warn")
    fails = sum(1 for c in checks if c["status"] == "fail")

    return {
        "checks": checks,
        "summary": {"pass": passes, "warn": warns, "fail": fails},
        "healthy": fails == 0,
    }


# ---------------------------------------------------------------------------
# Tool: agents
# ---------------------------------------------------------------------------
@mcp.tool()
def penbot_agents() -> dict:
    """List all available security testing agents.

    Returns a dict with 'agents' array and 'total' count.
    """
    from src.cli.agents_cmd import AGENT_REGISTRY

    agents = [
        {
            "id": agent_id,
            "class": info["class"],
            "description": info["description"],
            "owasp": info["owasp"],
            "patterns": info["patterns"],
            "priority": info["priority"],
            "category": info["category"],
        }
        for agent_id, info in AGENT_REGISTRY.items()
    ]
    return {"agents": agents, "total": len(agents)}


# ---------------------------------------------------------------------------
# Tool: agent detail
# ---------------------------------------------------------------------------
@mcp.tool()
def penbot_agent_detail(agent_id: str) -> dict:
    """Show detailed information about a specific agent.

    Args:
        agent_id: One of: jailbreak, encoding, info_disclosure,
                  output_security, impersonation, compliance,
                  token_soup, evolutionary, rag_poisoning, tool_exploit
    """
    from src.cli.agents_cmd import AGENT_REGISTRY

    agent_id = agent_id.lower()
    if agent_id not in AGENT_REGISTRY:
        return {
            "error": "unknown_agent",
            "message": f"Unknown agent: {agent_id}",
            "available": list(AGENT_REGISTRY.keys()),
        }
    return {"id": agent_id, **AGENT_REGISTRY[agent_id]}


# ---------------------------------------------------------------------------
# Tool: init
# ---------------------------------------------------------------------------
@mcp.tool()
def penbot_init(
    name: str,
    connector_type: str,
    endpoint: str,
    agents: str = "",
    max_attacks: int = 30,
) -> dict:
    """Create a target YAML config file (non-interactive).

    Args:
        name: Target display name (e.g. "Staging Bot")
        connector_type: "rest" or "playwright"
        endpoint: API URL or web page URL
        agents: Comma-separated agent IDs (optional)
        max_attacks: Maximum attack count (default 30)
    """
    import yaml

    from src.cli.init_cmd import _build_config

    agent_list = [a.strip() for a in agents.split(",") if a.strip()] if agents else []
    config = _build_config(name, connector_type, endpoint, agent_list, max_attacks)

    slug = re.sub(r"[^a-z0-9]+", "_", name.lower()).strip("_")
    configs_dir = Path("configs/clients")
    configs_dir.mkdir(parents=True, exist_ok=True)
    config_path = configs_dir / f"{slug}.yaml"

    config_path.write_text(
        yaml.dump(config, default_flow_style=False, sort_keys=False),
        encoding="utf-8",
    )

    return {
        "config_path": str(config_path),
        "target_name": name,
        "connector_type": connector_type,
    }


# ---------------------------------------------------------------------------
# Tool: validate
# ---------------------------------------------------------------------------
@mcp.tool()
async def penbot_validate(config_path: str) -> dict:
    """Validate a target YAML config and test connectivity.

    Args:
        config_path: Path to the YAML config file
    """
    from src.cli.config_loader import load_config
    from src.connectors.factory import create_connector

    try:
        config = load_config(config_path)
    except Exception as e:
        return {"config": config_path, "valid": False, "error": str(e)}

    target_config = config.get("target", {})
    conn_type_raw = target_config.get("connection", {}).get("type", "rest")
    conn_type = "playwright" if conn_type_raw == "playwright" else "api"

    try:
        connector = create_connector(conn_type, target_config)
        if hasattr(connector, "initialize"):
            await connector.initialize()
        reachable = await connector.health_check()
        if hasattr(connector, "close"):
            await connector.close()
    except Exception as e:
        return {
            "config": config_path,
            "valid": True,
            "connector_type": conn_type_raw,
            "reachable": False,
            "error": str(e),
        }

    return {
        "config": config_path,
        "valid": True,
        "connector_type": conn_type_raw,
        "reachable": reachable,
    }


# ---------------------------------------------------------------------------
# Tool: sessions
# ---------------------------------------------------------------------------
@mcp.tool()
def penbot_sessions() -> dict:
    """List all past test sessions.

    Returns a dict with 'sessions' array and 'total' count.
    """
    sessions_dir = Path("sessions")
    if not sessions_dir.exists():
        return {"sessions": [], "total": 0}

    files = sorted(sessions_dir.glob("*.json"), key=lambda f: f.stat().st_mtime, reverse=True)
    sessions = []
    for f in files:
        try:
            data = json.loads(f.read_text(encoding="utf-8"))
            sessions.append(
                {
                    "session_id": f.stem,
                    "target_name": data.get("target_name", "Unknown"),
                    "attacks": len(data.get("attack_attempts", [])),
                    "findings": len(data.get("security_findings", [])),
                    "status": data.get("test_status", "unknown"),
                    "started_at": data.get("started_at", ""),
                }
            )
        except Exception:
            sessions.append({"session_id": f.stem, "error": "failed to load"})

    return {"sessions": sessions, "total": len(sessions)}


# ---------------------------------------------------------------------------
# Tool: session view
# ---------------------------------------------------------------------------
@mcp.tool()
def penbot_session_view(session_id: str) -> dict:
    """View full data for a specific session.

    Args:
        session_id: Session UUID or numeric index (1 = most recent)
    """
    sessions_dir = Path("sessions")
    if not sessions_dir.exists():
        return {"error": "no_sessions", "message": "No sessions directory found"}

    if session_id.isdigit():
        files = sorted(sessions_dir.glob("*.json"), key=lambda f: f.stat().st_mtime, reverse=True)
        idx = int(session_id) - 1
        if 0 <= idx < len(files):
            session_id = files[idx].stem
        else:
            return {"error": "not_found", "message": f"Session #{session_id} not found"}

    file_path = sessions_dir / f"{session_id}.json"
    if not file_path.exists():
        return {"error": "not_found", "message": f"Session not found: {session_id}"}

    return json.loads(file_path.read_text(encoding="utf-8"))


# ---------------------------------------------------------------------------
# Tool: patterns
# ---------------------------------------------------------------------------
@mcp.tool()
def penbot_patterns() -> dict:
    """List all attack pattern libraries.

    Returns a dict with 'libraries' array and 'total_patterns' count.
    """
    lib_dir = Path("src/attack_library")
    if not lib_dir.exists():
        return {"libraries": [], "total_patterns": 0}

    libraries = []
    total = 0
    for f in sorted(lib_dir.glob("*.json")):
        try:
            data = json.loads(f.read_text(encoding="utf-8"))
            count = len(data) if isinstance(data, list) else 1
            cats = (
                list({p.get("category", "unknown") for p in data})
                if isinstance(data, list)
                else [data.get("category", "unknown")]
            )
            total += count
            libraries.append({"name": f.stem, "patterns": count, "categories": cats})
        except Exception as e:
            libraries.append({"name": f.stem, "error": str(e)})

    return {"libraries": libraries, "total_patterns": total}


# ---------------------------------------------------------------------------
# Tool: pattern search
# ---------------------------------------------------------------------------
@mcp.tool()
def penbot_pattern_search(query: str, limit: int = 20) -> dict:
    """Search attack patterns by keyword.

    Args:
        query: Search keyword
        limit: Maximum results (default 20)
    """
    lib_dir = Path("src/attack_library")
    query_lower = query.lower()
    results = []

    if not lib_dir.exists():
        return {"query": query, "results": [], "total": 0}

    for f in lib_dir.glob("*.json"):
        try:
            data = json.loads(f.read_text(encoding="utf-8"))
            patterns = data if isinstance(data, list) else [data]
            for p in patterns:
                name = p.get("name", "")
                desc = p.get("description", "")
                cat = p.get("category", "")
                if (
                    query_lower in name.lower()
                    or query_lower in desc.lower()
                    or query_lower in cat.lower()
                ):
                    results.append(
                        {
                            "library": f.stem,
                            "name": name,
                            "description": desc,
                            "category": cat,
                            "priority": p.get("priority", 3),
                            "severity": p.get("severity_if_success", "unknown"),
                        }
                    )
        except Exception:
            continue

    results.sort(key=lambda x: x["priority"], reverse=True)
    results = results[:limit]
    return {"query": query, "results": results, "total": len(results)}


# ---------------------------------------------------------------------------
# Tool: test
# ---------------------------------------------------------------------------
@mcp.tool()
async def penbot_test(
    config_path: str,
    agents: str = "",
    max_attacks: int = 0,
    quick: bool = False,
) -> dict:
    """Run an adversarial security test campaign against a target chatbot.

    This is a long-running operation. Returns the final session summary
    once complete.

    Args:
        config_path: Path to the YAML target config
        agents: Comma-separated agent IDs (optional, uses config default)
        max_attacks: Override max attack count (0 = use config default)
        quick: Quick mode — only 3 attacks
    """
    import sys

    cmd = [sys.executable, "-m", "src.cli.main", "test", "--config", config_path, "--machine"]

    if agents:
        cmd += ["--agents", agents]
    if max_attacks > 0:
        cmd += ["--max-attacks", str(max_attacks)]
    if quick:
        cmd.append("--quick")

    proc = await asyncio.create_subprocess_exec(
        *cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )

    events = []
    assert proc.stdout is not None
    async for raw_line in proc.stdout:
        line = raw_line.decode("utf-8", errors="replace").strip()
        if line:
            try:
                events.append(json.loads(line))
            except json.JSONDecodeError:
                events.append({"raw": line})

    await proc.wait()

    if events:
        last = events[-1]
        return {
            "status": "completed" if proc.returncode == 0 else "failed",
            "exit_code": proc.returncode,
            "events_count": len(events),
            "final_event": last,
        }

    stderr_out = (await proc.stderr.read()).decode("utf-8", errors="replace") if proc.stderr else ""
    return {
        "status": "failed",
        "exit_code": proc.returncode,
        "error": stderr_out[:500],
    }


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------
def run_mcp_server(transport: str = "stdio", host: str = "127.0.0.1", port: int = 9100):
    """Start the MCP server with the specified transport."""
    if transport == "stdio":
        mcp.run(transport="stdio")
    else:
        mcp.run(transport="sse", host=host, port=port)
