"""MCP Integration Test Generator.

Generates pytest-compatible integration tests for MCP servers.
Analyzes server source code to detect tools, resources, and prompts,
then generates test stubs covering: invocation, error handling, and permissions.

Zero external dependencies. Python 3.11+ stdlib only.
"""

from __future__ import annotations

import ast
import json
import re
import textwrap
from dataclasses import dataclass, field
from typing import Any


# ---------------------------------------------------------------------------
# Data Models
# ---------------------------------------------------------------------------

@dataclass
class DetectedTool:
    """A tool detected in MCP server code."""
    name: str
    line: int
    parameters: list[str] = field(default_factory=list)
    description: str = ""
    has_async: bool = False

    def to_dict(self) -> dict[str, Any]:
        return {
            "name": self.name, "line": self.line,
            "parameters": self.parameters,
            "description": self.description,
            "has_async": self.has_async,
        }


@dataclass
class DetectedResource:
    """A resource detected in MCP server code."""
    name: str
    uri_pattern: str = ""
    line: int = 0
    description: str = ""

    def to_dict(self) -> dict[str, Any]:
        return {
            "name": self.name, "uri_pattern": self.uri_pattern,
            "line": self.line, "description": self.description,
        }


@dataclass
class DetectedPrompt:
    """A prompt detected in MCP server code."""
    name: str
    line: int = 0
    arguments: list[str] = field(default_factory=list)
    description: str = ""

    def to_dict(self) -> dict[str, Any]:
        return {
            "name": self.name, "line": self.line,
            "arguments": self.arguments,
            "description": self.description,
        }


@dataclass
class ServerProfile:
    """Profile of detected MCP server components."""
    tools: list[DetectedTool] = field(default_factory=list)
    resources: list[DetectedResource] = field(default_factory=list)
    prompts: list[DetectedPrompt] = field(default_factory=list)
    framework: str = "unknown"
    has_async: bool = False
    server_name: str = ""

    @property
    def total_components(self) -> int:
        return len(self.tools) + len(self.resources) + len(self.prompts)

    def to_dict(self) -> dict[str, Any]:
        return {
            "tools": [t.to_dict() for t in self.tools],
            "resources": [r.to_dict() for r in self.resources],
            "prompts": [p.to_dict() for p in self.prompts],
            "framework": self.framework,
            "has_async": self.has_async,
            "server_name": self.server_name,
            "total_components": self.total_components,
        }


@dataclass
class GeneratedTest:
    """A single generated test case."""
    name: str
    category: str   # "invocation", "error_handling", "permissions", "edge_case"
    component: str  # tool/resource/prompt name
    code: str

    def to_dict(self) -> dict[str, Any]:
        return {
            "name": self.name, "category": self.category,
            "component": self.component, "code": self.code,
        }


@dataclass
class GeneratedSuite:
    """Collection of generated tests with metadata."""
    tests: list[GeneratedTest] = field(default_factory=list)
    profile: ServerProfile = field(default_factory=ServerProfile)
    total_tests: int = 0

    def to_dict(self) -> dict[str, Any]:
        return {
            "total_tests": self.total_tests,
            "profile": self.profile.to_dict(),
            "tests": [t.to_dict() for t in self.tests],
        }


# ---------------------------------------------------------------------------
# Detection: find tools, resources, prompts in source code
# ---------------------------------------------------------------------------

def _detect_framework(code: str) -> str:
    """Detect MCP framework from source code."""
    if "from fastmcp" in code or "from mcp.server.fastmcp" in code:
        return "fastmcp"
    if "from mcp.server" in code or "import mcp" in code:
        return "mcp-sdk"
    if "@server.tool" in code or "@app.tool" in code:
        return "fastmcp"
    if "McpServer" in code or "MCP.Server" in code:
        return "mcp-sdk"
    return "unknown"


