"""Checker for SQLAlchemy sessions used unsafely across thread boundaries.

SQLAlchemy sessions are NOT thread-safe. Using a session from multiple
threads (e.g., via ThreadPoolExecutor workers) causes race conditions
that corrupt transaction state, leading to errors like:
  "This session is in 'prepared' state; no further SQL can be emitted"

This checker detects two patterns:

1. DIRECT: A function submitted to ThreadPoolExecutor calls
   session.commit(), session.execute(), session.add(), etc.

2. INDIRECT: A function submitted to ThreadPoolExecutor calls
   self._progress_callback (or similar callback) that may write
   to the DB session.
"""

import ast
import re
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union

from stablestack.checkers.base import BaseChecker
from stablestack.models import Finding, Severity, Tier

# Session methods that are NOT thread-safe
SESSION_WRITE_METHODS = {
    "commit",
    "execute",
    "add",
    "add_all",
    "delete",
    "flush",
    "rollback",
}

# Attribute names that typically hold sessions
SESSION_ATTR_NAMES = {"db_session", "session", "db", "_session"}

# Callback attribute names that might write to DB
CALLBACK_ATTR_NAMES = {"_progress_callback", "progress_callback"}


class ThreadSessionSafetyChecker(BaseChecker):
    """Detects SQLAlchemy sessions used across ThreadPoolExecutor boundaries."""

    rule_id = "SESS002"
    rule_name = "thread-session-safety"
    description = "SQLAlchemy session used across ThreadPoolExecutor boundary."
    why = (
        "SQLAlchemy sessions are NOT thread-safe. When a function submitted to "
        "ThreadPoolExecutor calls session.commit(), execute(), or similar methods, "
        "it races with the main thread's session usage, corrupting transaction state. "
        "This causes 'prepared state' errors, lost data, and cascading failures. "
        "Worker threads should only do I/O; all DB writes belong on the main thread."
    )
    severity = Severity.ERROR
    tier = Tier.CRITICAL
    file_extensions = {".py"}

    def check_file(self, path: Path, content: str) -> List[Finding]:
        findings: List[Finding] = []

        # Quick pre-filter: only parse files that use ThreadPoolExecutor
        if "ThreadPoolExecutor" not in content:
            return findings

        try:
            tree = ast.parse(content, filename=str(path))
        except SyntaxError:
            return findings

        lines = content.split("\n")

        # Find all methods/functions in the file
        method_bodies = self._extract_method_bodies(tree)

        # Find ThreadPoolExecutor usage and the submitted callables
        for node in ast.walk(tree):
            if not isinstance(node, ast.With):
                continue

            executor_var = self._get_executor_var(node)
            if not executor_var:
                continue

            submitted_callables = self._find_submitted_callables(node, executor_var)

            for callable_name, submit_line in submitted_callables:
                # Check direct session access
                session_accesses = self._find_session_access_in_callable(
                    callable_name, method_bodies
                )
                for _access_type, access_line, detail in session_accesses:
                    report_line = access_line or submit_line
                    findings.append(
                        self.create_finding(
                            path=path,
                            line=report_line,
                            column=1,
                            code=(
                                lines[report_line - 1].strip()[:80]
                                if report_line <= len(lines)
                                else ""
                            ),
                            message=(
                                f"Function '{callable_name}' is submitted to "
                                f"ThreadPoolExecutor (line {submit_line}) but "
                                f"{detail}. SQLAlchemy sessions are NOT thread-safe."
                            ),
                            fix_hint=(
                                "Do all DB writes from the main thread. Worker "
                                "threads should only perform I/O and return results."
                            ),
                        )
                    )

                # Check indirect callback access
                callback_findings = self._check_callback_in_callable(
                    callable_name, method_bodies, tree, submit_line, path, lines
                )
                findings.extend(callback_findings)

        return findings

    def _get_executor_var(self, with_node: ast.With) -> Optional[str]:
        """Extract the variable name from `with ThreadPoolExecutor() as var`."""
        for item in with_node.items:
            ctx = item.context_expr
            if isinstance(ctx, ast.Call):
                func = ctx.func
                if isinstance(func, ast.Name) and func.id == "ThreadPoolExecutor":
                    if item.optional_vars and isinstance(item.optional_vars, ast.Name):
                        return item.optional_vars.id
                if (
                    isinstance(func, ast.Attribute)
                    and func.attr == "ThreadPoolExecutor"
                ):
                    if item.optional_vars and isinstance(item.optional_vars, ast.Name):
                        return item.optional_vars.id
        return None

    def _find_submitted_callables(
        self, with_node: ast.With, executor_var: str
    ) -> List[Tuple[str, int]]:
        """Find all executor.submit(callable, ...) calls and return (name, line)."""
        results = []
        for node in ast.walk(with_node):
            if not isinstance(node, ast.Call):
                continue
            func = node.func
            if (
                isinstance(func, ast.Attribute)
                and func.attr == "submit"
                and isinstance(func.value, ast.Name)
                and func.value.id == executor_var
            ):
                if node.args:
                    callable_arg = node.args[0]
                    name = self._get_callable_name(callable_arg)
                    if name:
                        results.append((name, node.lineno))
        return results

    def _get_callable_name(self, node: ast.expr) -> Optional[str]:
        """Extract the function/method name from a submit() argument."""
        if isinstance(node, ast.Name):
            return node.id
        if isinstance(node, ast.Attribute):
            return node.attr
        return None

    def _extract_method_bodies(self, tree: ast.Module) -> Dict[str, ast.AST]:
        """Map function/method names to their AST nodes."""
        methods: Dict[str, ast.AST] = {}
        for node in ast.walk(tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                methods[node.name] = node
        return methods

    def _find_session_access_in_callable(
        self,
        callable_name: str,
        method_bodies: Dict[str, ast.AST],
    ) -> List[Tuple[str, Optional[int], str]]:
        """Check if a callable directly accesses session write methods.

        Returns list of (access_type, line_number, description).
        """
        results: List[Tuple[str, Optional[int], str]] = []

        func_node = method_bodies.get(callable_name)
        if not func_node:
            return results

        for node in ast.walk(func_node):
            if not isinstance(node, ast.Call):
                continue
            func = node.func
            if isinstance(func, ast.Attribute) and func.attr in SESSION_WRITE_METHODS:
                if isinstance(func.value, ast.Attribute):
                    attr_name = func.value.attr
                    if attr_name in SESSION_ATTR_NAMES:
                        results.append(
                            (
                                "direct",
                                node.lineno,
                                f"calls self.{attr_name}.{func.attr}()",
                            )
                        )
                elif isinstance(func.value, ast.Name):
                    var_name = func.value.id
                    if var_name in SESSION_ATTR_NAMES:
                        results.append(
                            ("direct", node.lineno, f"calls {var_name}.{func.attr}()")
                        )

        return results

    def _check_callback_in_callable(
        self,
        callable_name: str,
        method_bodies: Dict[str, ast.AST],
        tree: ast.Module,
        submit_line: int,
        file_path: Path,
        lines: List[str],
    ) -> List[Finding]:
        """Check if a submitted callable calls self._progress_callback."""
        findings: List[Finding] = []

        func_node = method_bodies.get(callable_name)
        if not func_node:
            return findings

        calls_callback = False
        callback_line = None

        for node in ast.walk(func_node):
            if not isinstance(node, ast.Call):
                continue
            func = node.func
            if (
                isinstance(func, ast.Attribute)
                and func.attr in CALLBACK_ATTR_NAMES
                and isinstance(func.value, ast.Name)
                and func.value.id == "self"
            ):
                calls_callback = True
                callback_line = node.lineno
                break

        if not calls_callback:
            # Also check methods called by this callable
            called_methods = self._find_self_method_calls(func_node)
            for method_name in called_methods:
                sub_func = method_bodies.get(method_name)
                if sub_func:
                    for node in ast.walk(sub_func):
                        if not isinstance(node, ast.Call):
                            continue
                        func = node.func
                        if (
                            isinstance(func, ast.Attribute)
                            and func.attr in CALLBACK_ATTR_NAMES
                            and isinstance(func.value, ast.Name)
                            and func.value.id == "self"
                        ):
                            calls_callback = True
                            callback_line = node.lineno
                            break
                    if calls_callback:
                        break

        if not calls_callback:
            return findings

        enclosing = self._find_enclosing_method(tree, submit_line)
        if enclosing:
            has_guard = self._has_callback_guard(enclosing, lines)
            if not has_guard:
                report_line = callback_line or submit_line
                findings.append(
                    self.create_finding(
                        path=file_path,
                        line=report_line,
                        column=1,
                        code=(
                            lines[report_line - 1].strip()[:80]
                            if report_line <= len(lines)
                            else ""
                        ),
                        message=(
                            f"Function '{callable_name}' is submitted to "
                            f"ThreadPoolExecutor (line {submit_line}) and calls "
                            f"self._progress_callback, which may write to the "
                            f"DB session."
                        ),
                        fix_hint=(
                            "Replace _progress_callback with a log-only version "
                            "before entering the thread pool, and restore after."
                        ),
                    )
                )

        return findings

    def _find_self_method_calls(self, func_node: ast.AST) -> List[str]:
        """Find all self.method_name() calls in a function."""
        methods = []
        for node in ast.walk(func_node):
            if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
                if (
                    isinstance(node.func.value, ast.Name)
                    and node.func.value.id == "self"
                ):
                    methods.append(node.func.attr)
        return methods

    def _find_enclosing_method(
        self, tree: ast.Module, line: int
    ) -> Optional[Union[ast.FunctionDef, ast.AsyncFunctionDef]]:
        """Find the method that contains the given line number."""
        best = None
        for node in ast.walk(tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                if hasattr(node, "end_lineno") and node.end_lineno:
                    if node.lineno <= line <= node.end_lineno:
                        if best is None or node.lineno > best.lineno:
                            best = node
        return best

    def _has_callback_guard(
        self, method: Union[ast.FunctionDef, ast.AsyncFunctionDef], lines: List[str]
    ) -> bool:
        """Check if the method guards _progress_callback before ThreadPoolExecutor.

        Looks for patterns like:
            saved_callback = self._progress_callback
            self._progress_callback = lambda msg: logging.info(msg)
        before the ThreadPoolExecutor block.
        """
        method_source = "\n".join(
            lines[method.lineno - 1 : method.end_lineno]
            if hasattr(method, "end_lineno") and method.end_lineno
            else lines[method.lineno - 1 : method.lineno + 100]
        )

        guard_pattern = re.compile(
            r"saved_callback\s*=\s*self\._progress_callback.*\n"
            r".*self\._progress_callback\s*=",
            re.MULTILINE,
        )
        return bool(guard_pattern.search(method_source))
