"""Tool Permission Mapper — Maps MCP tools to required permissions via AST.

Analyzes MCP server source code to determine what API categories each tool
uses (filesystem, network, shell, database, etc.) and generates:
- Permission matrix: tools x API access
- Least privilege recommendations per tool
- Risk assessment for overprivileged tools

Zero external dependencies — pure Python stdlib.
"""

from __future__ import annotations

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


# ---------- API Categories ----------

class APICategory:
    """Standard API access categories."""

    FILESYSTEM = "filesystem"
    NETWORK = "network"
    SHELL = "shell"
    DATABASE = "database"
    CRYPTO = "crypto"
    ENVIRONMENT = "environment"
    SYSTEM = "system"
    SERIALIZATION = "serialization"
    DYNAMIC_CODE = "dynamic_code"


ALL_API_CATEGORIES: list[str] = [
    APICategory.FILESYSTEM,
    APICategory.NETWORK,
    APICategory.SHELL,
    APICategory.DATABASE,
    APICategory.CRYPTO,
    APICategory.ENVIRONMENT,
    APICategory.SYSTEM,
    APICategory.SERIALIZATION,
    APICategory.DYNAMIC_CODE,
]


# ---------- AST Pattern Registry ----------

@dataclass
class APIPattern:
    """A code pattern mapped to an API category."""

    category: str
    modules: list[str]  # module names
    functions: list[str]  # function/method names
    description: str
    risk_level: str = "medium"  # low, medium, high, critical

    def to_dict(self) -> dict[str, Any]:
        return {
            "category": self.category,
            "modules": self.modules,
            "functions": self.functions,
            "description": self.description,
            "risk_level": self.risk_level,
        }


# Built-in API patterns for detection
_API_PATTERNS: list[APIPattern] = [
    # Filesystem
    APIPattern(APICategory.FILESYSTEM, ["os", "os.path", "pathlib", "shutil", "glob", "tempfile", "fnmatch"],
               ["open", "read", "write", "readlines", "mkdir", "makedirs", "remove", "unlink",
                "rmdir", "rename", "listdir", "scandir", "walk", "glob", "iglob",
                "read_text", "write_text", "read_bytes", "write_bytes",
                "exists", "isfile", "isdir", "getsize", "copy", "copy2", "copytree",
                "rmtree", "move", "touch", "iterdir", "rglob", "stat"],
               "Filesystem access", "medium"),
    # Network
    APIPattern(APICategory.NETWORK, ["urllib", "urllib.request", "http", "http.client", "http.server",
                                      "socket", "requests", "httpx", "aiohttp", "smtplib",
                                      "ftplib", "xmlrpc", "ssl"],
               ["urlopen", "Request", "get", "post", "put", "delete", "patch",
                "connect", "create_connection", "socket", "ClientSession",
                "SMTP", "FTP", "ServerProxy"],
               "Network access", "high"),
    # Shell
    APIPattern(APICategory.SHELL, ["subprocess", "os", "shlex"],
               ["run", "Popen", "call", "check_output", "check_call",
                "system", "popen", "execl", "execle", "execlp", "execv", "execve",
                "execvp", "spawnl", "spawnle", "spawnlp", "spawnv", "spawnve", "spawnvp"],
               "Shell/subprocess execution", "critical"),
    # Database
    APIPattern(APICategory.DATABASE, ["sqlite3", "pymysql", "psycopg2", "asyncpg", "cx_Oracle",
                                       "pymongo", "redis"],
               ["connect", "Connection", "cursor", "execute", "executemany",
                "fetchall", "fetchone", "fetchmany", "commit", "rollback"],
               "Database access", "high"),
    # Crypto
    APIPattern(APICategory.CRYPTO, ["hashlib", "hmac", "secrets", "ssl", "cryptography"],
               ["sha256", "sha512", "md5", "sha1", "pbkdf2_hmac",
                "new", "digest", "compare_digest", "token_hex", "token_bytes",
                "token_urlsafe", "create_default_context", "SSLContext"],
               "Cryptographic operations", "low"),
    # Environment
    APIPattern(APICategory.ENVIRONMENT, ["os"],
               ["environ", "getenv", "putenv"],
               "Environment variable access", "medium"),
    # System
    APIPattern(APICategory.SYSTEM, ["platform", "os", "sys", "resource", "signal"],
               ["system", "node", "release", "machine", "processor",
                "getpid", "getppid", "getuid", "getgid", "uname", "cpu_count",
                "platform", "version", "executable", "argv",
                "signal", "alarm"],
               "System information access", "low"),
    # Serialization (note: detection patterns for security scanning, not usage of unsafe libs)
    APIPattern(APICategory.SERIALIZATION, ["shelve", "marshal", "yaml", "json", "xml"],
               ["load", "loads", "dump", "dumps", "unsafe_load",
                "parse", "fromstring"],
               "Data serialization", "medium"),
    # Dynamic code
    APIPattern(APICategory.DYNAMIC_CODE, ["importlib", "builtins"],
               ["eval", "exec", "compile", "__import__", "import_module", "reload"],
               "Dynamic code execution", "critical"),
]


