"""Tests for MCP Tool Permission Mapper."""

from __future__ import annotations

import json
import textwrap

import pytest

from mcpaudit.tool_perm_mapper import (
    APICategory,
    APIPattern,
    APIUsage,
    ALL_API_CATEGORIES,
    Recommendation,
    ToolMapping,
    ToolPermissionMap,
    _analyze_code_apis,
    _analyze_code_ast,
    _compute_tool_risk,
    _extract_tool_functions_ast,
    _extract_tool_functions_regex,
    _get_call_name,
    format_perm_map_json,
    format_perm_map_text,
    format_recommendations_json,
    format_recommendations_text,
    get_api_categories,
    get_api_patterns,
    map_tool_permissions,
    recommend_permissions,
)


# ---- APICategory ----

class TestAPICategory:
    def test_filesystem(self):
        assert APICategory.FILESYSTEM == "filesystem"

    def test_network(self):
        assert APICategory.NETWORK == "network"

    def test_shell(self):
        assert APICategory.SHELL == "shell"

    def test_database(self):
        assert APICategory.DATABASE == "database"

    def test_crypto(self):
        assert APICategory.CRYPTO == "crypto"

    def test_environment(self):
        assert APICategory.ENVIRONMENT == "environment"

    def test_system(self):
        assert APICategory.SYSTEM == "system"

    def test_serialization(self):
        assert APICategory.SERIALIZATION == "serialization"

    def test_dynamic_code(self):
        assert APICategory.DYNAMIC_CODE == "dynamic_code"


# ---- ALL_API_CATEGORIES ----

class TestAllCategories:
    def test_count(self):
        assert len(ALL_API_CATEGORIES) == 9

    def test_contains_filesystem(self):
        assert "filesystem" in ALL_API_CATEGORIES

    def test_contains_network(self):
        assert "network" in ALL_API_CATEGORIES


# ---- APIPattern ----

class TestAPIPattern:
    def test_to_dict(self):
        p = APIPattern(category="test", modules=["os"], functions=["open"], description="desc")
        d = p.to_dict()
        assert d["category"] == "test"
        assert d["modules"] == ["os"]

    def test_default_risk(self):
        p = APIPattern(category="x", modules=[], functions=[], description="d")
        assert p.risk_level == "medium"


# ---- APIUsage ----

class TestAPIUsage:
    def test_defaults(self):
        u = APIUsage(category="fs", line=1, code="open('f')", description="open")
        assert u.risk_level == "medium"

    def test_to_dict(self):
        u = APIUsage(category="net", line=5, code="requests.get()", description="http", risk_level="high")
        d = u.to_dict()
        assert d["category"] == "net"
        assert d["line"] == 5
        assert d["risk_level"] == "high"


# ---- ToolMapping ----

class TestToolMapping:
    def test_defaults(self):
        t = ToolMapping(tool_name="my_tool")
        assert t.api_categories == []
        assert t.risk_level == "low"

    def test_to_dict(self):
        t = ToolMapping(tool_name="t", api_categories=["filesystem"], risk_level="medium")
        d = t.to_dict()
        assert d["tool_name"] == "t"
        assert "filesystem" in d["api_categories"]


# ---- ToolPermissionMap ----

class TestToolPermissionMap:
    def test_empty(self):
        m = ToolPermissionMap()
        assert m.tool_count == 0
        assert m.category_count == 0

    def test_properties(self):
        m = ToolPermissionMap(
            tools=[ToolMapping(tool_name="a", api_categories=["filesystem", "network"])],
            all_categories=["filesystem", "network"],
        )
        assert m.tool_count == 1
        assert m.category_count == 2

    def test_get_matrix(self):
        m = ToolPermissionMap(
            tools=[
                ToolMapping(tool_name="a", api_categories=["filesystem"]),
                ToolMapping(tool_name="b", api_categories=["network"]),
            ]
        )
        matrix = m.get_matrix()
        assert matrix["a"] == ["filesystem"]
        assert matrix["b"] == ["network"]

    def test_to_dict(self):
        m = ToolPermissionMap(tools=[], all_categories=["filesystem"])
        d = m.to_dict()
        assert d["tool_count"] == 0
        assert d["category_count"] == 1

    def test_errors(self):
        m = ToolPermissionMap(errors=["test error"])
        d = m.to_dict()
        assert "test error" in d["errors"]


# ---- Recommendation ----

class TestRecommendation:
    def test_to_dict(self):
        r = Recommendation(
            tool_name="t",
            current_categories=["shell"],
            recommended_categories=["shell"],
            unnecessary=[],
            risk_level="critical",
            message="Risky",
        )
        d = r.to_dict()
        assert d["tool_name"] == "t"
        assert d["risk_level"] == "critical"


# ---- _compute_tool_risk ----

