# fitz_ai/engines/fitz_krag/ingestion/import_graph_store.py
"""CRUD operations for krag_import_graph table."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

from fitz_ai.engines.fitz_krag.ingestion.schema import TABLE_PREFIX

if TYPE_CHECKING:
    from fitz_ai.storage.postgres import PostgresConnectionManager

logger = logging.getLogger(__name__)

TABLE = f"{TABLE_PREFIX}import_graph"


class ImportGraphStore:
    """CRUD for the import graph."""

    def __init__(self, connection_manager: "PostgresConnectionManager", collection: str):
        self._cm = connection_manager
        self._collection = collection

    def upsert_batch(self, edges: list[dict[str, Any]]) -> None:
        """Insert or update a batch of import edges."""
        if not edges:
            return

        sql = f"""
            INSERT INTO {TABLE}
                ("source_file_id", "target_module", "target_file_id", "import_names")
            VALUES (%s, %s, %s, %s)
            ON CONFLICT ("source_file_id", "target_module") DO UPDATE SET
                "target_file_id" = EXCLUDED."target_file_id",
                "import_names" = EXCLUDED."import_names"
        """
        with self._cm.connection(self._collection) as conn:
            for edge in edges:
                conn.execute(
                    sql,
                    (
                        edge["source_file_id"],
                        edge["target_module"],
                        edge.get("target_file_id"),
                        edge.get("import_names", []),
                    ),
                )
            conn.commit()

    def get_imports(self, file_id: str) -> list[dict[str, Any]]:
        """Get all imports from a file."""
        sql = f"""
            SELECT "source_file_id", "target_module", "target_file_id", "import_names"
            FROM {TABLE} WHERE "source_file_id" = %s
        """
        with self._cm.connection(self._collection) as conn:
            rows = conn.execute(sql, (file_id,)).fetchall()
        return [
            {
                "source_file_id": row[0],
                "target_module": row[1],
                "target_file_id": row[2],
                "import_names": row[3] or [],
            }
            for row in rows
        ]

    def get_importers(self, file_id: str) -> list[dict[str, Any]]:
        """Get all files that import this file."""
        sql = f"""
            SELECT "source_file_id", "target_module", "target_file_id", "import_names"
            FROM {TABLE} WHERE "target_file_id" = %s
        """
        with self._cm.connection(self._collection) as conn:
            rows = conn.execute(sql, (file_id,)).fetchall()
        return [
            {
                "source_file_id": row[0],
                "target_module": row[1],
                "target_file_id": row[2],
                "import_names": row[3] or [],
            }
            for row in rows
        ]

    def get_reverse_counts(self) -> dict[str, int]:
        """Get the number of files that import each file.

        Returns ``{file_id: count}`` — useful for ranking files by connectivity.
        """
        sql = f"""
            SELECT "target_file_id", COUNT(*)
            FROM {TABLE}
            WHERE "target_file_id" IS NOT NULL
            GROUP BY "target_file_id"
        """
        with self._cm.connection(self._collection) as conn:
            rows = conn.execute(sql).fetchall()
        return {row[0]: int(row[1]) for row in rows}

    def delete_by_file(self, file_id: str) -> None:
        """Delete all import edges for a file (as source)."""
        sql = f'DELETE FROM {TABLE} WHERE "source_file_id" = %s'
        with self._cm.connection(self._collection) as conn:
            conn.execute(sql, (file_id,))
            conn.commit()

    def resolve_targets(self, path_to_id: dict[str, str]) -> int:
        """Resolve ``target_file_id`` for all unresolved import edges.

        Builds a module->file_id mapping from *path_to_id* (as returned by
        ``RawFileStore.list_ids_by_path()``), then UPDATEs every edge whose
        ``target_file_id IS NULL`` when the ``target_module`` matches a
        known file.

        Returns the number of edges resolved.
        """
        module_to_id: dict[str, str] = {}
        for path, file_id in path_to_id.items():
            module = _path_to_module(path)
            module_to_id[module] = file_id

        select_sql = f"""
            SELECT "source_file_id", "target_module"
            FROM {TABLE}
            WHERE "target_file_id" IS NULL
        """
        update_sql = f"""
            UPDATE {TABLE}
            SET "target_file_id" = %s
            WHERE "source_file_id" = %s AND "target_module" = %s
        """
        resolved = 0
        with self._cm.connection(self._collection) as conn:
            rows = conn.execute(select_sql).fetchall()
            for source_file_id, target_module in rows:
                target_id = module_to_id.get(target_module)
                if target_id:
                    conn.execute(update_sql, (target_id, source_file_id, target_module))
                    resolved += 1
            if resolved:
                conn.commit()
        return resolved


def _path_to_module(file_path: str) -> str:
    """Convert a file path to a Python-style module name.

    Intentionally duplicates the 6-line helper in each language strategy
    to avoid coupling import_graph_store to any specific strategy.
    """
    path = file_path.replace("\\", "/")
    if path.startswith("./"):
        path = path[2:]
    if path.endswith(".py"):
        path = path[:-3]
    if path.endswith("/__init__"):
        path = path[: -len("/__init__")]
    return path.replace("/", ".")