# ---------- Regex patterns for non-AST fallback ----------

_REGEX_PATTERNS: list[tuple[str, str]] = [
    (APICategory.FILESYSTEM, r'\b(open|Path|os\.path|shutil|glob)\s*[\.(]'),
    (APICategory.NETWORK, r'\b(urllib|requests|httpx|aiohttp|socket|http\.client|smtplib)\b'),
    (APICategory.SHELL, r'\b(subprocess\.(run|Popen|call|check_output)|os\.(system|popen|exec))\b'),
    (APICategory.DATABASE, r'\b(sqlite3|pymysql|psycopg2|asyncpg|cursor\.(execute|fetch))\b'),
    (APICategory.CRYPTO, r'\b(hashlib|hmac|secrets|ssl\.create_default_context)\b'),
    (APICategory.ENVIRONMENT, r'\b(os\.environ|os\.getenv|os\.putenv)\b'),
    (APICategory.SYSTEM, r'\b(platform\.(system|node)|os\.(getpid|uname)|sys\.(platform|argv))\b'),
    (APICategory.SERIALIZATION, r'\b(yaml\.(load|unsafe_load)|marshal)\b'),
    (APICategory.DYNAMIC_CODE, r'\b(eval|exec|compile|__import__|importlib)\s*\('),
]


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

@dataclass
class APIUsage:
    """A detected API usage in a tool."""

    category: str
    line: int
    code: str
    description: str
    risk_level: str = "medium"

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


@dataclass
class ToolMapping:
    """Permission mapping for a single tool."""

    tool_name: str
    start_line: int = 0
    end_line: int = 0
    api_categories: list[str] = field(default_factory=list)
    usages: list[APIUsage] = field(default_factory=list)
    risk_level: str = "low"  # overall risk level

    def to_dict(self) -> dict[str, Any]:
        return {
            "tool_name": self.tool_name,
            "start_line": self.start_line,
            "end_line": self.end_line,
            "api_categories": self.api_categories,
            "usages": [u.to_dict() for u in self.usages],
            "risk_level": self.risk_level,
        }


@dataclass
class ToolPermissionMap:
    """Complete permission map for all tools in a codebase."""

    tools: list[ToolMapping] = field(default_factory=list)
    global_usages: list[APIUsage] = field(default_factory=list)
    all_categories: list[str] = field(default_factory=list)
    errors: list[str] = field(default_factory=list)

    @property
    def tool_count(self) -> int:
        return len(self.tools)

    @property
    def category_count(self) -> int:
        return len(self.all_categories)

    def get_matrix(self) -> dict[str, list[str]]:
        """Return tool-name -> list-of-categories mapping."""
        return {t.tool_name: t.api_categories for t in self.tools}

    def to_dict(self) -> dict[str, Any]:
        return {
            "tool_count": self.tool_count,
            "category_count": self.category_count,
            "all_categories": self.all_categories,
            "tools": [t.to_dict() for t in self.tools],
            "global_usages": [u.to_dict() for u in self.global_usages],
            "errors": self.errors,
        }


@dataclass
class Recommendation:
    """A least-privilege recommendation for a tool."""

    tool_name: str
    current_categories: list[str]
    recommended_categories: list[str]
    unnecessary: list[str]
    risk_level: str
    message: str

    def to_dict(self) -> dict[str, Any]:
        return {
            "tool_name": self.tool_name,
            "current_categories": self.current_categories,
            "recommended_categories": self.recommended_categories,
            "unnecessary": self.unnecessary,
            "risk_level": self.risk_level,
            "message": self.message,
        }


# ---------- AST Analysis ----------

