"""SQLite-backed reference graph for Jedi-based call edge indexing."""

from __future__ import annotations

import hashlib
import logging
import os
import sqlite3
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Callable, Iterable

    from jedi.api.classes import Name  # type: ignore[import-untyped]

    from .._model import FunctionSource
    from ._call_graph import CallGraph, IndexResult

log = logging.getLogger(__name__)

_PARALLEL_THRESHOLD = 8

_worker_jedi_project: object | None = None
_worker_project_root_str: str | None = None


def get_qualified_name(
    module_name: str,
    full_qualified_name: str,
) -> str:
    """Strip the module prefix from a fully qualified name."""
    if not full_qualified_name:
        msg = "full_qualified_name cannot be empty"
        raise ValueError(msg)
    if not full_qualified_name.startswith(module_name):
        msg = f"{full_qualified_name} does not start with {module_name}"
        raise ValueError(msg)
    if module_name == full_qualified_name:
        msg = f"{full_qualified_name} is the same as {module_name}"
        raise ValueError(msg)
    return full_qualified_name[len(module_name) + 1 :]


def path_belongs_to_site_packages(file_path: Path) -> bool:
    """Return True if *file_path* is under a site-packages directory."""
    import site  # noqa: PLC0415

    file_path_resolved = file_path.resolve()
    site_packages = [Path(p).resolve() for p in site.getsitepackages()]
    return any(file_path_resolved.is_relative_to(sp) for sp in site_packages)


def belongs_to_function_qualified(
    name: Name,
    qualified_function_name: str,
) -> bool:
    """Return True if *name* is defined inside *qualified_function_name*."""
    try:
        if (
            name.full_name.startswith(name.module_name)
            and get_qualified_name(name.module_name, name.full_name)
            == qualified_function_name
        ):
            return False
        parent = name.parent()
        if parent and parent.type == "function":
            return (
                get_qualified_name(parent.module_name, parent.full_name)
                == qualified_function_name
            )
    except (ValueError, AttributeError):
        return False
    return False


def _resolve_definitions(ref: Name) -> list[Name]:
    """Resolve a Jedi reference to its definitions."""
    try:
        inferred = ref.infer()
        valid = [
            d for d in inferred if d.type in ("function", "class", "statement")
        ]
        if valid:
            return valid
    except Exception:  # noqa: BLE001, S110
        pass
    try:
        result: list[Name] = ref.goto(
            follow_imports=True, follow_builtin_imports=False
        )
        return result  # noqa: TRY300
    except Exception:  # noqa: BLE001
        return []


def _is_valid_definition(  # noqa: PLR0911
    definition: Name,
    caller_qualified_name: str,
    project_root_str: str,
) -> bool:
    """Return True if *definition* is a valid in-project callee."""
    definition_path = definition.module_path
    if definition_path is None:
        return False
    if not str(definition_path).startswith(project_root_str + os.sep):
        return False
    if path_belongs_to_site_packages(definition_path):
        return False
    if not definition.full_name or not definition.full_name.startswith(
        definition.module_name
    ):
        return False
    if definition.type not in ("function", "class", "statement"):
        return False
    try:
        def_qn = get_qualified_name(
            definition.module_name, definition.full_name
        )
        if def_qn == caller_qualified_name:
            return False
    except ValueError:
        return False
    try:
        if belongs_to_function_qualified(definition, caller_qualified_name):
            return False
    except Exception:  # noqa: BLE001, S110
        pass
    return True


def _get_enclosing_function_qn(ref: Name) -> str | None:
    """Return the qualified name of the function enclosing *ref*."""
    try:
        parent = ref.parent()
        if parent is None or parent.type != "function":
            return None
        if not parent.full_name or not parent.full_name.startswith(
            parent.module_name
        ):
            return None
        return get_qualified_name(parent.module_name, parent.full_name)
    except (ValueError, AttributeError):
        return None


