"""Repository methods for code graph data."""

from __future__ import annotations

from dataclasses import asdict, dataclass
from pathlib import Path

from sqlalchemy import case, delete, func, insert, or_, select
from sqlalchemy.orm import aliased

from codectx.db.models import CodeVersion, Dependency, Entity, EntityRelation, File
from codectx.parser.types import DependencyRecord, EntityRecord, RelationRecord
from common.db import Base, create_engine_with_wal, create_session_factory, session_scope


@dataclass(slots=True)
class CodeStats:
    files: int = 0
    entities: int = 0
    dependencies: int = 0
    relations: int = 0
    lines_of_code: int = 0
    tokens: int = 0
    versions: int = 0
    latest_version: str = ""

    def to_dict(self) -> dict[str, object]:
        return asdict(self)


class CodeRepository:
    """Persistence wrapper for codectx."""

    def __init__(self, db_path: str | Path, echo: bool = False) -> None:
        self.db_path = Path(db_path)
        self.engine = create_engine_with_wal(f"sqlite:///{self.db_path}", echo=echo)
        self.session_factory = create_session_factory(self.engine)

    def create_schema(self) -> None:
        Base.metadata.create_all(self.engine, tables=[
            CodeVersion.__table__, File.__table__, Entity.__table__,
            Dependency.__table__, EntityRelation.__table__,
        ])

    # ------------------------------------------------------------------
    # Version management
    # ------------------------------------------------------------------

    def create_version(
        self, version_tag: str, release_label: str = "", base_commit: str = "",
    ) -> int:
        """Create a new code version and mark it as latest."""
        with session_scope(self.session_factory) as session:
            # Unmark previous latest
            prev = session.scalar(
                select(CodeVersion).where(CodeVersion.is_latest == True)  # noqa: E712
            )
            if prev is not None:
                prev.is_latest = False
            ver = CodeVersion(
                version_tag=version_tag,
                release_label=release_label,
                base_commit=base_commit,
                is_latest=True,
            )
            session.add(ver)
            session.flush()
            return ver.id

    def get_version(self, version_tag: str | None = None) -> CodeVersion | None:
        """Get a version by tag, or the latest version if tag is None."""
        with session_scope(self.session_factory) as session:
            if version_tag:
                return session.scalar(
                    select(CodeVersion).where(CodeVersion.version_tag == version_tag)
                )
            return session.scalar(
                select(CodeVersion).where(CodeVersion.is_latest == True)  # noqa: E712
            )

    def list_versions(self) -> list[dict[str, object]]:
        """List all indexed versions."""
        with session_scope(self.session_factory) as session:
            rows = list(session.scalars(
                select(CodeVersion).order_by(CodeVersion.indexed_at.desc())
            ))
            return [
                {
                    "id": v.id,
                    "version_tag": v.version_tag,
                    "release_label": v.release_label,
                    "base_commit": v.base_commit,
                    "file_count": v.file_count,
                    "entity_count": v.entity_count,
                    "is_latest": v.is_latest,
                    "indexed_at": v.indexed_at.isoformat() if v.indexed_at else None,
                }
                for v in rows
            ]

    def update_version_counts(self, version_id: int) -> None:
        """Update file_count and entity_count for a version."""
        with session_scope(self.session_factory) as session:
            ver = session.get(CodeVersion, version_id)
            if ver is None:
                return
            ver.file_count = session.scalar(
                select(func.count()).select_from(File).where(File.version_id == version_id)
            ) or 0
            ver.entity_count = session.scalar(
                select(func.count()).select_from(Entity)
                .join(File, Entity.file_id == File.id)
                .where(File.version_id == version_id)
            ) or 0

    def get_changed_files(self, old_version_tag: str, directory: str = "") -> list[dict[str, object]]:
        """Compare two versions and return files that changed (added/modified/removed)."""
        with session_scope(self.session_factory) as session:
            old_ver = session.scalar(
                select(CodeVersion).where(CodeVersion.version_tag == old_version_tag)
            )
            new_ver = session.scalar(
                select(CodeVersion).where(CodeVersion.is_latest == True)  # noqa: E712
            )
            if not old_ver or not new_ver:
                return []
            old_files = {
                f.path: f.hash
                for f in session.scalars(select(File).where(File.version_id == old_ver.id))
            }
            new_files = {
                f.path: f.hash
                for f in session.scalars(select(File).where(File.version_id == new_ver.id))
            }
            changes = []
            for path, h in new_files.items():
                if path not in old_files:
                    changes.append({"path": path, "change": "added"})
                elif old_files[path] != h:
                    changes.append({"path": path, "change": "modified"})
            for path in old_files:
                if path not in new_files:
                    changes.append({"path": path, "change": "removed"})
            return changes

    def upsert_file(
        self, path: str, language: str, file_hash: str = "",
        lines_of_code: int = 0, tokens: int = 0, version_id: int | None = None,
    ) -> int:
        with session_scope(self.session_factory) as session:
            # When versioned, scope uniqueness to (path, version_id)
            if version_id is not None:
                existing = session.scalar(
                    select(File).where(File.path == path, File.version_id == version_id)
                )
            else:
                existing = session.scalar(select(File).where(File.path == path))
            if existing is None:
                file_row = File(
                    path=path, language=language, hash=file_hash or None,
                    lines_of_code=lines_of_code, tokens=tokens, version_id=version_id,
                )
                session.add(file_row)
                session.flush()
                return file_row.id

            existing.language = language
            existing.hash = file_hash or None
            existing.lines_of_code = lines_of_code
            existing.tokens = tokens
            session.flush()
            return existing.id

    def file_needs_update(self, path: str, new_hash: str) -> bool:
        """Check if a file needs updating by comparing its current hash with the stored hash."""
        with session_scope(self.session_factory) as session:
            existing = session.scalar(select(File).where(File.path == path))
            if existing is None:
                return True  # New file
            if existing.hash != new_hash:
                return True  # File changed
            return False  # File unchanged

    def get_file_by_path(self, path: str) -> File | None:
        with session_scope(self.session_factory) as session:
            return session.scalar(select(File).where(File.path == path))

    def get_entities_by_name(self, name: str) -> list[Entity]:
        with session_scope(self.session_factory) as session:
            return list(
                session.scalars(
                    select(Entity)
                    .where(or_(Entity.name == name, Entity.qualified_name == name))
                    .order_by(Entity.file_id, Entity.start_line)
                )
            )

    def get_dependency_graph(self, path: str) -> dict[str, object] | None:
        with session_scope(self.session_factory) as session:
            file_row = session.scalar(select(File).where(File.path == path))
            if file_row is None:
                return None
            dependencies = list(
                session.scalars(
                    select(Dependency).where(Dependency.source_file_id == file_row.id).order_by(Dependency.line_number, Dependency.target_path)
                )
            )
            return {
                "file": file_row.path,
                "dependencies": [dependency.target_path for dependency in dependencies],
            }

    def get_call_graph(self, name: str) -> dict[str, object] | None:
        with session_scope(self.session_factory) as session:
            entity = session.scalar(
                select(Entity)
                .where(or_(Entity.name == name, Entity.qualified_name == name))
                .order_by(Entity.start_line)
            )
            if entity is None:
                return None

            source_entity = aliased(Entity)
            target_entity = aliased(Entity)
            outgoing_rows = session.execute(
                select(EntityRelation, target_entity)
                .join(target_entity, EntityRelation.target_entity_id == target_entity.id)
                .where(EntityRelation.source_entity_id == entity.id, EntityRelation.relation_type == "calls")
                .order_by(EntityRelation.line_number, target_entity.qualified_name, target_entity.name)
            )
            incoming_rows = session.execute(
                select(EntityRelation, source_entity)
                .join(source_entity, EntityRelation.source_entity_id == source_entity.id)
                .where(EntityRelation.target_entity_id == entity.id, EntityRelation.relation_type == "calls")
                .order_by(EntityRelation.line_number, source_entity.qualified_name, source_entity.name)
            )
            return {
                "entity": entity.qualified_name or entity.name,
                "outgoing": [
                    {
                        "target": target.qualified_name or target.name,
                        "line_number": relation.line_number,
                        "context": relation.context or "",
                    }
                    for relation, target in outgoing_rows
                ],
                "incoming": [
                    {
                        "source": source.qualified_name or source.name,
                        "line_number": relation.line_number,
                        "context": relation.context or "",
                    }
                    for relation, source in incoming_rows
                ],
            }

    def delete_entities_by_file(self, file_id: int) -> None:
        with session_scope(self.session_factory) as session:
            entity_ids = list(session.scalars(select(Entity.id).where(Entity.file_id == file_id)))
            if entity_ids:
                session.execute(delete(EntityRelation).where(
                    or_(EntityRelation.source_entity_id.in_(entity_ids),
                        EntityRelation.target_entity_id.in_(entity_ids))
                ))
            session.execute(delete(Entity).where(Entity.file_id == file_id))

    def delete_dependencies_by_file(self, file_id: int) -> None:
        with session_scope(self.session_factory) as session:
            session.execute(delete(Dependency).where(Dependency.source_file_id == file_id))

    def delete_relations_by_file(self, file_id: int) -> None:
        with session_scope(self.session_factory) as session:
            entity_ids = list(session.scalars(select(Entity.id).where(Entity.file_id == file_id)))
            if not entity_ids:
                return
            session.execute(delete(EntityRelation).where(
                or_(EntityRelation.source_entity_id.in_(entity_ids),
                    EntityRelation.target_entity_id.in_(entity_ids))
            ))

    def insert_entity(self, file_id: int, entity: EntityRecord) -> int:
        with session_scope(self.session_factory) as session:
            row = Entity(
                file_id=file_id,
                name=entity.name,
                qualified_name=entity.qualified_name or None,
                type=entity.type,
                kind=entity.kind or None,
                signature=entity.signature or None,
                start_line=entity.start_line,
                end_line=entity.end_line,
                column_start=entity.column_start,
                column_end=entity.column_end,
                documentation=entity.docs or None,
                parent=entity.parent or None,
                visibility=entity.visibility or None,
                scope=entity.scope or None,
                language=entity.language or None,
                attributes=None,
            )
            session.add(row)
            session.flush()
            return row.id

    def insert_dependency(self, file_id: int, dependency: DependencyRecord) -> int:
        with session_scope(self.session_factory) as session:
            row = Dependency(
                source_file_id=file_id,
                target_path=dependency.path,
                import_type=dependency.type or None,
                line_number=dependency.line_number,
                resolved=dependency.resolved or None,
                is_local=dependency.is_local,
            )
            session.add(row)
            session.flush()
            return row.id

    def insert_relation(self, file_id: int, relation: RelationRecord) -> int | None:
        with session_scope(self.session_factory) as session:
            source = session.scalar(
                select(Entity)
                .where(
                    Entity.file_id == file_id,
                    or_(Entity.qualified_name == relation.source_qualified_name, Entity.name == relation.source_qualified_name),
                )
                .order_by(Entity.start_line)
            )
            if source is None:
                return None

            target = self._resolve_target_entity(session, file_id, relation.target_qualified_name)
            if target is None:
                return None

            row = EntityRelation(
                source_entity_id=source.id,
                target_entity_id=target.id,
                relation_type=relation.relation_type,
                line_number=relation.line_number,
                context=relation.context or None,
            )
            session.add(row)
            session.flush()
            return row.id

    # ------------------------------------------------------------------
    # Bulk write (high-performance batch path used by the indexer)
    # ------------------------------------------------------------------

    def bulk_write_file_data(
        self,
        records: list[tuple[int, object, list]],
    ) -> None:
        """Delete stale rows and bulk-insert entities + dependencies for all
        files in *records* inside a **single transaction**.

        Args:
            records: list of (file_id, parse_result, resolved_dependencies)
                     where parse_result has .entities and .relations attributes.
        """
        if not records:
            return

        file_ids = [file_id for file_id, _, _ in records]

        with session_scope(self.session_factory) as session:
            # --- Step 1: bulk-delete stale data ---
            existing_entity_ids = list(
                session.scalars(select(Entity.id).where(Entity.file_id.in_(file_ids)))
            )
            if existing_entity_ids:
                session.execute(
                    delete(EntityRelation).where(
                        or_(
                            EntityRelation.source_entity_id.in_(existing_entity_ids),
                            EntityRelation.target_entity_id.in_(existing_entity_ids),
                        )
                    )
                )
            session.execute(delete(Entity).where(Entity.file_id.in_(file_ids)))
            session.execute(delete(Dependency).where(Dependency.source_file_id.in_(file_ids)))

            # --- Step 2: bulk-insert entities (ORM add_all → single flush) ---
            entity_rows: list[Entity] = []
            for file_id, parse_result, _ in records:
                for entity in parse_result.entities:
                    entity_rows.append(Entity(
                        file_id=file_id,
                        name=entity.name,
                        qualified_name=entity.qualified_name or None,
                        type=entity.type,
                        kind=entity.kind or None,
                        signature=entity.signature or None,
                        start_line=entity.start_line,
                        end_line=entity.end_line,
                        column_start=entity.column_start,
                        column_end=entity.column_end,
                        documentation=entity.docs or None,
                        parent=entity.parent or None,
                        visibility=entity.visibility or None,
                        scope=entity.scope or None,
                        language=entity.language or None,
                        attributes=None,
                    ))
            session.add_all(entity_rows)
            session.flush()  # single round-trip; populates .id on all rows

            # --- Step 3: bulk-insert dependencies using core INSERT ---
            dep_dicts: list[dict] = []
            for file_id, parse_result, resolved_deps in records:
                for dep in resolved_deps:
                    dep_dicts.append(dict(
                        source_file_id=file_id,
                        target_path=dep.path,
                        import_type=dep.type or None,
                        line_number=dep.line_number,
                        resolved=dep.resolved or None,
                        is_local=dep.is_local,
                    ))
            if dep_dicts:
                session.execute(insert(Dependency), dep_dicts)

            # --- Step 4: bulk-insert relations via in-batch entity lookup ---
            entity_lookup: dict[tuple[int, str], int] = {}
            for row in entity_rows:
                fid = row.file_id
                if row.qualified_name:
                    entity_lookup[(fid, row.qualified_name)] = row.id
                entity_lookup[(fid, row.name)] = row.id

            relation_dicts: list[dict] = []
            for file_id, parse_result, _ in records:
                for relation in parse_result.relations:
                    src_name = relation.source_qualified_name
                    src_id = (
                        entity_lookup.get((file_id, src_name))
                        or entity_lookup.get((file_id, src_name.rsplit(".", 1)[-1]))
                    )
                    tgt_name = relation.target_qualified_name
                    tgt_id = (
                        entity_lookup.get((file_id, tgt_name))
                        or entity_lookup.get((file_id, tgt_name.rsplit(".", 1)[-1]))
                    )
                    # Cross-file fallback
                    if src_id is None:
                        for (_, name), eid in entity_lookup.items():
                            if name == src_name:
                                src_id = eid
                                break
                    if tgt_id is None:
                        tgt_short = tgt_name.rsplit(".", 1)[-1]
                        for (_, name), eid in entity_lookup.items():
                            if name in (tgt_name, tgt_short):
                                tgt_id = eid
                                break
                    if src_id is None or tgt_id is None:
                        continue
                    relation_dicts.append(dict(
                        source_entity_id=src_id,
                        target_entity_id=tgt_id,
                        relation_type=relation.relation_type,
                        line_number=relation.line_number,
                        context=relation.context or None,
                    ))
            if relation_dicts:
                session.execute(insert(EntityRelation), relation_dicts)

    def _resolve_target_entity(self, session, file_id: int, target_name: str) -> Entity | None:
        module_hint, _, symbol_name = target_name.rpartition(".")
        fallback_symbol = symbol_name or target_name

        imported_file_ids: list[int] = []
        if module_hint:
            imported_file_ids = self._resolve_imported_file_ids(session, file_id, module_hint)
            if imported_file_ids:
                target = session.scalar(
                    select(Entity)
                    .where(
                        Entity.file_id.in_(imported_file_ids),
                        or_(
                            Entity.qualified_name == target_name,
                            Entity.name == fallback_symbol,
                            Entity.qualified_name == fallback_symbol,
                        ),
                    )
                    .order_by(Entity.start_line.asc())
                )
                if target is not None:
                    return target

        prefer_same_file = case((Entity.file_id == file_id, 0), else_=1)
        prefer_exact_qualified = case((Entity.qualified_name == target_name, 0), else_=1)
        return session.scalar(
            select(Entity)
            .where(or_(Entity.qualified_name == target_name, Entity.name == fallback_symbol, Entity.name == target_name))
            .order_by(prefer_exact_qualified.asc(), prefer_same_file.asc(), Entity.start_line.asc())
        )

    def _resolve_imported_file_ids(self, session, file_id: int, module_hint: str) -> list[int]:
        dependency_rows = list(session.scalars(select(Dependency).where(Dependency.source_file_id == file_id)))
        matched_paths = []
        normalized_hint = module_hint.strip(".")
        hint_tail = normalized_hint.rsplit(".", 1)[-1]
        for dependency in dependency_rows:
            dependency_module = (dependency.target_path or "").strip(".")
            dependency_tail = dependency_module.rsplit(".", 1)[-1] if dependency_module else ""
            if normalized_hint in {dependency_module, dependency_tail} or dependency_module.endswith(f".{normalized_hint}"):
                if dependency.resolved:
                    matched_paths.append(dependency.resolved)

        if not matched_paths:
            return []

        return list(
            session.scalars(select(File.id).where(File.path.in_(matched_paths)))
        )

    def list_files(self, limit: int = 100, offset: int = 0) -> list[dict[str, object]]:
        with session_scope(self.session_factory) as session:
            rows = list(
                session.scalars(
                    select(File).order_by(File.path).limit(limit).offset(offset)
                )
            )
            return [
                {
                    "id": row.id,
                    "path": row.path,
                    "language": row.language,
                    "lines_of_code": row.lines_of_code,
                    "tokens": row.tokens,
                }
                for row in rows
            ]

    def search_entities(self, query: str, limit: int = 20) -> list[dict[str, object]]:
        with session_scope(self.session_factory) as session:
            pattern = f"%{query}%"
            rows = list(
                session.scalars(
                    select(Entity)
                    .where(
                        or_(
                            Entity.name.ilike(pattern),
                            Entity.qualified_name.ilike(pattern),
                        )
                    )
                    .order_by(Entity.name, Entity.start_line)
                    .limit(limit)
                )
            )
            return [
                {
                    "id": row.id,
                    "name": row.name,
                    "qualified_name": row.qualified_name,
                    "type": row.type,
                    "kind": row.kind,
                    "file_id": row.file_id,
                    "start_line": row.start_line,
                    "end_line": row.end_line,
                    "documentation": row.documentation,
                    "language": row.language,
                }
                for row in rows
            ]

    def get_entity_by_id(self, entity_id: int) -> dict[str, object] | None:
        with session_scope(self.session_factory) as session:
            row = session.get(Entity, entity_id)
            if row is None:
                return None
            return {
                "id": row.id,
                "name": row.name,
                "qualified_name": row.qualified_name,
                "type": row.type,
                "kind": row.kind,
                "signature": row.signature,
                "file_id": row.file_id,
                "start_line": row.start_line,
                "end_line": row.end_line,
                "documentation": row.documentation,
                "visibility": row.visibility,
                "scope": row.scope,
                "language": row.language,
                "parent": row.parent,
            }

    def get_entity_code(self, entity_id: int, source_root: str = ".") -> dict[str, object] | None:
        """Return the source lines for an entity, reading from disk."""
        with session_scope(self.session_factory) as session:
            entity = session.get(Entity, entity_id)
            if entity is None:
                return None
            file_row = session.get(File, entity.file_id)
            if file_row is None:
                return None

            file_path = Path(source_root) / file_row.path if not Path(file_row.path).is_absolute() else Path(file_row.path)
            code = ""
            if file_path.exists():
                try:
                    lines = file_path.read_text(encoding="utf-8", errors="replace").splitlines()
                    start = max(0, (entity.start_line or 1) - 1)
                    end = min(len(lines), entity.end_line or len(lines))
                    code = "\n".join(lines[start:end])
                except OSError:
                    code = ""

            return {
                "entity_id": entity_id,
                "name": entity.name,
                "qualified_name": entity.qualified_name,
                "file": file_row.path,
                "start_line": entity.start_line,
                "end_line": entity.end_line,
                "code": code,
            }

    def get_file_imports(self, file_path: str) -> list[dict[str, object]]:
        with session_scope(self.session_factory) as session:
            file_row = session.scalar(select(File).where(File.path == file_path))
            if file_row is None:
                return []
            rows = list(
                session.scalars(
                    select(Dependency)
                    .where(Dependency.source_file_id == file_row.id)
                    .order_by(Dependency.line_number, Dependency.target_path)
                )
            )
            return [
                {
                    "target_path": row.target_path,
                    "import_type": row.import_type,
                    "line_number": row.line_number,
                    "resolved": row.resolved,
                    "is_local": row.is_local,
                }
                for row in rows
            ]

    def get_file_entities(self, file_path: str) -> list[dict[str, object]]:
        with session_scope(self.session_factory) as session:
            file_row = session.scalar(select(File).where(File.path == file_path))
            if file_row is None:
                return []
            rows = list(
                session.scalars(
                    select(Entity)
                    .where(Entity.file_id == file_row.id)
                    .order_by(Entity.start_line)
                )
            )
            return [
                {
                    "id": row.id,
                    "name": row.name,
                    "qualified_name": row.qualified_name,
                    "type": row.type,
                    "kind": row.kind,
                    "start_line": row.start_line,
                    "end_line": row.end_line,
                    "documentation": row.documentation,
                    "visibility": row.visibility,
                }
                for row in rows
            ]

    def get_stats(self, version_tag: str | None = None) -> CodeStats:
        with session_scope(self.session_factory) as session:
            # Optionally scope to a specific version
            file_q = select(func.count()).select_from(File)
            entity_q = select(func.count()).select_from(Entity)
            dep_q = select(func.count()).select_from(Dependency)
            rel_q = select(func.count()).select_from(EntityRelation)
            loc_q = select(func.coalesce(func.sum(File.lines_of_code), 0))
            tok_q = select(func.coalesce(func.sum(File.tokens), 0))

            if version_tag:
                ver = session.scalar(
                    select(CodeVersion).where(CodeVersion.version_tag == version_tag)
                )
                if ver:
                    vid = ver.id
                    file_q = file_q.where(File.version_id == vid)
                    entity_q = entity_q.join(File, Entity.file_id == File.id).where(File.version_id == vid)
                    dep_q = dep_q.join(File, Dependency.source_file_id == File.id).where(File.version_id == vid)
                    rel_q = rel_q.join(Entity, EntityRelation.source_entity_id == Entity.id).join(File, Entity.file_id == File.id).where(File.version_id == vid)
                    loc_q = loc_q.where(File.version_id == vid)
                    tok_q = tok_q.where(File.version_id == vid)

            version_count = session.scalar(select(func.count()).select_from(CodeVersion)) or 0
            latest = session.scalar(
                select(CodeVersion.version_tag).where(CodeVersion.is_latest == True)  # noqa: E712
            )
            return CodeStats(
                files=int(session.scalar(file_q) or 0),
                entities=int(session.scalar(entity_q) or 0),
                dependencies=int(session.scalar(dep_q) or 0),
                relations=int(session.scalar(rel_q) or 0),
                lines_of_code=int(session.scalar(loc_q) or 0),
                tokens=int(session.scalar(tok_q) or 0),
                versions=int(version_count),
                latest_version=latest or "",
            )

    # ------------------------------------------------------------------
    # Batch query methods
    # ------------------------------------------------------------------

    def get_file_entity_counts(self, file_ids: list[int]) -> dict[int, int]:
        """Return {file_id: entity_count} for the given file IDs."""
        if not file_ids:
            return {}
        with session_scope(self.session_factory) as session:
            rows = session.execute(
                select(Entity.file_id, func.count(Entity.id))
                .where(Entity.file_id.in_(file_ids))
                .group_by(Entity.file_id)
            )
            return {file_id: int(count) for file_id, count in rows}

    def get_importers_of_file(self, file_path: str) -> list[dict[str, object]]:
        """Return files that import the given file path."""
        with session_scope(self.session_factory) as session:
            rows = session.execute(
                select(File, Dependency)
                .join(Dependency, Dependency.source_file_id == File.id)
                .where(or_(Dependency.target_path == file_path, Dependency.resolved == file_path))
                .order_by(File.path)
            )
            return [
                {
                    "file_id": file_row.id,
                    "path": file_row.path,
                    "language": file_row.language,
                    "import_type": dep_row.import_type,
                    "line_number": dep_row.line_number,
                }
                for file_row, dep_row in rows
            ]

    def get_entities_by_file_ids(self, file_ids: list[int]) -> list[dict[str, object]]:
        """Return all entities for a set of file IDs."""
        if not file_ids:
            return []
        with session_scope(self.session_factory) as session:
            rows = list(session.scalars(
                select(Entity)
                .where(Entity.file_id.in_(file_ids))
                .order_by(Entity.file_id, Entity.start_line)
            ))
            return [
                {
                    "id": row.id,
                    "file_id": row.file_id,
                    "name": row.name,
                    "qualified_name": row.qualified_name,
                    "type": row.type,
                    "kind": row.kind,
                    "start_line": row.start_line,
                    "end_line": row.end_line,
                }
                for row in rows
            ]

    def get_dependencies_by_file_ids(self, file_ids: list[int]) -> list[dict[str, object]]:
        """Return all dependencies for a set of file IDs."""
        if not file_ids:
            return []
        with session_scope(self.session_factory) as session:
            rows = list(session.scalars(
                select(Dependency)
                .where(Dependency.source_file_id.in_(file_ids))
                .order_by(Dependency.source_file_id, Dependency.line_number)
            ))
            return [
                {
                    "file_id": row.source_file_id,
                    "target_path": row.target_path,
                    "import_type": row.import_type,
                    "line_number": row.line_number,
                    "resolved": row.resolved,
                    "is_local": row.is_local,
                }
                for row in rows
            ]

    def get_relations_by_entity_ids(self, entity_ids: list[int]) -> list[dict[str, object]]:
        """Return all relations where source or target is in entity_ids."""
        if not entity_ids:
            return []
        with session_scope(self.session_factory) as session:
            rows = list(session.scalars(
                select(EntityRelation)
                .where(
                    or_(
                        EntityRelation.source_entity_id.in_(entity_ids),
                        EntityRelation.target_entity_id.in_(entity_ids),
                    )
                )
                .order_by(EntityRelation.source_entity_id, EntityRelation.line_number)
            ))
            return [
                {
                    "id": row.id,
                    "source_entity_id": row.source_entity_id,
                    "target_entity_id": row.target_entity_id,
                    "relation_type": row.relation_type,
                    "line_number": row.line_number,
                    "context": row.context,
                }
                for row in rows
            ]

    def get_entities_by_ids(self, entity_ids: list[int]) -> list[dict[str, object]]:
        """Return entities for a list of entity IDs."""
        if not entity_ids:
            return []
        with session_scope(self.session_factory) as session:
            rows = list(session.scalars(
                select(Entity).where(Entity.id.in_(entity_ids)).order_by(Entity.id)
            ))
            return [
                {
                    "id": row.id,
                    "name": row.name,
                    "qualified_name": row.qualified_name,
                    "type": row.type,
                    "kind": row.kind,
                    "file_id": row.file_id,
                    "start_line": row.start_line,
                    "end_line": row.end_line,
                    "documentation": row.documentation,
                    "visibility": row.visibility,
                    "language": row.language,
                }
                for row in rows
            ]

    def get_directory_tree(self, root_prefix: str = "") -> dict[str, object]:
        """Return a nested directory tree of all indexed files."""
        with session_scope(self.session_factory) as session:
            rows = list(session.scalars(select(File).order_by(File.path)))
        tree: dict[str, object] = {}
        for file_row in rows:
            path = file_row.path
            if root_prefix and not path.startswith(root_prefix):
                continue
            parts = path.replace("\\", "/").split("/")
            node = tree
            for part in parts[:-1]:
                node = node.setdefault(part, {})  # type: ignore[assignment]
            filename = parts[-1]
            node[filename] = {
                "id": file_row.id,
                "language": file_row.language,
                "lines_of_code": file_row.lines_of_code,
            }
        return tree

    def get_file_by_id(self, file_id: int) -> dict | None:
        """Return a single file record by its primary key."""
        with session_scope(self.session_factory) as session:
            row = session.get(File, file_id)
            if row is None:
                return None
            return {
                "id": row.id,
                "path": row.path,
                "language": row.language,
                "hash": row.hash,
                "lines_of_code": row.lines_of_code,
                "tokens": row.tokens,
            }

    def get_all_entities(self, limit: int = 1000) -> list[dict]:
        """Return up to *limit* entity records across all files."""
        with session_scope(self.session_factory) as session:
            rows = list(session.scalars(select(Entity).order_by(Entity.id).limit(limit)))
            return [
                {
                    "id": row.id,
                    "name": row.name,
                    "qualified_name": row.qualified_name,
                    "type": row.type,
                    "kind": row.kind,
                    "file_id": row.file_id,
                    "start_line": row.start_line,
                    "end_line": row.end_line,
                    "documentation": row.documentation,
                    "visibility": row.visibility,
                    "language": row.language,
                }
                for row in rows
            ]

    def get_all_dependencies(self, limit: int = 1000) -> list[dict]:
        """Return up to *limit* dependency records across all files."""
        with session_scope(self.session_factory) as session:
            rows = list(session.scalars(select(Dependency).order_by(Dependency.id).limit(limit)))
            return [
                {
                    "id": row.id,
                    "source_file_id": row.source_file_id,
                    "target_path": row.target_path,
                    "import_type": row.import_type,
                    "line_number": row.line_number,
                    "resolved": row.resolved,
                    "is_local": row.is_local,
                }
                for row in rows
            ]

    def get_all_relations(self, limit: int = 1000) -> list[dict]:
        """Return up to *limit* entity-relation records across all files."""
        with session_scope(self.session_factory) as session:
            rows = list(session.scalars(select(EntityRelation).order_by(EntityRelation.id).limit(limit)))
            return [
                {
                    "id": row.id,
                    "source_entity_id": row.source_entity_id,
                    "target_entity_id": row.target_entity_id,
                    "relation_type": row.relation_type,
                    "line_number": row.line_number,
                }
                for row in rows
            ]
