from __future__ import annotations

import ast
import re
from typing import Dict, List, Set, Tuple

from bastion.analyzers.base import BaseAnalyzer
from bastion.models import Capability, Evidence

SYSTEM_EXEC_MODULES = {"subprocess", "ctypes"}
SYSTEM_EXEC_FUNCS: Dict[str, Set[str]] = {
    "subprocess": {"run", "Popen", "call", "check_call", "check_output", "getoutput", "getstatusoutput"},
    "os": {"system", "popen", "execl", "execle", "execlp", "execlpe", "execv", "execve", "execvp", "execvpe"},
    "shutil": {"rmtree", "move"},
}

NETWORK_MODULES = {"requests", "urllib", "httpx", "aiohttp", "socket"}
NETWORK_FUNCS: Dict[str, Set[str]] = {
    "requests": {"get", "post", "put", "delete", "patch", "head", "options", "request", "Session"},
    "urllib.request": {"urlopen", "urlretrieve", "Request"},
    "httpx": {"get", "post", "put", "delete", "patch", "head", "options", "request", "Client", "AsyncClient"},
    "aiohttp": {"ClientSession", "request"},
    "socket": {"socket", "create_connection", "getaddrinfo"},
}

SECRET_PATTERNS = {"os.environ", "os.getenv", "dotenv"}
SECRET_FILE_PATTERNS = {".env", "credentials", "secrets", ".secret", ".credentials"}

DYNAMIC_EXEC_BUILTINS = {"eval", "exec", "compile", "__import__"}

OBFUSCATION_ENCODINGS = {"rot_13", "rot13", "hex_codec", "unicode_escape", "raw_unicode_escape"}

FILE_SYSTEM_FUNCS: Dict[str, Set[str]] = {
    "os": {"makedirs", "mkdir", "remove", "unlink", "rename", "listdir", "scandir", "walk", "path"},
    "shutil": {"copy", "copy2", "copytree"},
    "glob": {"glob", "iglob"},
    "pathlib": {"Path"},
}
FILE_SYSTEM_MODULES = {"pathlib", "glob"}