class TestComputeToolRisk:
    def test_empty(self):
        assert _compute_tool_risk([]) == "low"

    def test_filesystem_only(self):
        assert _compute_tool_risk(["filesystem"]) == "medium"

    def test_shell(self):
        assert _compute_tool_risk(["shell"]) == "critical"

    def test_dynamic_code(self):
        assert _compute_tool_risk(["dynamic_code"]) == "critical"

    def test_network(self):
        assert _compute_tool_risk(["network"]) == "high"

    def test_highest_wins(self):
        assert _compute_tool_risk(["crypto", "shell"]) == "critical"

    def test_system_low(self):
        assert _compute_tool_risk(["system"]) == "low"


# ---- _analyze_code_apis (regex-based security pattern detection) ----

class TestAnalyzeCodeAPIs:
    def test_detects_open(self):
        code = 'f = open("test.txt")\n'
        usages = _analyze_code_apis(code)
        cats = [u.category for u in usages]
        assert "filesystem" in cats

    def test_detects_subprocess_pattern(self):
        # Security scanner detecting subprocess usage patterns
        code = 'subprocess.run(["ls"])\n'
        usages = _analyze_code_apis(code)
        cats = [u.category for u in usages]
        assert "shell" in cats

    def test_detects_requests(self):
        code = 'requests.get("http://example.com")\n'
        usages = _analyze_code_apis(code)
        cats = [u.category for u in usages]
        assert "network" in cats

    def test_detects_sqlite(self):
        code = 'conn = sqlite3.connect("db.sqlite")\n'
        usages = _analyze_code_apis(code)
        cats = [u.category for u in usages]
        assert "database" in cats

    def test_detects_os_environ(self):
        code = 'key = os.environ["API_KEY"]\n'
        usages = _analyze_code_apis(code)
        cats = [u.category for u in usages]
        assert "environment" in cats

    def test_skips_comments(self):
        code = '# open("test.txt")\n'
        usages = _analyze_code_apis(code)
        assert len(usages) == 0

    def test_skips_empty(self):
        code = '\n\n\n'
        usages = _analyze_code_apis(code)
        assert len(usages) == 0

    def test_offset(self):
        code = 'subprocess.run(["ls"])\n'
        usages = _analyze_code_apis(code, offset=10)
        assert usages[0].line == 11

    def test_detects_dynamic_code_pattern(self):
        # Security scanner detecting dynamic code execution patterns
        code = 'result = eval("1+1")\n'
        usages = _analyze_code_apis(code)
        cats = [u.category for u in usages]
        assert "dynamic_code" in cats

    def test_detects_hashlib(self):
        code = 'h = hashlib.sha256(data)\n'
        usages = _analyze_code_apis(code)
        cats = [u.category for u in usages]
        assert "crypto" in cats


# ---- _extract_tool_functions_ast ----

class TestExtractToolFunctionsAST:
    def test_decorated_tool(self):
        code = textwrap.dedent("""\
            import os

            @server.tool()
            def read_file(path: str) -> str:
                return open(path).read()
        """)
        tools = _extract_tool_functions_ast(code)
        assert len(tools) >= 1
        assert tools[0][0] == "read_file"

    def test_named_tool_decorator(self):
        code = textwrap.dedent("""\
            @server.tool("my_tool")
            def handler(x):
                return x
        """)
        tools = _extract_tool_functions_ast(code)
        assert len(tools) >= 1
        assert tools[0][0] == "my_tool"

    def test_tool_prefix_heuristic(self):
        code = textwrap.dedent("""\
            def tool_read(path):
                return open(path).read()
        """)
        tools = _extract_tool_functions_ast(code)
        assert len(tools) >= 1

    def test_handle_prefix(self):
        code = textwrap.dedent("""\
            def handle_request(req):
                pass
        """)
        tools = _extract_tool_functions_ast(code)
        assert len(tools) >= 1

    def test_no_tools(self):
        code = textwrap.dedent("""\
            def helper(x):
                return x + 1
        """)
        tools = _extract_tool_functions_ast(code)
        assert len(tools) == 0

    def test_async_tool(self):
        code = textwrap.dedent("""\
            @server.tool()
            async def async_tool(x):
                return x
        """)
        tools = _extract_tool_functions_ast(code)
        assert len(tools) >= 1


# ---- map_tool_permissions (main API) ----

class TestMapToolPermissions:
    def test_empty_code(self):
        result = map_tool_permissions("")
        assert result.tool_count == 0
        assert "Empty code" in result.errors[0]

    def test_whitespace_only(self):
        result = map_tool_permissions("   \n\n  ")
        assert result.tool_count == 0

    def test_simple_tool_filesystem(self):
        code = textwrap.dedent("""\
            @server.tool()
            def read_file(path: str) -> str:
                return open(path).read()
        """)
        result = map_tool_permissions(code)
        assert result.tool_count >= 1

    def test_tool_with_network(self):
        code = textwrap.dedent("""\
            @server.tool()
            def fetch_url(url: str) -> str:
                import urllib.request
                return urllib.request.urlopen(url).read()
        """)
        result = map_tool_permissions(code)
        assert result.tool_count >= 1

    def test_global_fallback(self):
        code = textwrap.dedent("""\
            import os
            result = os.environ["KEY"]
        """)
        result = map_tool_permissions(code)
        assert result.tool_count >= 1
        assert "<global>" in [t.tool_name for t in result.tools]

    def test_multiple_tools(self):
        code = textwrap.dedent("""\
            @server.tool()
            def tool_read(path):
                return open(path).read()

            @server.tool()
            def tool_check(key):
                import os
                return os.environ[key]
        """)
        result = map_tool_permissions(code)
        assert result.tool_count >= 2

    def test_all_categories_populated(self):
        code = textwrap.dedent("""\
            @server.tool()
            def tool_multi():
                import os
                os.environ["X"]
                open("test.txt")
        """)
        result = map_tool_permissions(code)
        assert len(result.all_categories) >= 1

    def test_to_dict(self):
        code = textwrap.dedent("""\
            @server.tool()
            def tool_test():
                open("f")
        """)
        result = map_tool_permissions(code)
        d = result.to_dict()
        assert "tool_count" in d
        assert "tools" in d