def _detect_tools_ast(code: str) -> list[DetectedTool]:
    """Detect tools via AST analysis."""
    tools: list[DetectedTool] = []
    try:
        tree = ast.parse(code)
    except SyntaxError:
        return _detect_tools_regex(code)

    for node in ast.walk(tree):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            # Check decorators for @tool, @server.tool, etc.
            for dec in node.decorator_list:
                dec_name = _get_decorator_name(dec)
                if dec_name and "tool" in dec_name.lower():
                    params = [
                        a.arg for a in node.args.args
                        if a.arg != "self" and a.arg != "ctx"
                    ]
                    desc = ast.get_docstring(node) or ""
                    tools.append(DetectedTool(
                        name=node.name, line=node.lineno,
                        parameters=params, description=desc[:200],
                        has_async=isinstance(node, ast.AsyncFunctionDef),
                    ))
                    break

    # Also detect tools registered via add_tool or register_tool patterns
    for node in ast.walk(tree):
        if isinstance(node, ast.Call):
            func_name = _get_call_name(node)
            if func_name and ("add_tool" in func_name or "register_tool" in func_name):
                if node.args:
                    arg = node.args[0]
                    if isinstance(arg, ast.Name):
                        name = arg.id
                        if not any(t.name == name for t in tools):
                            tools.append(DetectedTool(name=name, line=node.lineno))

    return tools


def _detect_tools_regex(code: str) -> list[DetectedTool]:
    """Fallback: detect tools via regex patterns."""
    tools: list[DetectedTool] = []
    patterns = [
        re.compile(r'@\w*\.?tool\s*(?:\([^)]*\))?\s*\n\s*(?:async\s+)?def\s+(\w+)\s*\(([^)]*)\)', re.MULTILINE),
        re.compile(r'add_tool\(\s*(\w+)', re.MULTILINE),
    ]
    for pat in patterns:
        for m in pat.finditer(code):
            name = m.group(1)
            params = []
            if m.lastindex and m.lastindex >= 2:
                param_str = m.group(2)
                params = [p.strip().split(":")[0].strip()
                          for p in param_str.split(",")
                          if p.strip() and p.strip() not in ("self", "ctx")]
            if not any(t.name == name for t in tools):
                line = code[:m.start()].count("\n") + 1
                tools.append(DetectedTool(name=name, line=line, parameters=params))
    return tools


def _detect_resources(code: str) -> list[DetectedResource]:
    """Detect resources in MCP server code."""
    resources: list[DetectedResource] = []

    # Decorator pattern: @resource, @server.resource
    pat = re.compile(
        r'@\w*\.?resource\s*\(\s*["\']([^"\']+)["\']\s*\)\s*\n\s*(?:async\s+)?def\s+(\w+)',
        re.MULTILINE,
    )
    for m in pat.finditer(code):
        uri = m.group(1)
        name = m.group(2)
        line = code[:m.start()].count("\n") + 1
        resources.append(DetectedResource(name=name, uri_pattern=uri, line=line))

    # add_resource pattern
    pat2 = re.compile(r'add_resource\(\s*["\']([^"\']+)["\']', re.MULTILINE)
    for m in pat2.finditer(code):
        uri = m.group(1)
        line = code[:m.start()].count("\n") + 1
        if not any(r.uri_pattern == uri for r in resources):
            resources.append(DetectedResource(
                name=uri.split("/")[-1] or "resource",
                uri_pattern=uri, line=line,
            ))

    return resources


def _detect_prompts(code: str) -> list[DetectedPrompt]:
    """Detect prompts in MCP server code."""
    prompts: list[DetectedPrompt] = []

    # Decorator: @prompt, @server.prompt
    pat = re.compile(
        r'@\w*\.?prompt\s*(?:\([^)]*\))?\s*\n\s*(?:async\s+)?def\s+(\w+)\s*\(([^)]*)\)',
        re.MULTILINE,
    )
    for m in pat.finditer(code):
        name = m.group(1)
        args_str = m.group(2)
        args = [
            a.strip().split(":")[0].strip()
            for a in args_str.split(",")
            if a.strip() and a.strip() not in ("self", "ctx")
        ]
        line = code[:m.start()].count("\n") + 1
        prompts.append(DetectedPrompt(name=name, line=line, arguments=args))

    return prompts


def _detect_server_name(code: str) -> str:
    """Try to detect server name from code."""
    patterns = [
        re.compile(r'FastMCP\s*\(\s*["\']([^"\']+)["\']'),
        re.compile(r'name\s*=\s*["\']([^"\']+)["\']'),
        re.compile(r'server_name\s*=\s*["\']([^"\']+)["\']'),
    ]
    for pat in patterns:
        m = pat.search(code)
        if m:
            return m.group(1)
    return "mcp_server"