def _extract_tool_functions_ast(code: str) -> list[tuple[str, int, int]]:
    """Extract tool function boundaries using AST.

    Returns list of (tool_name, start_line, end_line).
    """
    tools: list[tuple[str, int, int]] = []

    try:
        tree = ast.parse(code)
    except SyntaxError:
        return _extract_tool_functions_regex(code)

    for node in ast.walk(tree):
        if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            continue

        is_tool = False
        tool_name = node.name

        # Check decorators
        for dec in node.decorator_list:
            dec_name = ""
            if isinstance(dec, ast.Call):
                if isinstance(dec.func, ast.Attribute):
                    dec_name = dec.func.attr
                    # Extract tool name from decorator args if present
                    if dec.args and isinstance(dec.args[0], ast.Constant):
                        tool_name = str(dec.args[0].value)
                elif isinstance(dec.func, ast.Name):
                    dec_name = dec.func.id
            elif isinstance(dec, ast.Attribute):
                dec_name = dec.attr

            if dec_name in ("tool", "command", "handler"):
                is_tool = True

        # Heuristic: function names suggesting tool handlers
        if not is_tool:
            tool_keywords = ("tool_", "handle_", "execute_", "run_", "cmd_")
            if any(node.name.startswith(kw) for kw in tool_keywords):
                is_tool = True

        if is_tool:
            end_line = _ast_end_line(node)
            tools.append((tool_name, node.lineno, end_line))

    return tools


def _ast_end_line(node: ast.AST) -> int:
    """Get end line of an AST node."""
    if hasattr(node, "end_lineno") and node.end_lineno is not None:
        return node.end_lineno
    # Fallback: find max line in children
    max_line = getattr(node, "lineno", 0)
    for child in ast.walk(node):
        child_line = getattr(child, "lineno", 0)
        if child_line > max_line:
            max_line = child_line
    return max_line


def _extract_tool_functions_regex(code: str) -> list[tuple[str, int, int]]:
    """Fallback regex extraction for non-parseable code."""
    tools: list[tuple[str, int, int]] = []
    lines = code.splitlines()
    decorator_re = re.compile(r'@\w+\.tool\s*\(\s*(?:["\'](\w+)["\'])?\s*\)')
    func_re = re.compile(r'^(?:async\s+)?def\s+(\w+)\s*\(')

    i = 0
    while i < len(lines):
        line = lines[i].strip()
        m = decorator_re.search(line)
        if m:
            j = i + 1
            while j < len(lines) and j - i < 5:
                fm = func_re.match(lines[j].strip())
                if fm:
                    name = m.group(1) or fm.group(1)
                    end = _find_func_end_regex(lines, j)
                    tools.append((name, j + 1, end))
                    i = end
                    break
                j += 1
        i += 1

    if not tools:
        for i, line in enumerate(lines):
            fm = func_re.match(line.strip())
            if fm:
                name = fm.group(1)
                if any(kw in name.lower() for kw in ("tool", "handle", "execute")):
                    end = _find_func_end_regex(lines, i)
                    tools.append((name, i + 1, end))

    return tools


def _find_func_end_regex(lines: list[str], start: int) -> int:
    """Find function end by indentation."""
    if start >= len(lines):
        return start
    indent = len(lines[start]) - len(lines[start].lstrip())
    end = start + 1
    while end < len(lines):
        line = lines[end]
        stripped = line.strip()
        if stripped and not stripped.startswith("#"):
            line_indent = len(line) - len(line.lstrip())
            if line_indent <= indent:
                break
        end += 1
    return end


def _analyze_code_apis(code: str, offset: int = 0) -> list[APIUsage]:
    """Analyze code for API usages using regex patterns.

    Args:
        code: Source code to analyze.
        offset: Line number offset for absolute positioning.

    Returns:
        List of detected API usages.
    """
    usages: list[APIUsage] = []
    lines = code.splitlines()

    for line_num, line_text in enumerate(lines, 1):
        stripped = line_text.strip()
        if stripped.startswith("#") or not stripped:
            continue

        for category, pattern in _REGEX_PATTERNS:
            if re.search(pattern, line_text):
                # Find risk level from API patterns
                risk = "medium"
                desc = category
                for ap in _API_PATTERNS:
                    if ap.category == category:
                        risk = ap.risk_level
                        desc = ap.description
                        break

                usages.append(APIUsage(
                    category=category,
                    line=line_num + offset,
                    code=stripped[:120],
                    description=desc,
                    risk_level=risk,
                ))

    return usages


