"""Source document management."""
from __future__ import annotations

import json
import sqlite3
from datetime import UTC, datetime, timedelta
from typing import TYPE_CHECKING, Any

from .converters import row_to_source_document
from .models import SourceDocumentSnapshot

if TYPE_CHECKING:
    from .engine import RetrievalLearningEngine


class SourceDocumentsMixin:
    """Mixin for source document functionality."""

    @staticmethod
    def register_source_document(
        engine: "RetrievalLearningEngine",
        *,
        workspace_id: str | None,
        source_key: str,
        document_id: str,
        content: str,
        metadata: dict[str, Any],
        source_ref: str | None = None,
    ) -> SourceDocumentSnapshot:
        """Register a source document snapshot."""
        snapshot = SourceDocumentSnapshot(
            workspace_id=workspace_id,
            source_key=source_key,
            document_id=document_id,
            content=content,
            metadata=dict(metadata),
            source_ref=source_ref,
        )
        with engine._connect() as connection:
            engine._ensure_schema(connection)
            if engine.dsn:
                connection.execute(
                    f"""
                    INSERT INTO {engine._table('source_documents')}
                    (workspace_id, source_key, document_id, content, metadata_json, source_ref, updated_at)
                    VALUES (?, ?, ?, ?, ?, ?, ?)
                    ON CONFLICT (workspace_id, source_key, document_id)
                    DO UPDATE SET
                        content = EXCLUDED.content,
                        metadata_json = EXCLUDED.metadata_json,
                        source_ref = EXCLUDED.source_ref,
                        updated_at = EXCLUDED.updated_at
                    """,
                    (
                        engine._workspace_value(workspace_id),
                        source_key,
                        document_id,
                        content,
                        json.dumps(snapshot.metadata, sort_keys=True),
                        source_ref,
                        snapshot.updated_at,
                    ),
                )
            else:
                connection.execute(
                    f"""
                    INSERT OR REPLACE INTO {engine._table('source_documents')}
                    (workspace_id, source_key, document_id, content, metadata_json, source_ref, updated_at)
                    VALUES (?, ?, ?, ?, ?, ?, ?)
                    """,
                    (
                        engine._workspace_value(workspace_id),
                        source_key,
                        document_id,
                        content,
                        json.dumps(snapshot.metadata, sort_keys=True),
                        source_ref,
                        snapshot.updated_at,
                    ),
                )
                SourceDocumentsMixin._prune_source_documents(
                    connection,
                    engine=engine,
                    workspace_key=engine._workspace_value(workspace_id),
                    source_key=source_key,
                )
            connection.commit()
        return snapshot

    @staticmethod
    def list_source_documents(
        engine: "RetrievalLearningEngine",
        *,
        workspace_id: str | None = None,
        source_key: str | None = None,
    ) -> list[SourceDocumentSnapshot]:
        """List source document snapshots."""
        clauses = ["workspace_id = ?"]
        params: list[Any] = [engine._workspace_value(workspace_id)]
        if source_key is not None:
            clauses.append("source_key = ?")
            params.append(source_key)
        with engine._connect() as connection:
            engine._ensure_schema(connection)
            rows = connection.execute(
                f"""
                SELECT * FROM {engine._table('source_documents')}
                WHERE {' AND '.join(clauses)}
                ORDER BY updated_at DESC
                """,
                params,
            ).fetchall()
        return [row_to_source_document(row, workspace_id) for row in rows]

    @staticmethod
    def prune_source_documents(
        engine: "RetrievalLearningEngine",
        *,
        workspace_id: str | None = None,
        source_key: str | None = None,
        ttl_days: int | None = None,
        max_documents_per_source: int | None = None,
    ) -> int:
        """Prune old source documents."""
        deleted = 0
        with engine._connect() as connection:
            engine._ensure_schema(connection)
            deleted += SourceDocumentsMixin._prune_source_documents(
                connection,
                engine=engine,
                workspace_key=engine._workspace_value(workspace_id) if workspace_id is not None else None,
                source_key=source_key,
                ttl_days=ttl_days,
                max_documents_per_source=max_documents_per_source,
            )
            connection.commit()
        return deleted

    @staticmethod
    def _prune_source_documents(
        connection: sqlite3.Connection,
        *,
        engine: "RetrievalLearningEngine",
        workspace_key: str | None,
        source_key: str | None,
        ttl_days: int | None = None,
        max_documents_per_source: int | None = None,
    ) -> int:
        """Internal prune source documents."""
        deleted = 0
        effective_ttl_days = ttl_days if ttl_days is not None else engine.source_document_ttl_days
        effective_max = max_documents_per_source if max_documents_per_source is not None else engine.max_source_documents_per_source
        
        clauses: list[str] = []
        params: list[Any] = []
        if workspace_key is not None:
            clauses.append("workspace_id = ?")
            params.append(workspace_key)
        if source_key is not None:
            clauses.append("source_key = ?")
            params.append(source_key)
        
        where_sql = f"WHERE {' AND '.join(clauses)}" if clauses else ""
        cutoff = (datetime.now(UTC) - timedelta(days=effective_ttl_days)).isoformat()
        
        # Delete stale documents
        stale_rows = connection.execute(
            f"""
            SELECT workspace_id, source_key, document_id
            FROM {engine._table('source_documents')}
            {where_sql} {"AND" if where_sql else "WHERE"} updated_at < ?
            """,
            [*params, cutoff],
        ).fetchall()
        for row in stale_rows:
            connection.execute(
                f"""
                DELETE FROM {engine._table('source_documents')}
                WHERE workspace_id = ? AND source_key = ? AND document_id = ?
                """,
                (row["workspace_id"], row["source_key"], row["document_id"]),
            )
            deleted += 1

        # Enforce max documents per source
        grouped_rows = connection.execute(
            f"""
            SELECT workspace_id, source_key, document_id
            FROM {engine._table('source_documents')}
            {where_sql}
            ORDER BY workspace_id ASC, source_key ASC, updated_at DESC
            """,
            params,
        ).fetchall()
        seen: dict[tuple[str, str], int] = {}
        for row in grouped_rows:
            key = (row["workspace_id"], row["source_key"])
            seen[key] = seen.get(key, 0) + 1
            if seen[key] <= effective_max:
                continue
            connection.execute(
                f"""
                DELETE FROM {engine._table('source_documents')}
                WHERE workspace_id = ? AND source_key = ? AND document_id = ?
                """,
                (row["workspace_id"], row["source_key"], row["document_id"]),
            )
            deleted += 1
        return deleted
