from __future__ import annotations

import re
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Set, Tuple

from .config import CAPABILITY_FUNCTIONS, DEFAULT_SANITIZERS, DEFAULT_SINKS, DEFAULT_SOURCES, NONCE_FUNCTIONS
from .filesystem import path_is_ignored, safe_read_text
from .models import AttackSurfaceEntry, FindingDetail, Severity
from .php_ast import AstTaintAnalyzer, parse_php_ast


class FunctionScopeTracker:
    """
    Tracks function definitions and their line ranges using brace counting.
    """
    def __init__(self, content: str):
        self.lines = content.splitlines()
        self.functions: Dict[str, Tuple[int, int]] = {}  # name -> (start_line, end_line)
        self._parse()

    def _parse(self):
        # Very basic parser: find "function name(...)" and count braces
        # This is not robust against complex string interpolation or comments with braces,
        # but better than nothing for a heuristic tool.
        current_func = None
        brace_balance = 0
        start_line = 0

        for idx, line in enumerate(self.lines, start=1):
            # Remove comments for brace counting (simple approximation)
            clean_line = re.sub(r"//.*|/\*.*?\*/", "", line)
            
            # Detect function start
            if current_func is None:
                m = re.search(r"function\s+(\w+)\s*\(", clean_line)
                if m:
                    current_func = m.group(1)
                    start_line = idx
                    brace_balance = 0
            
            if current_func:
                brace_balance += clean_line.count("{")
                brace_balance -= clean_line.count("}")
                
                if brace_balance == 0 and "}" in clean_line:
                    # Function ended
                    self.functions[current_func] = (start_line, idx)
                    current_func = None

    def get_function_body(self, func_name: str) -> Optional[str]:
        if func_name not in self.functions:
            return None
        start, end = self.functions[func_name]
        # Return lines inclusive
        return "\n".join(self.lines[start-1:end])