def _analyze_code_ast(code: str) -> list[APIUsage]:
    """Analyze code using AST for precise detection."""
    usages: list[APIUsage] = []

    try:
        tree = ast.parse(code)
    except SyntaxError:
        return _analyze_code_apis(code)

    for node in ast.walk(tree):
        # Check imports
        if isinstance(node, ast.Import):
            for alias in node.names:
                for ap in _API_PATTERNS:
                    if alias.name in ap.modules:
                        usages.append(APIUsage(
                            category=ap.category,
                            line=node.lineno,
                            code=f"import {alias.name}",
                            description=ap.description,
                            risk_level=ap.risk_level,
                        ))

        elif isinstance(node, ast.ImportFrom):
            module = node.module or ""
            for ap in _API_PATTERNS:
                if module in ap.modules or any(module.startswith(m + ".") for m in ap.modules):
                    names = ", ".join(a.name for a in (node.names or []))
                    usages.append(APIUsage(
                        category=ap.category,
                        line=node.lineno,
                        code=f"from {module} import {names}",
                        description=ap.description,
                        risk_level=ap.risk_level,
                    ))

        # Check function calls
        elif isinstance(node, ast.Call):
            func_name = _get_call_name(node)
            if func_name:
                for ap in _API_PATTERNS:
                    parts = func_name.split(".")
                    mod_match = any(func_name.startswith(m) for m in ap.modules)
                    func_match = parts[-1] in ap.functions if parts else False
                    if mod_match or func_match:
                        usages.append(APIUsage(
                            category=ap.category,
                            line=getattr(node, "lineno", 0),
                            code=func_name,
                            description=ap.description,
                            risk_level=ap.risk_level,
                        ))

    return usages


def _get_call_name(node: ast.Call) -> str:
    """Extract the full dotted name from a Call node."""
    parts: list[str] = []
    current: ast.expr = node.func

    while True:
        if isinstance(current, ast.Attribute):
            parts.append(current.attr)
            current = current.value
        elif isinstance(current, ast.Name):
            parts.append(current.id)
            break
        else:
            break

    return ".".join(reversed(parts))


def _compute_tool_risk(categories: list[str]) -> str:
    """Compute overall risk level for a tool based on its API categories."""
    risk_map = {
        APICategory.DYNAMIC_CODE: "critical",
        APICategory.SHELL: "critical",
        APICategory.NETWORK: "high",
        APICategory.DATABASE: "high",
        APICategory.FILESYSTEM: "medium",
        APICategory.ENVIRONMENT: "medium",
        APICategory.SERIALIZATION: "medium",
        APICategory.CRYPTO: "low",
        APICategory.SYSTEM: "low",
    }
    risk_order = ["critical", "high", "medium", "low"]
    highest = "low"
    for cat in categories:
        cat_risk = risk_map.get(cat, "low")
        if risk_order.index(cat_risk) < risk_order.index(highest):
            highest = cat_risk
    return highest


# ---------- Public API ----------

def map_tool_permissions(code: str) -> ToolPermissionMap:
    """Map MCP tools to their required permissions via AST analysis.

    Analyzes the given source code to find tool functions and determine
    which API categories each tool accesses.

    Args:
        code: Source code of the MCP server.

    Returns:
        ToolPermissionMap with all detected tools and their permissions.
    """
    result = ToolPermissionMap()

    if not code or not code.strip():
        result.errors.append("Empty code provided")
        return result

    # Global API usage analysis
    global_usages = _analyze_code_ast(code)
    result.global_usages = global_usages

    # Extract tool functions
    tools = _extract_tool_functions_ast(code)
    lines = code.splitlines()

    all_cats: set[str] = set()

    for tool_name, start, end in tools:
        # Extract tool body
        tool_lines = lines[start - 1:end]
        tool_code = "\n".join(tool_lines)

        # Analyze tool-specific API usage
        tool_usages = _analyze_code_apis(tool_code, offset=start - 1)

        categories = sorted(set(u.category for u in tool_usages))
        all_cats.update(categories)

        risk = _compute_tool_risk(categories)

        result.tools.append(ToolMapping(
            tool_name=tool_name,
            start_line=start,
            end_line=end,
            api_categories=categories,
            usages=tool_usages,
            risk_level=risk,
        ))

    # If no tools found, treat entire code as one implicit tool
    if not tools and global_usages:
        categories = sorted(set(u.category for u in global_usages))
        all_cats.update(categories)
        result.tools.append(ToolMapping(
            tool_name="<global>",
            start_line=1,
            end_line=len(lines),
            api_categories=categories,
            usages=global_usages,
            risk_level=_compute_tool_risk(categories),
        ))

    result.all_categories = sorted(all_cats)

    return result