def _get_decorator_name(dec: ast.expr) -> str:
    """Extract decorator name from AST node."""
    if isinstance(dec, ast.Name):
        return dec.id
    if isinstance(dec, ast.Attribute):
        return f"{_get_decorator_name(dec.value)}.{dec.attr}"
    if isinstance(dec, ast.Call):
        return _get_decorator_name(dec.func)
    return ""


def _get_call_name(node: ast.Call) -> str:
    """Extract function call name from AST node."""
    if isinstance(node.func, ast.Name):
        return node.func.id
    if isinstance(node.func, ast.Attribute):
        return node.func.attr
    return ""


def detect_server_components(code: str) -> ServerProfile:
    """Detect all MCP server components from source code.

    Args:
        code: Python source code of an MCP server.

    Returns:
        ServerProfile with detected tools, resources, prompts.
    """
    profile = ServerProfile()
    profile.framework = _detect_framework(code)
    profile.server_name = _detect_server_name(code)
    profile.tools = _detect_tools_ast(code)
    profile.resources = _detect_resources(code)
    profile.prompts = _detect_prompts(code)
    profile.has_async = "async def" in code
    return profile


# ---------------------------------------------------------------------------
# Test generation
# ---------------------------------------------------------------------------

def _sanitize_name(name: str) -> str:
    """Convert name to valid Python test function name."""
    return re.sub(r'[^a-zA-Z0-9_]', '_', name).lower()


def _generate_tool_tests(tool: DetectedTool, framework: str) -> list[GeneratedTest]:
    """Generate tests for a single tool."""
    tests: list[GeneratedTest] = []
    safe_name = _sanitize_name(tool.name)
    async_prefix = "async " if tool.has_async else ""
    await_kw = "await " if tool.has_async else ""
    pytest_mark = "@pytest.mark.asyncio\n    " if tool.has_async else ""

    # 1. Invocation test
    param_dict = ", ".join(f'"{p}": "test_value"' for p in tool.parameters) if tool.parameters else ""
    param_arg = "{" + param_dict + "}" if param_dict else "{}"

    tests.append(GeneratedTest(
        name=f"test_{safe_name}_invocation",
        category="invocation",
        component=tool.name,
        code=textwrap.dedent(f"""\
    {pytest_mark}{async_prefix}def test_{safe_name}_invocation(mcp_client):
        \"\"\"Test that tool '{tool.name}' can be invoked successfully.\"\"\"
        result = {await_kw}mcp_client.call_tool("{tool.name}", {param_arg})
        assert result is not None
        assert "error" not in str(result).lower() or "isError" not in result
"""),
    ))

    # 2. Error handling test
    tests.append(GeneratedTest(
        name=f"test_{safe_name}_error_handling",
        category="error_handling",
        component=tool.name,
        code=textwrap.dedent(f"""\
    {pytest_mark}{async_prefix}def test_{safe_name}_error_handling(mcp_client):
        \"\"\"Test that tool '{tool.name}' handles invalid input gracefully.\"\"\"
        # Pass empty arguments
        result = {await_kw}mcp_client.call_tool("{tool.name}", {{}})
        # Should not crash — either returns error or handles gracefully
        assert result is not None

    {pytest_mark}{async_prefix}def test_{safe_name}_invalid_types(mcp_client):
        \"\"\"Test that tool '{tool.name}' rejects invalid parameter types.\"\"\"
        result = {await_kw}mcp_client.call_tool("{tool.name}", {{"__invalid__": None}})
        assert result is not None
"""),
    ))

    # 3. Permissions test
    tests.append(GeneratedTest(
        name=f"test_{safe_name}_permissions",
        category="permissions",
        component=tool.name,
        code=textwrap.dedent(f"""\
    {pytest_mark}{async_prefix}def test_{safe_name}_permissions(mcp_client):
        \"\"\"Test that tool '{tool.name}' respects permission boundaries.\"\"\"
        # Verify tool is listed in available tools
        tools = {await_kw}mcp_client.list_tools()
        tool_names = [t.name if hasattr(t, 'name') else t['name'] for t in tools]
        assert "{tool.name}" in tool_names

    {pytest_mark}{async_prefix}def test_{safe_name}_no_side_effects(mcp_client):
        \"\"\"Test that tool '{tool.name}' invocation is repeatable without side effects.\"\"\"
        result1 = {await_kw}mcp_client.call_tool("{tool.name}", {param_arg})
        result2 = {await_kw}mcp_client.call_tool("{tool.name}", {param_arg})
        # Both calls should succeed
        assert result1 is not None
        assert result2 is not None
"""),
    ))

    # 4. Edge case test
    if tool.parameters:
        edge_params = ", ".join(f'"{p}": ""' for p in tool.parameters)
        tests.append(GeneratedTest(
            name=f"test_{safe_name}_edge_empty_params",
            category="edge_case",
            component=tool.name,
            code=textwrap.dedent(f"""\
    {pytest_mark}{async_prefix}def test_{safe_name}_edge_empty_params(mcp_client):
        \"\"\"Test tool '{tool.name}' with empty string parameters.\"\"\"
        result = {await_kw}mcp_client.call_tool("{tool.name}", {{{edge_params}}})
        assert result is not None
"""),
        ))

    return tests


