"""
Database management for groknroll

Reliability Features (REL-006):
- Connection pooling via SQLAlchemy
- Explicit session cleanup
- Context manager support for safe transactions
"""

import hashlib
import json
import logging
from contextlib import contextmanager
from datetime import datetime
from pathlib import Path
from typing import Any, Generator, Optional

from sqlalchemy import create_engine, desc, event, or_
from sqlalchemy.orm import Session as DBSession
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import QueuePool

from groknroll.storage.models import (
    Analysis,
    Base,
    CodeRelation,
    CodeSummary,
    CodeTerm,
    ContextEntry,
    Execution,
    FileIndex,
    Project,
    Session,
    TermOccurrence,
)


# REL-006: Configure logger for database operations
logger = logging.getLogger(__name__)


class Database:
    """
    SQLite database for groknroll project state

    Stores:
    - Project metadata
    - File index with AST data
    - RLM execution history
    - Session history
    - Analysis results
    """

    def __init__(self, db_path: Optional[Path] = None):
        """
        Initialize database

        Args:
            db_path: Path to SQLite database file.
                    If None, uses ~/.groknroll/groknroll.db
        """
        if db_path is None:
            db_path = Path.home() / ".groknroll" / "groknroll.db"

        db_path.parent.mkdir(parents=True, exist_ok=True)

        self.db_path = db_path

        # REL-006: Configure connection pooling
        # Using QueuePool for thread-safe connection reuse
        self.engine = create_engine(
            f"sqlite:///{db_path}",
            echo=False,
            poolclass=QueuePool,
            pool_size=5,  # Number of persistent connections
            max_overflow=10,  # Additional connections when pool is full
            pool_timeout=30,  # Seconds to wait for connection
            pool_recycle=3600,  # Recycle connections after 1 hour
        )

        # Enable SQLite performance optimizations
        @event.listens_for(self.engine, "connect")
        def set_sqlite_pragma(dbapi_connection, connection_record):
            cursor = dbapi_connection.cursor()
            # WAL mode for 10x faster concurrent reads
            cursor.execute("PRAGMA journal_mode=WAL")
            # NORMAL sync is safe with WAL and faster than FULL
            cursor.execute("PRAGMA synchronous=NORMAL")
            # 64MB cache for better read performance
            cursor.execute("PRAGMA cache_size=-64000")
            # Memory-mapped I/O for large databases (256MB)
            cursor.execute("PRAGMA mmap_size=268435456")
            cursor.close()

        self.SessionLocal = sessionmaker(bind=self.engine)

        # Create tables
        Base.metadata.create_all(self.engine)

    def get_session(self) -> DBSession:
        """Get database session"""
        return self.SessionLocal()

    @contextmanager
    def session_scope(self) -> Generator[DBSession, None, None]:
        """
        Context manager for safe session handling.

        REL-006: Ensures sessions are properly committed or rolled back,
        and always closed, even if exceptions occur.

        Usage:
            with db.session_scope() as session:
                session.add(item)
                # Automatically commits on success, rolls back on error

        Yields:
            DBSession: Active database session
        """
        session = self.SessionLocal()
        try:
            yield session
            session.commit()
        except Exception as e:
            session.rollback()
            logger.error(f"Database error, rolling back: {e}")
            raise
        finally:
            session.close()

    def dispose(self) -> None:
        """
        Dispose of all connection pool connections.

        REL-006: Call this when shutting down to release all resources.
        """
        if self.engine:
            self.engine.dispose()
            logger.debug("Database connection pool disposed")

    def __enter__(self):
        """Context manager entry."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - dispose connections."""
        self.dispose()
        return False

    # =========================================================================
    # Project Operations
    # =========================================================================

    def get_or_create_project(self, project_path: Path) -> Project:
        """Get existing project or create new one"""
        with self.get_session() as session:
            project = session.query(Project).filter_by(path=str(project_path)).first()

            if project is None:
                project = Project(path=str(project_path), name=project_path.name)
                session.add(project)
                session.commit()
                session.refresh(project)

            return project

    def update_project_stats(self, project_id: int, total_files: int, total_lines: int) -> None:
        """Update project statistics"""
        with self.get_session() as session:
            project = session.query(Project).filter_by(id=project_id).first()
            if project:
                project.total_files = total_files
                project.total_lines = total_lines
                project.last_indexed = datetime.utcnow()
                session.commit()

    # =========================================================================
    # File Index Operations
    # =========================================================================

    def index_file(
        self, project_id: int, file_path: Path, relative_path: str, **metadata
    ) -> FileIndex:
        """Index a file"""
        with self.get_session() as session:
            file_index = (
                session.query(FileIndex)
                .filter_by(project_id=project_id, path=str(file_path))
                .first()
            )

            if file_index is None:
                file_index = FileIndex(
                    project_id=project_id, path=str(file_path), relative_path=relative_path
                )
                session.add(file_index)

            # Update metadata
            for key, value in metadata.items():
                if hasattr(file_index, key):
                    setattr(file_index, key, value)

            file_index.updated_at = datetime.utcnow()
            session.commit()
            session.refresh(file_index)

            return file_index

    def get_project_files(self, project_id: int) -> list[FileIndex]:
        """Get all indexed files for project"""
        with self.get_session() as session:
            return session.query(FileIndex).filter_by(project_id=project_id).all()

    def search_files(
        self, project_id: int, query: str, language: Optional[str] = None
    ) -> list[FileIndex]:
        """Search indexed files"""
        with self.get_session() as session:
            q = session.query(FileIndex).filter_by(project_id=project_id)

            if language:
                q = q.filter_by(language=language)

            q = q.filter(FileIndex.relative_path.like(f"%{query}%"))

            return q.all()

    # =========================================================================
    # Execution History
    # =========================================================================

    def log_execution(self, project_id: int, task: str, response: str, **metrics) -> Execution:
        """Log RLM execution"""
        with self.get_session() as session:
            # Extract status separately to avoid duplicate keyword argument
            status = metrics.pop("status", "success")
            execution = Execution(
                project_id=project_id,
                task=task,
                response=response,
                status=status,
                **metrics,
            )
            session.add(execution)
            session.commit()
            session.refresh(execution)

            return execution

    def get_recent_executions(self, project_id: int, limit: int = 10) -> list[Execution]:
        """Get recent executions"""
        with self.get_session() as session:
            return (
                session.query(Execution)
                .filter_by(project_id=project_id)
                .order_by(desc(Execution.started_at))
                .limit(limit)
                .all()
            )

    def get_execution_stats(self, project_id: int) -> dict[str, Any]:
        """Get execution statistics"""
        with self.get_session() as session:
            executions = session.query(Execution).filter_by(project_id=project_id).all()

            total_cost = sum(e.total_cost or 0.0 for e in executions)
            total_time = sum(e.total_time or 0.0 for e in executions)
            total_count = len(executions)
            success_count = sum(1 for e in executions if e.status == "success")

            return {
                "total_executions": total_count,
                "successful": success_count,
                "failed": total_count - success_count,
                "total_cost": total_cost,
                "total_time": total_time,
                "avg_cost": total_cost / total_count if total_count > 0 else 0.0,
                "avg_time": total_time / total_count if total_count > 0 else 0.0,
            }

    # =========================================================================
    # Session Management
    # =========================================================================

    def create_session(self, project_id: int, session_type: str) -> Session:
        """Create new session"""
        with self.get_session() as session:
            new_session = Session(project_id=project_id, session_type=session_type)
            session.add(new_session)
            session.commit()
            session.refresh(new_session)

            return new_session

    def end_session(self, session_id: int) -> None:
        """End session"""
        with self.get_session() as session:
            sess = session.query(Session).filter_by(id=session_id).first()
            if sess:
                sess.ended_at = datetime.utcnow()
                session.commit()

    # =========================================================================
    # Analysis Results
    # =========================================================================

    def save_analysis(
        self, project_id: int, analysis_type: str, results: dict[str, Any], **metadata
    ) -> Analysis:
        """Save analysis results"""
        with self.get_session() as session:
            analysis = Analysis(
                project_id=project_id, analysis_type=analysis_type, results=results, **metadata
            )
            session.add(analysis)
            session.commit()
            session.refresh(analysis)

            return analysis

    def get_latest_analysis(self, project_id: int, analysis_type: str) -> Optional[Analysis]:
        """Get latest analysis of given type"""
        with self.get_session() as session:
            return (
                session.query(Analysis)
                .filter_by(project_id=project_id, analysis_type=analysis_type)
                .order_by(desc(Analysis.created_at))
                .first()
            )

    # =========================================================================
    # Context Entry Operations (RLM Persistent Knowledge Base)
    # =========================================================================

    def store_context(
        self,
        project_id: int,
        context_key: str,
        context_type: str,
        content: Any,
        variable_name: Optional[str] = None,
        session_id: Optional[int] = None,
    ) -> ContextEntry:
        """
        Store context entry with deduplication.

        Args:
            project_id: Project ID
            context_key: Key like "context_0", "history_1", "my_var"
            context_type: "context", "history", or "variable"
            content: Content to store (will be JSON serialized)
            variable_name: Original variable name if applicable
            session_id: Optional session ID

        Returns:
            Created or updated ContextEntry
        """
        # Serialize content to JSON
        content_json = json.dumps(content, default=str)
        content_hash = hashlib.sha256(content_json.encode()).hexdigest()

        with self.get_session() as session:
            # Check for existing entry with same key
            existing = (
                session.query(ContextEntry)
                .filter_by(project_id=project_id, context_key=context_key)
                .first()
            )

            if existing:
                # Update existing entry
                existing.content = content_json
                existing.content_hash = content_hash
                existing.accessed_at = datetime.utcnow()
                if variable_name:
                    existing.variable_name = variable_name
                session.commit()
                session.refresh(existing)
                return existing
            else:
                # Create new entry
                entry = ContextEntry(
                    project_id=project_id,
                    session_id=session_id,
                    context_key=context_key,
                    context_type=context_type,
                    content=content_json,
                    content_hash=content_hash,
                    variable_name=variable_name,
                )
                session.add(entry)
                session.commit()
                session.refresh(entry)
                return entry

    def store_contexts_batch(
        self,
        project_id: int,
        entries: list[dict],
        session_id: Optional[int] = None,
    ) -> int:
        """
        Batch store multiple context entries in a single transaction.

        This is significantly faster than calling store_context() individually
        for each entry, as it minimizes transaction overhead.

        Args:
            project_id: Project ID
            entries: List of dicts with keys:
                - context_key: Key like "context_0", "my_var"
                - context_type: "context", "history", or "variable"
                - content: Content to store (will be JSON serialized)
                - variable_name: Optional original variable name
            session_id: Optional session ID

        Returns:
            Number of entries stored/updated
        """
        if not entries:
            return 0

        with self.get_session() as session:
            count = 0
            for entry_data in entries:
                content_json = json.dumps(entry_data["content"], default=str)
                content_hash = hashlib.sha256(content_json.encode()).hexdigest()

                # Check for existing entry
                existing = (
                    session.query(ContextEntry)
                    .filter_by(project_id=project_id, context_key=entry_data["context_key"])
                    .first()
                )

                if existing:
                    # Update existing
                    existing.content = content_json
                    existing.content_hash = content_hash
                    existing.accessed_at = datetime.utcnow()
                    if entry_data.get("variable_name"):
                        existing.variable_name = entry_data["variable_name"]
                else:
                    # Create new entry
                    entry = ContextEntry(
                        project_id=project_id,
                        session_id=session_id,
                        context_key=entry_data["context_key"],
                        context_type=entry_data["context_type"],
                        content=content_json,
                        content_hash=content_hash,
                        variable_name=entry_data.get("variable_name"),
                    )
                    session.add(entry)
                count += 1

            session.commit()
            return count

    def get_context(self, project_id: int, context_key: str) -> Optional[Any]:
        """
        Get context by key and update accessed_at timestamp.

        Args:
            project_id: Project ID
            context_key: Context key to retrieve

        Returns:
            Deserialized content or None if not found
        """
        with self.get_session() as session:
            entry = (
                session.query(ContextEntry)
                .filter_by(project_id=project_id, context_key=context_key)
                .first()
            )

            if entry:
                # Update access time
                entry.accessed_at = datetime.utcnow()
                session.commit()
                return json.loads(entry.content)

            return None

    def get_context_entry(self, project_id: int, context_key: str) -> Optional[ContextEntry]:
        """
        Get full ContextEntry object by key.

        Args:
            project_id: Project ID
            context_key: Context key to retrieve

        Returns:
            ContextEntry or None
        """
        with self.get_session() as session:
            return (
                session.query(ContextEntry)
                .filter_by(project_id=project_id, context_key=context_key)
                .first()
            )

    def search_contexts(
        self,
        project_id: int,
        query: Optional[str] = None,
        context_type: Optional[str] = None,
        limit: int = 100,
    ) -> list[ContextEntry]:
        """
        Search context entries.

        Args:
            project_id: Project ID
            query: Optional search term for key or variable name
            context_type: Optional filter by type ("context", "history", "variable")
            limit: Maximum results to return

        Returns:
            List of matching ContextEntry objects
        """
        with self.get_session() as session:
            q = session.query(ContextEntry).filter_by(project_id=project_id)

            if context_type:
                q = q.filter_by(context_type=context_type)

            if query:
                q = q.filter(
                    or_(
                        ContextEntry.context_key.like(f"%{query}%"),
                        ContextEntry.variable_name.like(f"%{query}%"),
                    )
                )

            return q.order_by(desc(ContextEntry.accessed_at)).limit(limit).all()

    def get_all_contexts(
        self,
        project_id: int,
        context_type: Optional[str] = None,
    ) -> dict[str, Any]:
        """
        Get all contexts as a dictionary for loading into REPL.

        Args:
            project_id: Project ID
            context_type: Optional filter by type

        Returns:
            Dict mapping context_key to deserialized content
        """
        with self.get_session() as session:
            q = session.query(ContextEntry).filter_by(project_id=project_id)

            if context_type:
                q = q.filter_by(context_type=context_type)

            entries = q.all()
            return {entry.context_key: json.loads(entry.content) for entry in entries}

    def get_variables(self, project_id: int) -> dict[str, Any]:
        """
        Get all variables for loading into REPL namespace.

        Args:
            project_id: Project ID

        Returns:
            Dict mapping variable_name to deserialized content
        """
        with self.get_session() as session:
            entries = (
                session.query(ContextEntry)
                .filter_by(project_id=project_id, context_type="variable")
                .all()
            )
            # Use variable_name as key if available, otherwise context_key
            return {
                (entry.variable_name or entry.context_key): json.loads(entry.content)
                for entry in entries
            }

    def delete_context(self, project_id: int, context_key: str) -> bool:
        """
        Delete a context entry.

        Args:
            project_id: Project ID
            context_key: Context key to delete

        Returns:
            True if deleted, False if not found
        """
        with self.get_session() as session:
            entry = (
                session.query(ContextEntry)
                .filter_by(project_id=project_id, context_key=context_key)
                .first()
            )

            if entry:
                session.delete(entry)
                session.commit()
                return True

            return False

    def get_context_stats(self, project_id: int) -> dict[str, Any]:
        """
        Get statistics about stored contexts.

        Args:
            project_id: Project ID

        Returns:
            Stats dict with counts by type
        """
        with self.get_session() as session:
            entries = session.query(ContextEntry).filter_by(project_id=project_id).all()

            type_counts = {"context": 0, "history": 0, "variable": 0}
            total_size = 0

            for entry in entries:
                type_counts[entry.context_type] = type_counts.get(entry.context_type, 0) + 1
                total_size += len(entry.content) if entry.content else 0

            return {
                "total_entries": len(entries),
                "by_type": type_counts,
                "total_size_bytes": total_size,
            }

    # =========================================================================
    # Code Summary Operations (Hybrid Retrieval System)
    # =========================================================================

    def store_summary(
        self,
        project_id: int,
        scope_type: str,
        scope_path: str,
        scope_name: Optional[str] = None,
        summary: Optional[str] = None,
        keywords: Optional[list[str]] = None,
        parent_id: Optional[int] = None,
        signature: Optional[str] = None,
        docstring: Optional[str] = None,
        content_hash: Optional[str] = None,
        line_start: Optional[int] = None,
        line_end: Optional[int] = None,
    ) -> CodeSummary:
        """
        Store or update a code summary.

        Args:
            project_id: Project ID
            scope_type: Type of scope ('package', 'module', 'class', 'function')
            scope_path: Path to the scope (file path for modules)
            scope_name: Name of the scope (class/function name)
            summary: LLM-generated summary
            keywords: Extracted keywords for BM25
            parent_id: Parent summary ID for hierarchy
            signature: Function signature or class inheritance
            docstring: Original docstring
            content_hash: Hash for change detection
            line_start: Starting line number
            line_end: Ending line number

        Returns:
            Created or updated CodeSummary
        """
        with self.get_session() as session:
            # Check for existing summary with same scope
            existing = (
                session.query(CodeSummary)
                .filter_by(
                    project_id=project_id,
                    scope_type=scope_type,
                    scope_path=scope_path,
                    scope_name=scope_name,
                )
                .first()
            )

            if existing:
                # Update existing
                if summary is not None:
                    existing.summary = summary
                if keywords is not None:
                    existing.keywords = keywords
                if parent_id is not None:
                    existing.parent_id = parent_id
                if signature is not None:
                    existing.signature = signature
                if docstring is not None:
                    existing.docstring = docstring
                if content_hash is not None:
                    existing.content_hash = content_hash
                if line_start is not None:
                    existing.line_start = line_start
                if line_end is not None:
                    existing.line_end = line_end
                if summary is not None:
                    existing.last_summarized = datetime.utcnow()
                session.commit()
                session.refresh(existing)
                return existing
            else:
                # Create new
                code_summary = CodeSummary(
                    project_id=project_id,
                    scope_type=scope_type,
                    scope_path=scope_path,
                    scope_name=scope_name,
                    summary=summary,
                    keywords=keywords,
                    parent_id=parent_id,
                    signature=signature,
                    docstring=docstring,
                    content_hash=content_hash,
                    line_start=line_start,
                    line_end=line_end,
                    last_summarized=datetime.utcnow() if summary else None,
                )
                session.add(code_summary)
                session.commit()
                session.refresh(code_summary)
                return code_summary

    def get_summary(
        self,
        project_id: int,
        scope_type: str,
        scope_path: str,
        scope_name: Optional[str] = None,
    ) -> Optional[CodeSummary]:
        """Get a specific code summary."""
        with self.get_session() as session:
            return (
                session.query(CodeSummary)
                .filter_by(
                    project_id=project_id,
                    scope_type=scope_type,
                    scope_path=scope_path,
                    scope_name=scope_name,
                )
                .first()
            )

    def get_summary_by_id(self, summary_id: int) -> Optional[CodeSummary]:
        """Get a code summary by ID."""
        with self.get_session() as session:
            return session.query(CodeSummary).filter_by(id=summary_id).first()

    def get_summaries_by_path(
        self, project_id: int, path_prefix: str
    ) -> list[CodeSummary]:
        """Get all summaries under a path prefix."""
        with self.get_session() as session:
            return (
                session.query(CodeSummary)
                .filter_by(project_id=project_id)
                .filter(CodeSummary.scope_path.like(f"{path_prefix}%"))
                .all()
            )

    def get_summaries_by_type(
        self, project_id: int, scope_type: str
    ) -> list[CodeSummary]:
        """Get all summaries of a specific type."""
        with self.get_session() as session:
            return (
                session.query(CodeSummary)
                .filter_by(project_id=project_id, scope_type=scope_type)
                .all()
            )

    def get_child_summaries(self, parent_id: int) -> list[CodeSummary]:
        """Get all child summaries of a parent."""
        with self.get_session() as session:
            return session.query(CodeSummary).filter_by(parent_id=parent_id).all()

    def get_summaries_needing_update(
        self, project_id: int, content_hashes: dict[str, str]
    ) -> list[CodeSummary]:
        """
        Get summaries that need updating based on content hash changes.

        Args:
            project_id: Project ID
            content_hashes: Dict mapping scope_path to current content hash

        Returns:
            List of summaries with stale content
        """
        with self.get_session() as session:
            summaries = (
                session.query(CodeSummary)
                .filter_by(project_id=project_id, scope_type="module")
                .all()
            )
            stale = []
            for summary in summaries:
                current_hash = content_hashes.get(summary.scope_path)
                if current_hash and summary.content_hash != current_hash:
                    stale.append(summary)
            return stale

    def delete_summaries_for_path(self, project_id: int, scope_path: str) -> int:
        """Delete all summaries for a specific path."""
        with self.get_session() as session:
            count = (
                session.query(CodeSummary)
                .filter_by(project_id=project_id, scope_path=scope_path)
                .delete()
            )
            session.commit()
            return count

    # =========================================================================
    # BM25 Index Operations
    # =========================================================================

    def get_or_create_term(self, project_id: int, term: str) -> CodeTerm:
        """Get or create a term in the index."""
        with self.get_session() as session:
            existing = (
                session.query(CodeTerm)
                .filter_by(project_id=project_id, term=term)
                .first()
            )
            if existing:
                return existing

            code_term = CodeTerm(project_id=project_id, term=term, document_freq=0)
            session.add(code_term)
            session.commit()
            session.refresh(code_term)
            return code_term

    def update_term_document_freq(self, term_id: int, document_freq: int) -> None:
        """Update the document frequency for a term."""
        with self.get_session() as session:
            term = session.query(CodeTerm).filter_by(id=term_id).first()
            if term:
                term.document_freq = document_freq
                session.commit()

    def add_term_occurrence(
        self,
        term_id: int,
        summary_id: int,
        term_freq: int,
        field_type: str,
    ) -> TermOccurrence:
        """Add a term occurrence linking a term to a summary."""
        with self.get_session() as session:
            # Check for existing
            existing = (
                session.query(TermOccurrence)
                .filter_by(term_id=term_id, summary_id=summary_id, field_type=field_type)
                .first()
            )
            if existing:
                existing.term_freq = term_freq
                session.commit()
                session.refresh(existing)
                return existing

            occurrence = TermOccurrence(
                term_id=term_id,
                summary_id=summary_id,
                term_freq=term_freq,
                field_type=field_type,
            )
            session.add(occurrence)
            session.commit()
            session.refresh(occurrence)
            return occurrence

    def search_terms(
        self, project_id: int, terms: list[str]
    ) -> list[tuple[CodeTerm, list[TermOccurrence]]]:
        """
        Search for terms and return their occurrences.

        Args:
            project_id: Project ID
            terms: List of terms to search for

        Returns:
            List of (CodeTerm, [TermOccurrence]) tuples
        """
        with self.get_session() as session:
            results = []
            for term in terms:
                code_term = (
                    session.query(CodeTerm)
                    .filter_by(project_id=project_id, term=term.lower())
                    .first()
                )
                if code_term:
                    occurrences = (
                        session.query(TermOccurrence)
                        .filter_by(term_id=code_term.id)
                        .all()
                    )
                    results.append((code_term, occurrences))
            return results

    def get_total_documents(self, project_id: int) -> int:
        """Get total number of indexed documents (summaries) for BM25."""
        with self.get_session() as session:
            return (
                session.query(CodeSummary)
                .filter_by(project_id=project_id)
                .count()
            )

    def get_average_document_length(self, project_id: int) -> float:
        """Get average document length for BM25 normalization."""
        with self.get_session() as session:
            summaries = (
                session.query(CodeSummary)
                .filter_by(project_id=project_id)
                .all()
            )
            if not summaries:
                return 0.0
            total_length = sum(
                len((s.summary or "") + " " + (s.docstring or ""))
                for s in summaries
            )
            return total_length / len(summaries)

    def delete_term_occurrences_for_summary(self, summary_id: int) -> int:
        """Delete all term occurrences for a summary (before re-indexing)."""
        with self.get_session() as session:
            count = (
                session.query(TermOccurrence)
                .filter_by(summary_id=summary_id)
                .delete()
            )
            session.commit()
            return count

    def clear_bm25_index(self, project_id: int) -> None:
        """Clear all BM25 index data for a project."""
        with self.get_session() as session:
            # Delete occurrences first (foreign key constraint)
            term_ids = [
                t.id for t in session.query(CodeTerm).filter_by(project_id=project_id).all()
            ]
            if term_ids:
                session.query(TermOccurrence).filter(
                    TermOccurrence.term_id.in_(term_ids)
                ).delete(synchronize_session=False)

            # Delete terms
            session.query(CodeTerm).filter_by(project_id=project_id).delete()
            session.commit()

    # =========================================================================
    # Code Relation (Graph) Operations
    # =========================================================================

    def add_relation(
        self,
        source_id: int,
        target_id: int,
        relation_type: str,
    ) -> CodeRelation:
        """Add a relation between two code summaries."""
        with self.get_session() as session:
            # Check for existing
            existing = (
                session.query(CodeRelation)
                .filter_by(
                    source_id=source_id,
                    target_id=target_id,
                    relation_type=relation_type,
                )
                .first()
            )
            if existing:
                return existing

            relation = CodeRelation(
                source_id=source_id,
                target_id=target_id,
                relation_type=relation_type,
            )
            session.add(relation)
            session.commit()
            session.refresh(relation)
            return relation

    def get_outgoing_relations(
        self, summary_id: int, relation_types: Optional[list[str]] = None
    ) -> list[CodeRelation]:
        """Get all outgoing relations from a summary."""
        with self.get_session() as session:
            query = session.query(CodeRelation).filter_by(source_id=summary_id)
            if relation_types:
                query = query.filter(CodeRelation.relation_type.in_(relation_types))
            return query.all()

    def get_incoming_relations(
        self, summary_id: int, relation_types: Optional[list[str]] = None
    ) -> list[CodeRelation]:
        """Get all incoming relations to a summary."""
        with self.get_session() as session:
            query = session.query(CodeRelation).filter_by(target_id=summary_id)
            if relation_types:
                query = query.filter(CodeRelation.relation_type.in_(relation_types))
            return query.all()

    def get_related_summaries(
        self,
        summary_id: int,
        relation_types: Optional[list[str]] = None,
        direction: str = "both",
    ) -> list[CodeSummary]:
        """
        Get summaries related to the given summary.

        Args:
            summary_id: Source summary ID
            relation_types: Filter by relation types (optional)
            direction: 'outgoing', 'incoming', or 'both'

        Returns:
            List of related CodeSummary objects
        """
        with self.get_session() as session:
            related_ids = set()

            if direction in ("outgoing", "both"):
                query = session.query(CodeRelation).filter_by(source_id=summary_id)
                if relation_types:
                    query = query.filter(CodeRelation.relation_type.in_(relation_types))
                for rel in query.all():
                    related_ids.add(rel.target_id)

            if direction in ("incoming", "both"):
                query = session.query(CodeRelation).filter_by(target_id=summary_id)
                if relation_types:
                    query = query.filter(CodeRelation.relation_type.in_(relation_types))
                for rel in query.all():
                    related_ids.add(rel.source_id)

            if not related_ids:
                return []

            return (
                session.query(CodeSummary)
                .filter(CodeSummary.id.in_(related_ids))
                .all()
            )

    def delete_relations_for_summary(self, summary_id: int) -> int:
        """Delete all relations involving a summary."""
        with self.get_session() as session:
            count = (
                session.query(CodeRelation)
                .filter(
                    (CodeRelation.source_id == summary_id)
                    | (CodeRelation.target_id == summary_id)
                )
                .delete(synchronize_session=False)
            )
            session.commit()
            return count

    def clear_relations(self, project_id: int) -> None:
        """Clear all relations for a project."""
        with self.get_session() as session:
            # Get all summary IDs for this project
            summary_ids = [
                s.id for s in session.query(CodeSummary).filter_by(project_id=project_id).all()
            ]
            if summary_ids:
                session.query(CodeRelation).filter(
                    CodeRelation.source_id.in_(summary_ids)
                ).delete(synchronize_session=False)
            session.commit()