# ---- recommend_permissions ----

class TestRecommendPermissions:
    def test_empty_map(self):
        m = ToolPermissionMap()
        recs = recommend_permissions(m)
        assert len(recs) == 0

    def test_low_risk_tool(self):
        m = ToolPermissionMap(tools=[
            ToolMapping(tool_name="safe", api_categories=["crypto"], risk_level="low"),
        ])
        recs = recommend_permissions(m)
        assert len(recs) == 1
        assert "minimal" in recs[0].message.lower()

    def test_high_risk_tool(self):
        m = ToolPermissionMap(tools=[
            ToolMapping(tool_name="risky", api_categories=["shell", "network"], risk_level="critical"),
        ])
        recs = recommend_permissions(m)
        assert len(recs) == 1
        assert "high-risk" in recs[0].message.lower()

    def test_broad_access(self):
        m = ToolPermissionMap(tools=[
            ToolMapping(tool_name="broad", api_categories=["filesystem", "network", "shell", "database"], risk_level="critical"),
        ])
        recs = recommend_permissions(m)
        assert len(recs) == 1
        assert "categories" in recs[0].message

    def test_shell_and_network_warning(self):
        m = ToolPermissionMap(tools=[
            ToolMapping(tool_name="danger", api_categories=["shell", "network"], risk_level="critical"),
        ])
        recs = recommend_permissions(m)
        assert any("remote code execution" in r.message.lower() for r in recs)

    def test_recommendation_to_dict(self):
        m = ToolPermissionMap(tools=[
            ToolMapping(tool_name="t", api_categories=["filesystem"], risk_level="medium"),
        ])
        recs = recommend_permissions(m)
        d = recs[0].to_dict()
        assert d["tool_name"] == "t"


# ---- get_api_patterns / get_api_categories ----

class TestGetters:
    def test_get_patterns_not_empty(self):
        patterns = get_api_patterns()
        assert len(patterns) > 0

    def test_patterns_have_category(self):
        for p in get_api_patterns():
            assert p.category in ALL_API_CATEGORIES

    def test_get_categories(self):
        cats = get_api_categories()
        assert len(cats) == 9


# ---- format_perm_map_text ----

class TestFormatText:
    def test_basic(self):
        m = ToolPermissionMap(
            tools=[ToolMapping(tool_name="t", api_categories=["filesystem"])],
            all_categories=["filesystem"],
        )
        text = format_perm_map_text(m)
        assert "Tool Permission Map" in text
        assert "filesystem" in text

    def test_with_errors(self):
        m = ToolPermissionMap(errors=["error1"])
        text = format_perm_map_text(m)
        assert "error1" in text


# ---- format_perm_map_json ----

class TestFormatJSON:
    def test_valid_json(self):
        m = ToolPermissionMap(
            tools=[ToolMapping(tool_name="t", api_categories=["network"])],
            all_categories=["network"],
        )
        result = format_perm_map_json(m)
        data = json.loads(result)
        assert data["tool_count"] == 1

    def test_empty_map(self):
        m = ToolPermissionMap()
        result = format_perm_map_json(m)
        data = json.loads(result)
        assert data["tool_count"] == 0


# ---- format_recommendations_text ----

class TestFormatRecommendationsText:
    def test_basic(self):
        recs = [Recommendation(tool_name="t", current_categories=["shell"],
                               recommended_categories=["shell"], unnecessary=[],
                               risk_level="critical", message="Risky tool")]
        text = format_recommendations_text(recs)
        assert "Least Privilege" in text
        assert "CRITICAL" in text

    def test_with_unnecessary(self):
        recs = [Recommendation(tool_name="t", current_categories=["shell", "crypto"],
                               recommended_categories=["shell"], unnecessary=["crypto"],
                               risk_level="high", message="Remove crypto")]
        text = format_recommendations_text(recs)
        assert "crypto" in text


# ---- format_recommendations_json ----

class TestFormatRecommendationsJSON:
    def test_valid_json(self):
        recs = [Recommendation(tool_name="t", current_categories=[], recommended_categories=[],
                               unnecessary=[], risk_level="low", message="OK")]
        result = format_recommendations_json(recs)
        data = json.loads(result)
        assert len(data["recommendations"]) == 1