def _generate_resource_tests(resource: DetectedResource, framework: str) -> list[GeneratedTest]:
    """Generate tests for a single resource."""
    tests: list[GeneratedTest] = []
    safe_name = _sanitize_name(resource.name)

    tests.append(GeneratedTest(
        name=f"test_resource_{safe_name}_read",
        category="invocation",
        component=resource.name,
        code=textwrap.dedent(f"""\
    async def test_resource_{safe_name}_read(mcp_client):
        \"\"\"Test that resource '{resource.name}' can be read.\"\"\"
        result = await mcp_client.read_resource("{resource.uri_pattern or resource.name}")
        assert result is not None
"""),
    ))

    tests.append(GeneratedTest(
        name=f"test_resource_{safe_name}_error",
        category="error_handling",
        component=resource.name,
        code=textwrap.dedent(f"""\
    async def test_resource_{safe_name}_error(mcp_client):
        \"\"\"Test that invalid resource URI returns proper error.\"\"\"
        try:
            result = await mcp_client.read_resource("{resource.uri_pattern or resource.name}/__invalid__")
            # Should handle gracefully
        except Exception:
            pass  # Expected for invalid URIs
"""),
    ))

    tests.append(GeneratedTest(
        name=f"test_resource_{safe_name}_listed",
        category="permissions",
        component=resource.name,
        code=textwrap.dedent(f"""\
    async def test_resource_{safe_name}_listed(mcp_client):
        \"\"\"Test that resource '{resource.name}' appears in resource listing.\"\"\"
        resources = await mcp_client.list_resources()
        assert len(resources) > 0
"""),
    ))

    return tests


def _generate_prompt_tests(prompt: DetectedPrompt, framework: str) -> list[GeneratedTest]:
    """Generate tests for a single prompt."""
    tests: list[GeneratedTest] = []
    safe_name = _sanitize_name(prompt.name)

    args_dict = ", ".join(f'"{a}": "test"' for a in prompt.arguments) if prompt.arguments else ""
    args = "{" + args_dict + "}" if args_dict else "{}"

    tests.append(GeneratedTest(
        name=f"test_prompt_{safe_name}_get",
        category="invocation",
        component=prompt.name,
        code=textwrap.dedent(f"""\
    async def test_prompt_{safe_name}_get(mcp_client):
        \"\"\"Test that prompt '{prompt.name}' can be retrieved.\"\"\"
        result = await mcp_client.get_prompt("{prompt.name}", {args})
        assert result is not None
"""),
    ))

    tests.append(GeneratedTest(
        name=f"test_prompt_{safe_name}_error",
        category="error_handling",
        component=prompt.name,
        code=textwrap.dedent(f"""\
    async def test_prompt_{safe_name}_error(mcp_client):
        \"\"\"Test prompt '{prompt.name}' with missing arguments.\"\"\"
        try:
            result = await mcp_client.get_prompt("{prompt.name}", {{}})
            assert result is not None
        except Exception:
            pass  # Expected if arguments are required
"""),
    ))

    tests.append(GeneratedTest(
        name=f"test_prompt_{safe_name}_listed",
        category="permissions",
        component=prompt.name,
        code=textwrap.dedent(f"""\
    async def test_prompt_{safe_name}_listed(mcp_client):
        \"\"\"Test that prompt '{prompt.name}' appears in prompt listing.\"\"\"
        prompts = await mcp_client.list_prompts()
        prompt_names = [p.name if hasattr(p, 'name') else p['name'] for p in prompts]
        assert "{prompt.name}" in prompt_names
"""),
    ))

    return tests