def _analyze_file(  # noqa: C901, PLR0912
    file_path: Path,
    jedi_project: object,
    project_root_str: str,
) -> tuple[set[tuple[str, ...]], bool]:
    """Analyze a Python file with Jedi and return call edges."""
    import jedi  # type: ignore[import-untyped]  # noqa: PLC0415

    resolved = str(file_path.resolve())
    try:
        script = jedi.Script(path=file_path, project=jedi_project)
        refs = script.get_names(
            all_scopes=True, definitions=False, references=True
        )
    except Exception:  # noqa: BLE001
        return set(), True
    edges: set[tuple[str, ...]] = set()
    for ref in refs:
        try:
            caller_qn = _get_enclosing_function_qn(ref)
            if caller_qn is None:
                continue
            definitions = _resolve_definitions(ref)
            if not definitions:
                continue
            definition = definitions[0]
            definition_path = definition.module_path
            if definition_path is None:
                continue
            if not _is_valid_definition(
                definition, caller_qn, project_root_str
            ):
                continue
            edge_base = (resolved, caller_qn, str(definition_path))
            if definition.type == "function":
                callee_qn = get_qualified_name(
                    definition.module_name, definition.full_name
                )
                if len(callee_qn.split(".")) > 2:  # noqa: PLR2004
                    continue
                edges.add(
                    (
                        *edge_base,
                        callee_qn,
                        definition.full_name,
                        definition.name,
                        definition.type,
                        definition.get_line_code(),
                    )
                )
            elif definition.type == "class":
                init_qn = get_qualified_name(
                    definition.module_name,
                    f"{definition.full_name}.__init__",
                )
                if len(init_qn.split(".")) > 2:  # noqa: PLR2004
                    continue
                edges.add(
                    (
                        *edge_base,
                        init_qn,
                        f"{definition.full_name}.__init__",
                        "__init__",
                        definition.type,
                        definition.get_line_code(),
                    )
                )
            elif definition.type == "statement":
                callee_qn = get_qualified_name(
                    definition.module_name, definition.full_name
                )
                if len(callee_qn.split(".")) > 2:  # noqa: PLR2004
                    continue
                edges.add(
                    (
                        *edge_base,
                        callee_qn,
                        definition.full_name,
                        definition.name,
                        definition.type,
                        definition.get_line_code(),
                    )
                )
        except Exception:  # noqa: BLE001, S112
            continue
    return edges, False


def _init_index_worker(project_root: str) -> None:
    """Initialize the Jedi project for a worker process."""
    import jedi  # noqa: PLC0415

    global _worker_jedi_project, _worker_project_root_str  # noqa: PLW0603
    _worker_jedi_project = jedi.Project(path=project_root)
    _worker_project_root_str = project_root


def _index_file_worker(
    args: tuple[str, str],
) -> tuple[str, str, set[tuple[str, ...]], bool]:
    """Worker entry point for parallel file indexing."""
    file_path_str, file_hash = args
    assert _worker_project_root_str is not None  # noqa: S101
    edges, had_error = _analyze_file(
        Path(file_path_str),
        _worker_jedi_project,
        _worker_project_root_str,
    )
    return file_path_str, file_hash, edges, had_error