class PythonAnalyzer(BaseAnalyzer):

    def analyze(self, file_path: str) -> List[Capability]:
        try:
            with open(file_path, "r", encoding="utf-8", errors="replace") as f:
                source = f.read()
        except (OSError, IOError):
            return []

        lines = source.splitlines()

        try:
            tree = ast.parse(source, filename=file_path)
        except SyntaxError:
            return self._regex_fallback(file_path, lines)

        capabilities: Dict[str, Capability] = {}
        imported_names: Dict[str, str] = {}

        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for alias in node.names:
                    name = alias.asname if alias.asname else alias.name
                    imported_names[name] = alias.name
                    self._check_import(alias.name, node, file_path, lines, capabilities)

            elif isinstance(node, ast.ImportFrom):
                module = node.module or ""
                for alias in node.names:
                    name = alias.asname if alias.asname else alias.name
                    imported_names[name] = f"{module}.{alias.name}"
                    self._check_import(module, node, file_path, lines, capabilities)
                    self._check_import(f"{module}.{alias.name}", node, file_path, lines, capabilities)

            elif isinstance(node, ast.Call):
                self._check_call(node, file_path, lines, capabilities, imported_names)

            elif isinstance(node, ast.Attribute):
                self._check_attribute(node, file_path, lines, capabilities)

        self._check_obfuscation(file_path, lines, capabilities)

        return list(capabilities.values())

    def _line_snippet(self, lines: List[str], lineno: int) -> str:
        if 1 <= lineno <= len(lines):
            return lines[lineno - 1].rstrip()
        return ""

    def _add_evidence(
        self,
        capabilities: Dict[str, Capability],
        category: str,
        severity: str,
        description: str,
        file_path: str,
        lineno: int,
        snippet: str,
        pattern_name: str,
    ) -> None:
        evidence = Evidence(
            file=file_path,
            line=lineno,
            code_snippet=snippet,
            pattern_name=pattern_name,
        )
        if category not in capabilities:
            capabilities[category] = Capability(
                category=category,
                description=description,
                evidence=[evidence],
                severity=severity,
            )
        else:
            cap = capabilities[category]
            cap.evidence.append(evidence)
            severity_order = {"Low": 0, "Medium": 1, "High": 2, "Critical": 3}
            if severity_order.get(severity, 0) > severity_order.get(cap.severity, 0):
                cap.severity = severity

    def _check_import(
        self,
        module_name: str,
        node: ast.AST,
        file_path: str,
        lines: List[str],
        capabilities: Dict[str, Capability],
    ) -> None:
        lineno = getattr(node, "lineno", 0)
        snippet = self._line_snippet(lines, lineno)
        top_module = module_name.split(".")[0]

        if top_module in SYSTEM_EXEC_MODULES:
            self._add_evidence(
                capabilities, "system_execution", "Critical",
                "System command execution capability detected",
                file_path, lineno, snippet, f"import_{top_module}",
            )

        if top_module in NETWORK_MODULES:
            self._add_evidence(
                capabilities, "network_access", "High",
                "Network access capability detected",
                file_path, lineno, snippet, f"import_{top_module}",
            )

        if "dotenv" in module_name:
            self._add_evidence(
                capabilities, "secret_access", "High",
                "Secret/environment variable access detected",
                file_path, lineno, snippet, "import_dotenv",
            )

        if top_module in FILE_SYSTEM_MODULES:
            self._add_evidence(
                capabilities, "file_system_access", "Medium",
                "File system access capability detected",
                file_path, lineno, snippet, f"import_{top_module}",
            )

    def _check_call(
        self,
        node: ast.Call,
        file_path: str,
        lines: List[str],
        capabilities: Dict[str, Capability],
        imported_names: Dict[str, str],
    ) -> None:
        lineno = getattr(node, "lineno", 0)
        snippet = self._line_snippet(lines, lineno)

        func_name = self._get_call_name(node)
        if not func_name:
            return

        parts = func_name.split(".")

        if len(parts) == 2:
            mod, fn = parts
            resolved_mod = imported_names.get(mod, mod)
            resolved_top = resolved_mod.split(".")[0]

            if resolved_top in SYSTEM_EXEC_FUNCS or mod in SYSTEM_EXEC_FUNCS:
                target_fns = SYSTEM_EXEC_FUNCS.get(resolved_top, set()) | SYSTEM_EXEC_FUNCS.get(mod, set())
                if fn in target_fns:
                    self._add_evidence(
                        capabilities, "system_execution", "Critical",
                        "System command execution capability detected",
                        file_path, lineno, snippet, f"{mod}.{fn}",
                    )

            if resolved_top in NETWORK_FUNCS or mod in NETWORK_FUNCS:
                target_fns = NETWORK_FUNCS.get(resolved_top, set()) | NETWORK_FUNCS.get(mod, set()) | NETWORK_FUNCS.get(resolved_mod, set())
                if fn in target_fns:
                    self._add_evidence(
                        capabilities, "network_access", "High",
                        "Network access capability detected",
                        file_path, lineno, snippet, f"{mod}.{fn}",
                    )

            if mod == "os" and fn == "getenv":
                self._add_evidence(
                    capabilities, "secret_access", "High",
                    "Secret/environment variable access detected",
                    file_path, lineno, snippet, "os.getenv",
                )

            if mod == "os" and fn == "environ":
                self._add_evidence(
                    capabilities, "secret_access", "High",
                    "Secret/environment variable access detected",
                    file_path, lineno, snippet, "os.environ",
                )

            if mod == "importlib" and fn == "import_module":
                self._add_evidence(
                    capabilities, "dynamic_code", "Critical",
                    "Dynamic code execution capability detected",
                    file_path, lineno, snippet, "importlib.import_module",
                )

            if mod == "codecs" and fn == "decode":
                self._add_evidence(
                    capabilities, "obfuscation", "Critical",
                    "Code obfuscation patterns detected",
                    file_path, lineno, snippet, "codecs.decode",
                )

            if resolved_top in FILE_SYSTEM_FUNCS or mod in FILE_SYSTEM_FUNCS:
                target_fns = FILE_SYSTEM_FUNCS.get(resolved_top, set()) | FILE_SYSTEM_FUNCS.get(mod, set())
                if fn in target_fns:
                    self._add_evidence(
                        capabilities, "file_system_access", "Medium",
                        "File system access capability detected",
                        file_path, lineno, snippet, f"{mod}.{fn}",
                    )

        if len(parts) == 1:
            fn = parts[0]
            if fn in DYNAMIC_EXEC_BUILTINS:
                self._add_evidence(
                    capabilities, "dynamic_code", "Critical",
                    "Dynamic code execution capability detected",
                    file_path, lineno, snippet, f"builtin_{fn}",
                )

            if fn == "getattr" and len(node.args) >= 2:
                second_arg = node.args[1]
                if not isinstance(second_arg, ast.Constant):
                    self._add_evidence(
                        capabilities, "dynamic_code", "Critical",
                        "Dynamic code execution capability detected",
                        file_path, lineno, snippet, "dynamic_getattr",
                    )

            if fn == "open":
                self._check_file_open(node, file_path, lineno, lines, capabilities)

    def _check_attribute(
        self,
        node: ast.Attribute,
        file_path: str,
        lines: List[str],
        capabilities: Dict[str, Capability],
    ) -> None:
        lineno = getattr(node, "lineno", 0)
        snippet = self._line_snippet(lines, lineno)

        if isinstance(node.value, ast.Attribute):
            if isinstance(node.value.value, ast.Name) and node.value.value.id == "os":
                if node.value.attr == "environ":
                    self._add_evidence(
                        capabilities, "secret_access", "High",
                        "Secret/environment variable access detected",
                        file_path, lineno, snippet, "os.environ_access",
                    )
        elif isinstance(node.value, ast.Name):
            if node.value.id == "os" and node.attr == "environ":
                self._add_evidence(
                    capabilities, "secret_access", "High",
                    "Secret/environment variable access detected",
                    file_path, lineno, snippet, "os.environ",
                )

    def _check_file_open(
        self,
        node: ast.Call,
        file_path: str,
        lineno: int,
        lines: List[str],
        capabilities: Dict[str, Capability],
    ) -> None:
        snippet = self._line_snippet(lines, lineno)
        matched_secret = False
        if node.args:
            first_arg = node.args[0]
            if isinstance(first_arg, ast.Constant) and isinstance(first_arg.value, str):
                fname = first_arg.value.lower()
                for pat in SECRET_FILE_PATTERNS:
                    if pat in fname:
                        self._add_evidence(
                            capabilities, "secret_access", "High",
                            "Secret/environment variable access detected",
                            file_path, lineno, snippet, f"file_read_{pat}",
                        )
                        matched_secret = True
                        break

        if not matched_secret:
            self._add_evidence(
                capabilities, "file_system_access", "Medium",
                "File system access capability detected",
                file_path, lineno, snippet, "file_open",
            )

    def _check_obfuscation(
        self,
        file_path: str,
        lines: List[str],
        capabilities: Dict[str, Capability],
    ) -> None:
        for i, line in enumerate(lines, 1):
            stripped = line.strip()

            if re.search(r'b(?:ase)?64[._]?(?:b64)?decode\s*\(', stripped) and re.search(r'\b(eval|exec)\s*\(', stripped):
                self._add_evidence(
                    capabilities, "obfuscation", "Critical",
                    "Code obfuscation patterns detected",
                    file_path, i, stripped, "base64_exec",
                )

            if re.search(r'bytes\.fromhex\s*\(', stripped) and re.search(r'\b(eval|exec)\s*\(', stripped):
                self._add_evidence(
                    capabilities, "obfuscation", "Critical",
                    "Code obfuscation patterns detected",
                    file_path, i, stripped, "hex_exec",
                )

            if re.search(r'\[::\s*-1\s*\]', stripped) and re.search(r'\b(eval|exec)\s*\(', stripped):
                self._add_evidence(
                    capabilities, "obfuscation", "Critical",
                    "Code obfuscation patterns detected",
                    file_path, i, stripped, "reverse_exec",
                )

            if len(stripped) > 500 and re.search(r'^[A-Za-z0-9+/=]{500,}', stripped):
                self._add_evidence(
                    capabilities, "obfuscation", "Critical",
                    "Code obfuscation patterns detected",
                    file_path, i, stripped[:200] + "...[truncated]",
                    "long_encoded_string",
                )

            if re.search(r'codecs\.decode\s*\(', stripped):
                for enc in OBFUSCATION_ENCODINGS:
                    if enc in stripped.lower():
                        self._add_evidence(
                            capabilities, "obfuscation", "Critical",
                            "Code obfuscation patterns detected",
                            file_path, i, stripped, f"codecs_{enc}",
                        )
                        break

    def _get_call_name(self, node: ast.Call) -> str | None:
        if isinstance(node.func, ast.Name):
            return node.func.id
        if isinstance(node.func, ast.Attribute):
            value = node.func.value
            if isinstance(value, ast.Name):
                return f"{value.id}.{node.func.attr}"
            if isinstance(value, ast.Attribute) and isinstance(value.value, ast.Name):
                return f"{value.value.id}.{value.attr}.{node.func.attr}"
        return None

    def _regex_fallback(self, file_path: str, lines: List[str]) -> List[Capability]:
        capabilities: Dict[str, Capability] = {}

        patterns: List[Tuple[str, str, str, str]] = [
            (r'\bsubprocess\.\w+\s*\(', "system_execution", "Critical", "subprocess_call"),
            (r'\bos\.system\s*\(', "system_execution", "Critical", "os.system"),
            (r'\bos\.popen\s*\(', "system_execution", "Critical", "os.popen"),
            (r'\bos\.exec\w*\s*\(', "system_execution", "Critical", "os.exec"),
            (r'\bshutil\.rmtree\s*\(', "system_execution", "Critical", "shutil.rmtree"),
            (r'\bctypes\b', "system_execution", "Critical", "ctypes_usage"),
            (r'\brequests\.\w+\s*\(', "network_access", "High", "requests_call"),
            (r'\burllib\.\w+', "network_access", "High", "urllib_usage"),
            (r'\bhttpx\.\w+', "network_access", "High", "httpx_usage"),
            (r'\baiohttp\.\w+', "network_access", "High", "aiohttp_usage"),
            (r'\bsocket\.\w+', "network_access", "High", "socket_usage"),
            (r'\bos\.environ\b', "secret_access", "High", "os.environ"),
            (r'\bos\.getenv\s*\(', "secret_access", "High", "os.getenv"),
            (r'\bdotenv\b', "secret_access", "High", "dotenv_usage"),
            (r'\beval\s*\(', "dynamic_code", "Critical", "eval_call"),
            (r'\bexec\s*\(', "dynamic_code", "Critical", "exec_call"),
            (r'\bcompile\s*\(', "dynamic_code", "Critical", "compile_call"),
            (r'\b__import__\s*\(', "dynamic_code", "Critical", "__import___call"),
            (r'\bopen\s*\(', "file_system_access", "Medium", "file_open"),
            (r'\bpathlib\.\w+', "file_system_access", "Medium", "pathlib_usage"),
            (r'\bglob\.\w+', "file_system_access", "Medium", "glob_usage"),
            (r'\bos\.(makedirs|mkdir|remove|unlink|rename|listdir|scandir)\s*\(', "file_system_access", "Medium", "os_fs"),
            (r'\bshutil\.(copy|copy2|copytree)\s*\(', "file_system_access", "Medium", "shutil_copy"),
        ]

        category_descriptions = {
            "system_execution": "System command execution capability detected",
            "network_access": "Network access capability detected",
            "secret_access": "Secret/environment variable access detected",
            "dynamic_code": "Dynamic code execution capability detected",
            "obfuscation": "Code obfuscation patterns detected",
            "file_system_access": "File system access capability detected",
        }

        for i, line in enumerate(lines, 1):
            for pattern, category, severity, name in patterns:
                if re.search(pattern, line):
                    self._add_evidence(
                        capabilities, category, severity,
                        category_descriptions[category],
                        file_path, i, line.rstrip(), name,
                    )

        self._check_obfuscation(file_path, lines, capabilities)

        return list(capabilities.values())
