"""MCP-001/MCP-002: MCP-specific tool safety rules.

These rules detect patterns specific to MCP server implementations,
checking for unsafe tool patterns that are unique to the MCP protocol.
"""

from __future__ import annotations

import ast
import re

from . import Rule
from ..models import Finding, Severity


class MCPToolDescriptionRule(Rule):
    """Detects MCP tool functions missing descriptions or docstrings."""

    @property
    def rule_id(self) -> str:
        return "MCP-001"

    @property
    def name(self) -> str:
        return "MCP Tool Missing Description"

    @property
    def severity(self) -> Severity:
        return Severity.LOW

    @property
    def description(self) -> str:
        return "MCP tool handler function has no docstring."

    def check(self, tree: ast.Module, source: str) -> list[Finding]:
        findings: list[Finding] = []
        for node in ast.walk(tree):
            if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                continue
            if not self._is_tool_handler(node):
                continue
            docstring = ast.get_docstring(node)
            if not docstring:
                findings.append(Finding(
                    rule_id=self.rule_id,
                    severity=self.severity,
                    title="MCP tool handler missing docstring",
                    description=(
                        f"Function '{node.name}' appears to be an MCP tool "
                        f"handler but has no docstring."
                    ),
                    file="",
                    line=node.lineno,
                    code_snippet=self._get_snippet(source, node.lineno),
                    recommendation=(
                        "Add a docstring describing the tool's purpose, "
                        "parameters, and security considerations."
                    ),
                    cwe="CWE-1059",
                    confidence="MEDIUM",
                ))
        return findings

    def _is_tool_handler(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> bool:
        """Check if function looks like an MCP tool handler."""
        name_lower = node.name.lower()
        if any(kw in name_lower for kw in ("tool", "handle")):
            return True
        for dec in node.decorator_list:
            dec_str = ast.dump(dec)
            if "tool" in dec_str.lower():
                return True
        return False


class MCPEnvAccessRule(Rule):
    """Detects direct environment variable access in tool handlers."""

    @property
    def rule_id(self) -> str:
        return "MCP-002"

    @property
    def name(self) -> str:
        return "Direct Environment Access in Tool"

    @property
    def severity(self) -> Severity:
        return Severity.MEDIUM

    @property
    def description(self) -> str:
        return "Tool handler accesses environment variables directly."

    def check(self, tree: ast.Module, source: str) -> list[Finding]:
        findings: list[Finding] = []
        env_patterns = {"os.environ", "os.getenv"}

        for node in ast.walk(tree):
            if not isinstance(node, ast.Call):
                continue
            func_name = self._get_func_name(node)
            if func_name not in env_patterns:
                continue

            parent_func = self._find_parent_function(tree, node)
            if parent_func and self._is_tool_context(parent_func, source):
                findings.append(Finding(
                    rule_id=self.rule_id,
                    severity=self.severity,
                    title="Environment variable access in tool handler",
                    description=(
                        f"{func_name}() called inside a tool handler. "
                        f"Consider loading env vars at startup instead."
                    ),
                    file="",
                    line=node.lineno,
                    code_snippet=self._get_snippet(source, node.lineno),
                    recommendation=(
                        "Load environment variables at module level or "
                        "startup, not inside tool handlers."
                    ),
                    cwe="CWE-526",
                    confidence="MEDIUM",
                ))
        return findings

    def _is_tool_context(
        self,
        func: ast.FunctionDef | ast.AsyncFunctionDef,
        source: str,
    ) -> bool:
        """Check if function is in a tool-handling context."""
        name_lower = func.name.lower()
        if any(kw in name_lower for kw in ("tool", "handle", "call_tool")):
            return True
        for dec in func.decorator_list:
            if "tool" in ast.dump(dec).lower():
                return True
        return False