class ReferenceGraph:
    """SQLite-backed call graph that indexes Python projects via Jedi."""

    SCHEMA_VERSION = 2

    def __init__(
        self,
        project_root: Path,
        language: str = "python",
        db_path: Path | None = None,
    ) -> None:
        """Initialize the graph with a Jedi project and SQLite connection."""
        import jedi  # noqa: PLC0415

        self.project_root = project_root.resolve()
        self.project_root_str = str(self.project_root)
        self.language = language
        self.jedi_project = jedi.Project(path=self.project_root)
        if db_path is None:
            db_path = Path.home() / ".codeflash" / "codeflash_cache.db"
        self.conn = sqlite3.connect(str(db_path))
        self.conn.execute("PRAGMA journal_mode=WAL")
        self.indexed_file_hashes: dict[str, str] = {}
        self._resolved_paths: dict[Path, str] = {}
        self._init_schema()

    def _init_schema(self) -> None:
        """Create or migrate SQLite tables for call edge storage."""
        cur = self.conn.cursor()
        cur.execute(
            "CREATE TABLE IF NOT EXISTS cg_schema_version "
            "(version INTEGER PRIMARY KEY)"
        )
        row = cur.execute(
            "SELECT version FROM cg_schema_version LIMIT 1"
        ).fetchone()
        if row is None:
            cur.execute(
                "INSERT INTO cg_schema_version (version) VALUES (?)",
                (self.SCHEMA_VERSION,),
            )
        elif row[0] != self.SCHEMA_VERSION:
            for table in [
                "cg_call_edges",
                "cg_indexed_files",
                "cg_languages",
                "cg_projects",
                "cg_project_meta",
                "indexed_files",
                "call_edges",
            ]:
                cur.execute(f"DROP TABLE IF EXISTS {table}")
            cur.execute("DELETE FROM cg_schema_version")
            cur.execute(
                "INSERT INTO cg_schema_version (version) VALUES (?)",
                (self.SCHEMA_VERSION,),
            )
        cur.execute("""
            CREATE TABLE IF NOT EXISTS indexed_files (
                project_root TEXT NOT NULL,
                language     TEXT NOT NULL,
                file_path    TEXT NOT NULL,
                file_hash    TEXT NOT NULL,
                PRIMARY KEY (project_root, language, file_path)
            )
        """)
        cur.execute("""
            CREATE TABLE IF NOT EXISTS call_edges (
                project_root                TEXT NOT NULL,
                language                    TEXT NOT NULL,
                caller_file                 TEXT NOT NULL,
                caller_qualified_name       TEXT NOT NULL,
                callee_file                 TEXT NOT NULL,
                callee_qualified_name       TEXT NOT NULL,
                callee_fully_qualified_name TEXT NOT NULL,
                callee_only_function_name   TEXT NOT NULL,
                callee_definition_type      TEXT NOT NULL,
                callee_source_line          TEXT NOT NULL,
                PRIMARY KEY (
                    project_root, language,
                    caller_file, caller_qualified_name,
                    callee_file, callee_qualified_name
                )
            )
        """)
        cur.execute("""
            CREATE INDEX IF NOT EXISTS idx_call_edges_caller
            ON call_edges (
                project_root, language,
                caller_file, caller_qualified_name
            )
        """)
        self.conn.commit()

    def resolve_path(self, file_path: Path) -> str:
        """Return a cached resolved string path for *file_path*."""
        cached = self._resolved_paths.get(file_path)
        if cached is not None:
            return cached
        resolved = str(file_path.resolve())
        self._resolved_paths[file_path] = resolved
        return resolved

    def get_callees(
        self,
        file_path_to_qualified_names: dict[Path, set[str]],
    ) -> tuple[dict[Path, set[FunctionSource]], list[FunctionSource]]:
        """Return callees as FunctionSource from the indexed graph."""
        from ._call_graph import callees_from_graph  # noqa: PLC0415

        graph = self.get_call_graph(
            file_path_to_qualified_names, include_metadata=True
        )
        return callees_from_graph(graph)

    def count_callees_per_function(
        self,
        file_path_to_qualified_names: dict[Path, set[str]],
    ) -> dict[tuple[Path, str], int]:
        """Count the number of callees for each (file, qualified_name)."""
        all_caller_keys: list[tuple[Path, str, str]] = []
        for file_path, qualified_names in file_path_to_qualified_names.items():
            resolved = self.resolve_path(file_path)
            self.ensure_file_indexed(file_path, resolved)
            all_caller_keys.extend(
                (file_path, resolved, qn) for qn in qualified_names
            )
        if not all_caller_keys:
            return {}
        cur = self.conn.cursor()
        cur.execute(
            "CREATE TEMP TABLE IF NOT EXISTS _count_keys "
            "(caller_file TEXT, caller_qualified_name TEXT)"
        )
        cur.execute("DELETE FROM _count_keys")
        cur.executemany(
            "INSERT INTO _count_keys VALUES (?, ?)",
            [(resolved, qn) for _, resolved, qn in all_caller_keys],
        )
        rows = cur.execute(
            """
            SELECT ck.caller_file, ck.caller_qualified_name,
                   COUNT(ce.rowid)
            FROM _count_keys ck
            LEFT JOIN call_edges ce
                ON ce.caller_file = ck.caller_file
                AND ce.caller_qualified_name =
                    ck.caller_qualified_name
                AND ce.project_root = ?
                AND ce.language = ?
            GROUP BY ck.caller_file, ck.caller_qualified_name
            """,
            (self.project_root_str, self.language),
        ).fetchall()
        resolved_to_path: dict[str, Path] = {
            resolved: fp for fp, resolved, _ in all_caller_keys
        }
        counts: dict[tuple[Path, str], int] = {}
        for caller_file, caller_qn, cnt in rows:
            counts[(resolved_to_path[caller_file], caller_qn)] = cnt
        return counts

    def ensure_file_indexed(
        self,
        file_path: Path,
        resolved: str | None = None,
    ) -> IndexResult:
        """Index *file_path* if it is not already cached."""
        from ._call_graph import IndexResult  # noqa: PLC0415

        if resolved is None:
            resolved = self.resolve_path(file_path)
        try:
            content = file_path.read_text(encoding="utf-8")
        except Exception:  # noqa: BLE001
            return IndexResult(
                file_path=file_path,
                cached=False,
                num_edges=0,
                edges=(),
                cross_file_edges=0,
                error=True,
            )
        file_hash = hashlib.sha256(content.encode("utf-8")).hexdigest()
        if self._is_file_cached(resolved, file_hash):
            return IndexResult(
                file_path=file_path,
                cached=True,
                num_edges=0,
                edges=(),
                cross_file_edges=0,
                error=False,
            )
        return self.index_file(file_path, file_hash, resolved)

    def index_file(
        self,
        file_path: Path,
        file_hash: str,
        resolved: str | None = None,
    ) -> IndexResult:
        """Analyze *file_path* with Jedi and persist edges."""
        if resolved is None:
            resolved = self.resolve_path(file_path)
        edges, had_error = _analyze_file(
            file_path, self.jedi_project, self.project_root_str
        )
        if had_error:
            log.debug("ReferenceGraph: failed to parse %s", file_path)
        return self._persist_edges(
            file_path, resolved, file_hash, edges, had_error
        )

    def _persist_edges(
        self,
        file_path: Path,
        resolved: str,
        file_hash: str,
        edges: set[tuple[str, ...]],
        had_error: bool,  # noqa: FBT001
    ) -> IndexResult:
        """Write call edges to SQLite and return an IndexResult."""
        from ._call_graph import IndexResult  # noqa: PLC0415

        cur = self.conn.cursor()
        scope = (self.project_root_str, self.language)
        cur.execute(
            "DELETE FROM call_edges "
            "WHERE project_root = ? AND language = ? "
            "AND caller_file = ?",
            (*scope, resolved),
        )
        cur.execute(
            "DELETE FROM indexed_files "
            "WHERE project_root = ? AND language = ? "
            "AND file_path = ?",
            (*scope, resolved),
        )
        if not had_error and edges:
            cur.executemany(
                """
                INSERT OR REPLACE INTO call_edges
                (project_root, language, caller_file,
                 caller_qualified_name, callee_file,
                 callee_qualified_name,
                 callee_fully_qualified_name,
                 callee_only_function_name,
                 callee_definition_type, callee_source_line)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                [(*scope, *edge) for edge in edges],
            )
        cur.execute(
            "INSERT OR REPLACE INTO indexed_files "
            "(project_root, language, file_path, file_hash) "
            "VALUES (?, ?, ?, ?)",
            (*scope, resolved, file_hash),
        )
        self.conn.commit()
        self.indexed_file_hashes[resolved] = file_hash
        edges_summary = tuple(
            (
                caller_qn,
                callee_name,
                caller_file != callee_file,
            )
            for (
                caller_file,
                caller_qn,
                callee_file,
                _,
                _,
                callee_name,
                _,
                _,
            ) in edges
        )
        cross_file_count = sum(
            is_cross_file for _, _, is_cross_file in edges_summary
        )
        return IndexResult(
            file_path=file_path,
            cached=False,
            num_edges=len(edges),
            edges=edges_summary,
            cross_file_edges=cross_file_count,
            error=had_error,
        )

    def build_index(
        self,
        file_paths: Iterable[Path],
        on_progress: Callable[[IndexResult], None] | None = None,
    ) -> None:
        """Batch-index multiple files, using parallelism when possible."""
        from ._call_graph import IndexResult  # noqa: PLC0415

        to_index: list[tuple[Path, str, str]] = []
        for file_path in file_paths:
            resolved = self.resolve_path(file_path)
            if resolved in self.indexed_file_hashes:
                self._report_progress(
                    on_progress,
                    IndexResult(
                        file_path=file_path,
                        cached=True,
                        num_edges=0,
                        edges=(),
                        cross_file_edges=0,
                        error=False,
                    ),
                )
                continue
            try:
                content = file_path.read_text(encoding="utf-8")
            except Exception:  # noqa: BLE001
                self._report_progress(
                    on_progress,
                    IndexResult(
                        file_path=file_path,
                        cached=False,
                        num_edges=0,
                        edges=(),
                        cross_file_edges=0,
                        error=True,
                    ),
                )
                continue
            file_hash = hashlib.sha256(content.encode("utf-8")).hexdigest()
            if self._is_file_cached(resolved, file_hash):
                self._report_progress(
                    on_progress,
                    IndexResult(
                        file_path=file_path,
                        cached=True,
                        num_edges=0,
                        edges=(),
                        cross_file_edges=0,
                        error=False,
                    ),
                )
                continue
            to_index.append((file_path, resolved, file_hash))
        if not to_index:
            return
        if len(to_index) >= _PARALLEL_THRESHOLD:
            self._build_index_parallel(to_index, on_progress)
        else:
            for file_path, resolved, file_hash in to_index:
                result = self.index_file(file_path, file_hash, resolved)
                self._report_progress(on_progress, result)

    def _is_file_cached(self, resolved: str, file_hash: str) -> bool:
        """Return True if *resolved* is already indexed with *file_hash*."""
        if self.indexed_file_hashes.get(resolved) == file_hash:
            return True
        row = self.conn.execute(
            "SELECT file_hash FROM indexed_files "
            "WHERE project_root = ? AND language = ? "
            "AND file_path = ?",
            (self.project_root_str, self.language, resolved),
        ).fetchone()
        if row and row[0] == file_hash:
            self.indexed_file_hashes[resolved] = file_hash
            return True
        return False

    def _report_progress(
        self,
        on_progress: Callable[[IndexResult], None] | None,
        result: IndexResult,
    ) -> None:
        """Invoke the progress callback if provided."""
        if on_progress is not None:
            on_progress(result)

    def _build_index_parallel(
        self,
        to_index: list[tuple[Path, str, str]],
        on_progress: Callable[[IndexResult], None] | None,
    ) -> None:
        """Index files in parallel with a ProcessPoolExecutor."""
        from concurrent.futures import (  # noqa: PLC0415
            ProcessPoolExecutor,
            as_completed,
        )

        from ._call_graph import IndexResult  # noqa: PLC0415

        max_workers = min(os.cpu_count() or 1, len(to_index), 8)
        path_info: dict[str, tuple[Path, str]] = {
            resolved: (fp, fh) for fp, resolved, fh in to_index
        }
        worker_args = [(resolved, fh) for _fp, resolved, fh in to_index]
        log.debug(
            "ReferenceGraph: indexing %d files across %d workers",
            len(to_index),
            max_workers,
        )
        try:
            with ProcessPoolExecutor(
                max_workers=max_workers,
                initializer=_init_index_worker,
                initargs=(self.project_root_str,),
            ) as executor:
                futures = {
                    executor.submit(_index_file_worker, args): args[0]
                    for args in worker_args
                }
                for future in as_completed(futures):
                    resolved = futures[future]
                    file_path, file_hash = path_info[resolved]
                    try:
                        _, _, edges, had_error = future.result()
                    except Exception:  # noqa: BLE001
                        log.debug(
                            "ReferenceGraph: worker failed for %s",
                            file_path,
                        )
                        self._persist_edges(
                            file_path,
                            resolved,
                            file_hash,
                            set(),
                            had_error=True,
                        )
                        self._report_progress(
                            on_progress,
                            IndexResult(
                                file_path=file_path,
                                cached=False,
                                num_edges=0,
                                edges=(),
                                cross_file_edges=0,
                                error=True,
                            ),
                        )
                        continue
                    if had_error:
                        log.debug(
                            "ReferenceGraph: failed to parse %s",
                            file_path,
                        )
                    result = self._persist_edges(
                        file_path,
                        resolved,
                        file_hash,
                        edges,
                        had_error,
                    )
                    self._report_progress(on_progress, result)
        except Exception:  # noqa: BLE001
            log.debug(
                "ReferenceGraph: parallel indexing failed, "
                "falling back to sequential"
            )
            self._fallback_sequential_index(to_index, on_progress)

    def _fallback_sequential_index(
        self,
        to_index: list[tuple[Path, str, str]],
        on_progress: Callable[[IndexResult], None] | None,
    ) -> None:
        """Fall back to sequential indexing on parallel failure."""
        for file_path, resolved, file_hash in to_index:
            if resolved in self.indexed_file_hashes:
                continue
            result = self.index_file(file_path, file_hash, resolved)
            self._report_progress(on_progress, result)

    def get_call_graph(
        self,
        file_path_to_qualified_names: dict[Path, set[str]],
        *,
        include_metadata: bool = False,
    ) -> CallGraph:
        """Build a CallGraph from indexed edges in the database."""
        from ._call_graph import (  # noqa: PLC0415
            CallEdge,
            CalleeMetadata,
            CallGraph,
            FunctionNode,
        )

        all_caller_keys: list[tuple[Path, str, str]] = []
        for file_path, qualified_names in file_path_to_qualified_names.items():
            resolved = self.resolve_path(file_path)
            self.ensure_file_indexed(file_path, resolved)
            all_caller_keys.extend(
                (file_path, resolved, qn) for qn in qualified_names
            )
        if not all_caller_keys:
            return CallGraph(edges=[])
        cur = self.conn.cursor()
        cur.execute(
            "CREATE TEMP TABLE IF NOT EXISTS _graph_keys "
            "(caller_file TEXT, caller_qualified_name TEXT)"
        )
        cur.execute("DELETE FROM _graph_keys")
        cur.executemany(
            "INSERT INTO _graph_keys VALUES (?, ?)",
            [(resolved, qn) for _, resolved, qn in all_caller_keys],
        )
        if include_metadata:
            rows = cur.execute(
                """
                SELECT ce.caller_file,
                       ce.caller_qualified_name,
                       ce.callee_file,
                       ce.callee_qualified_name,
                       ce.callee_fully_qualified_name,
                       ce.callee_only_function_name,
                       ce.callee_definition_type,
                       ce.callee_source_line
                FROM call_edges ce
                INNER JOIN _graph_keys gk
                    ON ce.caller_file = gk.caller_file
                    AND ce.caller_qualified_name =
                        gk.caller_qualified_name
                WHERE ce.project_root = ?
                    AND ce.language = ?
                """,
                (self.project_root_str, self.language),
            ).fetchall()
            edges: list[CallEdge] = []
            for (
                caller_file,
                caller_qn,
                callee_file,
                callee_qn,
                callee_fqn,
                callee_name,
                callee_type,
                callee_src,
            ) in rows:
                edges.append(
                    CallEdge(
                        caller=FunctionNode(
                            file_path=Path(caller_file),
                            qualified_name=caller_qn,
                        ),
                        callee=FunctionNode(
                            file_path=Path(callee_file),
                            qualified_name=callee_qn,
                        ),
                        is_cross_file=caller_file != callee_file,
                        callee_metadata=CalleeMetadata(
                            fully_qualified_name=callee_fqn,
                            only_function_name=callee_name,
                            definition_type=callee_type,
                            source_line=callee_src,
                        ),
                    )
                )
        else:
            rows = cur.execute(
                """
                SELECT ce.caller_file,
                       ce.caller_qualified_name,
                       ce.callee_file,
                       ce.callee_qualified_name
                FROM call_edges ce
                INNER JOIN _graph_keys gk
                    ON ce.caller_file = gk.caller_file
                    AND ce.caller_qualified_name =
                        gk.caller_qualified_name
                WHERE ce.project_root = ?
                    AND ce.language = ?
                """,
                (self.project_root_str, self.language),
            ).fetchall()
            edges = []
            for (
                caller_file,
                caller_qn,
                callee_file,
                callee_qn,
            ) in rows:
                edges.append(
                    CallEdge(
                        caller=FunctionNode(
                            file_path=Path(caller_file),
                            qualified_name=caller_qn,
                        ),
                        callee=FunctionNode(
                            file_path=Path(callee_file),
                            qualified_name=callee_qn,
                        ),
                        is_cross_file=caller_file != callee_file,
                    )
                )
        return CallGraph(edges=edges)

    def close(self) -> None:
        """Close the database connection."""
        self.conn.close()