def _generate_fixture(profile: ServerProfile) -> str:
    """Generate pytest fixture for MCP client."""
    return textwrap.dedent(f"""\
import pytest


@pytest.fixture
def mcp_client():
    \"\"\"Create an MCP client connected to the server under test.

    TODO: Replace with actual MCP client initialization.
    Framework detected: {profile.framework}
    Server: {profile.server_name}
    \"\"\"
    # Example for FastMCP:
    # from mcp.client import ClientSession
    # from mcp.client.stdio import stdio_client
    # async with stdio_client("python", "server.py") as (read, write):
    #     async with ClientSession(read, write) as client:
    #         yield client
    raise NotImplementedError("Configure MCP client for your server")

""")


# ---------------------------------------------------------------------------
# Main generation function
# ---------------------------------------------------------------------------

def generate_tests(code: str) -> str:
    """Generate pytest-compatible integration tests for an MCP server.

    Analyzes the source code to detect tools, resources, and prompts,
    then generates test stubs for each component.

    Args:
        code: Python source code of an MCP server.

    Returns:
        Complete Python test file as a string.
    """
    profile = detect_server_components(code)
    suite = generate_test_suite(code)

    parts: list[str] = []
    parts.append(f'"""Auto-generated integration tests for MCP server: {profile.server_name}.')
    parts.append("")
    parts.append(f"Framework: {profile.framework}")
    parts.append(f"Tools: {len(profile.tools)}")
    parts.append(f"Resources: {len(profile.resources)}")
    parts.append(f"Prompts: {len(profile.prompts)}")
    parts.append(f'Total tests: {suite.total_tests}')
    parts.append('"""')
    parts.append("")
    parts.append("import pytest")
    parts.append("")
    parts.append("")

    # Fixture
    parts.append(_generate_fixture(profile))

    # Group tests by component
    current_component = ""
    for test in suite.tests:
        if test.component != current_component:
            current_component = test.component
            parts.append(f"\n# --- Tests for: {current_component} ---\n")
        parts.append(test.code)
        parts.append("")

    return "\n".join(parts)


def generate_test_suite(code: str) -> GeneratedSuite:
    """Generate a GeneratedSuite object from MCP server code.

    Args:
        code: Python source code of an MCP server.

    Returns:
        GeneratedSuite with all generated test cases.
    """
    profile = detect_server_components(code)
    tests: list[GeneratedTest] = []

    for tool in profile.tools:
        tests.extend(_generate_tool_tests(tool, profile.framework))

    for resource in profile.resources:
        tests.extend(_generate_resource_tests(resource, profile.framework))

    for prompt in profile.prompts:
        tests.extend(_generate_prompt_tests(prompt, profile.framework))

    return GeneratedSuite(
        tests=tests,
        profile=profile,
        total_tests=len(tests),
    )


# ---------------------------------------------------------------------------
# Formatting
# ---------------------------------------------------------------------------

def format_suite_text(suite: GeneratedSuite) -> str:
    """Format test suite summary as text."""
    lines = ["=" * 60, "MCP Integration Test Generator", "=" * 60]
    p = suite.profile
    lines.append(f"Server:     {p.server_name}")
    lines.append(f"Framework:  {p.framework}")
    lines.append(f"Tools:      {len(p.tools)}")
    lines.append(f"Resources:  {len(p.resources)}")
    lines.append(f"Prompts:    {len(p.prompts)}")
    lines.append(f"Tests:      {suite.total_tests}")
    lines.append("")

    cats: dict[str, int] = {}
    for t in suite.tests:
        cats[t.category] = cats.get(t.category, 0) + 1
    for cat, count in sorted(cats.items()):
        lines.append(f"  {cat:<20} {count} tests")

    return "\n".join(lines)


def format_suite_json(suite: GeneratedSuite) -> str:
    """Format test suite as JSON."""
    return json.dumps(suite.to_dict(), indent=2, ensure_ascii=False)