class PHPAnalyzer:
    """
    Lightweight heuristic analyzer for PHP files.
    This is not a full AST engine but approximates dataflow for common WordPress patterns.
    """

    def __init__(
        self,
        wp_version: str,
        php_version: str,
        strict: bool,
        include_suspicious: bool,
        ignore_paths: List[str] | None = None,
        allow_sinks: set[str] | None = None,
        deny_sinks: dict[str, str] | None = None,
        severity_threshold: str | None = None,
    ):
        self.wp_version = wp_version
        self.php_version = php_version
        self.strict = strict
        self.include_suspicious = include_suspicious
        self.ignore_paths = ignore_paths or []
        self.allow_sinks = allow_sinks or set()
        self.severity_threshold = (severity_threshold or "").lower() or None
        self.sinks = dict(DEFAULT_SINKS)
        for name, desc in (deny_sinks or {}).items():
            self.sinks[name] = desc or "Custom sink"
        for name in list(self.sinks.keys()):
            if name in self.allow_sinks:
                self.sinks.pop(name, None)
        self.ast_analyzer = AstTaintAnalyzer(self.sinks, set(DEFAULT_SANITIZERS))

    def analyze(self, root: Path) -> Tuple[List[FindingDetail], List[AttackSurfaceEntry]]:
        findings: List[FindingDetail] = []
        surface: List[AttackSurfaceEntry] = []
        for php in (p for p in root.rglob("*.php") if p.is_file()):
            if path_is_ignored(php, self.ignore_paths):
                continue
            content = safe_read_text(php, limit=300_000)
            file_findings, entries = self._analyze_php_file(php, content)
            findings.extend(file_findings)
            surface.extend(entries)
            # AST pass (if ext-ast/php available)
            ast_tree = parse_php_ast(php)
            if ast_tree:
                ast_result = self.ast_analyzer.analyze(ast_tree, php)
                findings.extend(ast_result.findings)
        findings = self._apply_threshold(findings)
        return findings, surface

    def _analyze_php_file(self, path: Path, content: str) -> Tuple[List[FindingDetail], List[AttackSurfaceEntry]]:
        findings: List[FindingDetail] = []
        surface: List[AttackSurfaceEntry] = []

        tainted_vars = set()
        sanitized_vars = set()
        
        scope_tracker = FunctionScopeTracker(content)
        ajax_actions: Dict[str, str] = {} # hook_name -> callback_function

        lines = content.splitlines()
        for idx, line in enumerate(lines, start=1):
            stripped = line.strip()
            
            # Attack surface & Hook detection
            entries = self._extract_hooks(line, path, idx)
            surface.extend(entries)
            
            # Track AJAX hooks for later verification
            for entry in entries:
                if entry.hook_type == "action" and entry.hook.startswith("wp_ajax_"):
                    ajax_actions[entry.hook] = entry.callback

            # Source tracking
            for src in DEFAULT_SOURCES:
                if f"${src}" in line or re.search(rf"{re.escape(src)}\s*\(", line):
                    m = re.search(r"\$(\w+)\s*=", line)
                    if m:
                        tainted_vars.add(m.group(1))

            # Sanitizers
            for sanitizer in DEFAULT_SANITIZERS:
                if sanitizer in line:
                    m = re.search(r"\$(\w+)\s*=\s*[^;]*" + re.escape(sanitizer) + r"\s*\(", line)
                    if m:
                        sanitized_vars.add(m.group(1))
                        # Best effort: consider assigned var untainted
                        if m.group(1) in tainted_vars:
                            tainted_vars.discard(m.group(1))

            # Direct echo/print with tainted
            for var in list(tainted_vars):
                if re.search(rf"(echo|print|printf)\s+[^;]*\${var}\b", line):
                    findings.append(
                        FindingDetail(
                            category="xss",
                            message="User input echoed without escaping",
                            severity=Severity.HIGH,
                            file=str(path),
                            line=idx,
                            evidence=stripped,
                            remediation="Escape output with esc_html/esc_attr before echoing.",
                            requires_auth="depends",
                        )
                    )

            # Sinks
            for sink, desc in self.sinks.items():
                if re.search(rf"\b{re.escape(sink)}\s*\(", line):
                    # find variables in call
                    vars_in_line = re.findall(r"\$(\w+)", line)
                    for var in vars_in_line:
                        if var in tainted_vars and var not in sanitized_vars:
                            severity = Severity.HIGH if self.strict else Severity.MEDIUM
                            findings.append(
                                FindingDetail(
                                    category="source-to-sink",
                                    message=f"Tainted input flows into {sink} ({desc})",
                                    severity=severity,
                                    file=str(path),
                                    line=idx,
                                    evidence=stripped,
                                    remediation="Sanitize or validate input before calling this sink.",
                                )
                            )
                        elif self.include_suspicious:
                            findings.append(
                                FindingDetail(
                                    category="suspicious-sink",
                                    message=f"{sink} called; manual review recommended",
                                    severity=Severity.LOW,
                                    file=str(path),
                                    line=idx,
                                    evidence=stripped,
                                )
                            )

            # SQL injection heuristic
            if "$wpdb" in line and re.search(r"->(query|get_results|get_row|get_var|prepare)\s*\(", line):
                if "prepare" in line:
                    continue
                if any(var in tainted_vars for var in re.findall(r"\$(\w+)", line)):
                    findings.append(
                        FindingDetail(
                            category="sql-injection",
                            message="Unprepared $wpdb call uses user input",
                            severity=Severity.CRITICAL if self.strict else Severity.HIGH,
                            file=str(path),
                            line=idx,
                            evidence=stripped,
                            remediation="Use $wpdb->prepare with placeholders before executing queries.",
                        )
                    )

            # unserialize with tainted input
            if "unserialize" in line:
                if any(var in tainted_vars for var in re.findall(r"\$(\w+)", line)):
                    findings.append(
                        FindingDetail(
                            category="object-injection",
                            message="unserialize on user input (object injection risk)",
                            severity=Severity.CRITICAL if self.strict else Severity.HIGH,
                            file=str(path),
                            line=idx,
                            evidence=stripped,
                            remediation="Avoid unserialize on untrusted data. Use json_decode or verify type.",
                        )
                    )

            # File upload checks
            if re.search(r"move_uploaded_file", line) and not re.search(r"(mime|type|extension)", content, re.IGNORECASE):
                findings.append(
                    FindingDetail(
                        category="file-upload",
                        message="File upload without extension/MIME validation detected",
                        severity=Severity.HIGH,
                        file=str(path),
                        line=idx,
                        evidence=stripped,
                        remediation="Validate MIME/extension and use wp_check_filetype_and_ext.",
                    )
                )

        # Post-file checks: Verify AJAX handlers have nonce/capability checks IN THEIR BODY
        for hook, callback in ajax_actions.items():
            func_body = scope_tracker.get_function_body(callback)
            if not func_body:
                # Could not find function body (maybe in another file or dynamic), skip for now to avoid FP
                continue
            
            nonce_present = bool(re.search(r"(check_(ajax_)?referer|wp_verify_nonce)", func_body))
            capability_present = bool(re.search(r"(current_user_can|user_can|is_user_logged_in)", func_body))
            
            if not nonce_present:
                findings.append(
                    FindingDetail(
                        category="csrf",
                        message=f"AJAX handler '{callback}' missing nonce verification",
                        severity=Severity.HIGH if self.strict else Severity.MEDIUM,
                        file=str(path),
                        evidence=f"Hook: {hook}, Callback: {callback}",
                        remediation="Call check_ajax_referer or wp_verify_nonce inside the handler.",
                    )
                )
            
            if not capability_present:
                # Only flag if it's a privileged ajax action (not nopriv)
                if "nopriv" not in hook:
                    findings.append(
                        FindingDetail(
                            category="authz",
                            message=f"AJAX handler '{callback}' missing capability check",
                            severity=Severity.HIGH if self.strict else Severity.MEDIUM,
                            file=str(path),
                            evidence=f"Hook: {hook}, Callback: {callback}",
                            remediation="Validate user capabilities (current_user_can) inside the handler.",
                        )
                    )

        return findings, surface

    def _apply_threshold(self, findings: List[FindingDetail]) -> List[FindingDetail]:
        if not self.severity_threshold:
            return findings
        order = {
            Severity.CRITICAL: 5,
            Severity.HIGH: 4,
            Severity.MEDIUM: 3,
            Severity.LOW: 2,
            Severity.INFO: 1,
            Severity.NOTE: 0,
        }
        threshold_val = order.get(Severity(self.severity_threshold), 0) if self.severity_threshold in Severity._value2member_map_ else 0
        filtered: List[FindingDetail] = []
        for f in findings:
            if order.get(f.severity, 0) >= threshold_val:
                filtered.append(f)
            elif self.include_suspicious:
                filtered.append(f)
        return filtered

    def _extract_hooks(self, line: str, path: Path, idx: int) -> List[AttackSurfaceEntry]:
        entries: List[AttackSurfaceEntry] = []
        hook_match = re.search(r"add_(action|filter)\s*\(\s*['\"]([^'\"]+)['\"]\s*,\s*['\"]?([\w:]+)", line)
        if hook_match:
            hook_type = hook_match.group(1)
            hook = hook_match.group(2)
            callback = hook_match.group(3)
            priority = None
            arg_match = re.search(r"add_(action|filter)\s*\(.*?,.*?,\s*(\d+)", line)
            if arg_match:
                try:
                    priority = int(arg_match.group(2))
                except ValueError:
                    priority = None
            entries.append(
                AttackSurfaceEntry(
                    hook_type=hook_type,
                    hook=hook,
                    callback=callback,
                    priority=priority,
                    file=str(path),
                    line=idx,
                )
            )

        rest_match = re.search(r"register_rest_route\s*\(\s*['\"]([^'\"]+)['\"]\s*,\s*['\"]([^'\"]+)['\"]", line)
        if rest_match:
            route = f"/{rest_match.group(1)}{rest_match.group(2)}"
            entries.append(
                AttackSurfaceEntry(
                    hook_type="rest",
                    hook=route,
                    callback="handler",
                    file=str(path),
                    line=idx,
                )
            )
        shortcode_match = re.search(r"add_shortcode\s*\(\s*['\"]([^'\"]+)['\"]\s*,\s*['\"]?([\w:]+)", line)
        if shortcode_match:
            entries.append(
                AttackSurfaceEntry(
                    hook_type="shortcode",
                    hook=shortcode_match.group(1),
                    callback=shortcode_match.group(2),
                    file=str(path),
                    line=idx,
                )
            )
        return entries