def recommend_permissions(perm_map: ToolPermissionMap) -> list[Recommendation]:
    """Generate least-privilege recommendations for each tool.

    Args:
        perm_map: Result from map_tool_permissions.

    Returns:
        List of Recommendation objects.
    """
    recommendations: list[Recommendation] = []

    high_risk_categories = {APICategory.SHELL, APICategory.DYNAMIC_CODE}
    medium_risk_categories = {APICategory.NETWORK, APICategory.DATABASE, APICategory.FILESYSTEM}

    for tool in perm_map.tools:
        unnecessary: list[str] = []
        msgs: list[str] = []

        # Check for high risk
        risky = [c for c in tool.api_categories if c in high_risk_categories]
        if risky:
            msgs.append(
                f"Tool '{tool.tool_name}' uses high-risk APIs: {', '.join(risky)}. "
                "Consider sandboxing or restricting."
            )

        # Check for broad access
        if len(tool.api_categories) > 3:
            msgs.append(
                f"Tool '{tool.tool_name}' accesses {len(tool.api_categories)} API categories. "
                "Consider splitting into focused tools."
            )

        # Check for unnecessary combos
        if APICategory.SHELL in tool.api_categories and APICategory.NETWORK in tool.api_categories:
            msgs.append(
                f"Tool '{tool.tool_name}' has both shell and network access — "
                "high risk for remote code execution chains."
            )

        # Minimal recommended set
        recommended = list(tool.api_categories)

        if not msgs:
            msgs.append(
                f"Tool '{tool.tool_name}' has minimal permissions ({len(tool.api_categories)} categories)."
            )

        recommendations.append(Recommendation(
            tool_name=tool.tool_name,
            current_categories=tool.api_categories,
            recommended_categories=recommended,
            unnecessary=unnecessary,
            risk_level=tool.risk_level,
            message=" ".join(msgs),
        ))

    return recommendations


def get_api_patterns() -> list[APIPattern]:
    """Return all known API patterns."""
    return list(_API_PATTERNS)


def get_api_categories() -> list[str]:
    """Return all known API categories."""
    return list(ALL_API_CATEGORIES)


# ---------- Formatters ----------

def format_perm_map_text(perm_map: ToolPermissionMap) -> str:
    """Format permission map as human-readable text."""
    lines: list[str] = []
    lines.append("\nTool Permission Map")
    lines.append("=" * 50)
    lines.append(f"Tools: {perm_map.tool_count}")
    lines.append(f"API categories: {perm_map.category_count}")
    lines.append("")

    if perm_map.tools:
        lines.append("Permission Matrix:")
        lines.append(f"  {'Tool':<30} {'Categories':<40} {'Risk':<10}")
        lines.append(f"  {'-'*30} {'-'*40} {'-'*10}")
        for t in perm_map.tools:
            cats = ", ".join(t.api_categories) if t.api_categories else "none"
            lines.append(f"  {t.tool_name:<30} {cats:<40} {t.risk_level:<10}")
        lines.append("")

    if perm_map.errors:
        lines.append("Errors:")
        for err in perm_map.errors:
            lines.append(f"  - {err}")
        lines.append("")

    return "\n".join(lines)


def format_perm_map_json(perm_map: ToolPermissionMap) -> str:
    """Format permission map as JSON."""
    return json.dumps(perm_map.to_dict(), indent=2, ensure_ascii=False)


def format_recommendations_text(recs: list[Recommendation]) -> str:
    """Format recommendations as human-readable text."""
    lines: list[str] = []
    lines.append("\nLeast Privilege Recommendations")
    lines.append("=" * 50)

    for rec in recs:
        lines.append(f"\n[{rec.risk_level.upper()}] {rec.tool_name}")
        lines.append(f"  Categories: {', '.join(rec.current_categories) if rec.current_categories else 'none'}")
        lines.append(f"  {rec.message}")
        if rec.unnecessary:
            lines.append(f"  Unnecessary: {', '.join(rec.unnecessary)}")

    lines.append("")
    return "\n".join(lines)


def format_recommendations_json(recs: list[Recommendation]) -> str:
    """Format recommendations as JSON."""
    return json.dumps(
        {"recommendations": [r.to_dict() for r in recs]},
        indent=2,
        ensure_ascii=False,
    )
