"""Service for orchestrating document deletion across PostgreSQL and Milvus."""

from uuid import UUID

import structlog
from sqlalchemy.ext.asyncio import AsyncSession

from src.core.exceptions import IndexingError
from src.extraction.models import ExtractedRecord
from src.knowledge.structured_store import StructuredStore
from src.knowledge.vector_store import VectorStore, delete_vectors_by_document_id

logger = structlog.get_logger(__name__)


class DocumentDeletionService:
    """Orchestrates deletion operations across PostgreSQL and Milvus.

    Ensures atomic deletion - both stores are updated or neither is changed.
    Handles rollback on failure to maintain data consistency.
    """

    def __init__(self, db_session: AsyncSession):
        """Initialize DocumentDeletionService.

        Args:
            db_session: Async SQLAlchemy session for database operations
        """
        self.db_session = db_session
        self.structured_store = StructuredStore(db_session)
        self.logger = logger.bind(component="document_deletion_service")

    async def delete_document(self, document_id: UUID) -> dict[str, int]:
        """Delete document and all associated records atomically.

        Deletes from both PostgreSQL and Milvus within a database transaction.
        If Milvus deletion fails, PostgreSQL transaction is rolled back.

        Args:
            document_id: Document UUID to delete

        Returns:
            Dict with deletion counts: {"records": int, "vectors": int}

        Raises:
            IndexingError: If deletion fails
        """
        try:
            # Start transaction
            async with self.db_session.begin():
                # Delete from PostgreSQL (CASCADE to records)
                records_deleted = await self.structured_store.delete_records_by_document_id(
                    document_id
                )

                # Delete from Milvus
                try:
                    vectors_deleted = await delete_vectors_by_document_id(document_id)
                except Exception as milvus_error:
                    # Milvus deletion failed - transaction will auto-rollback
                    self.logger.error(
                        "milvus_deletion_failed_rolling_back",
                        document_id=str(document_id),
                        error=str(milvus_error),
                    )
                    raise IndexingError(
                        f"Milvus deletion failed, rolling back: {milvus_error}"
                    ) from milvus_error

                self.logger.info(
                    "document_deleted",
                    document_id=str(document_id),
                    records_deleted=records_deleted,
                    vectors_deleted=vectors_deleted,
                )

                return {
                    "records": records_deleted,
                    "vectors": vectors_deleted,
                }

        except Exception as e:
            # Transaction will auto-rollback
            self.logger.error(
                "document_deletion_failed",
                document_id=str(document_id),
                error=str(e),
            )
            raise IndexingError(f"Failed to delete document: {e}") from e

    async def reindex_document(
        self,
        document_id: UUID,
        records: list[ExtractedRecord],
    ) -> dict[str, int]:
        """Re-index document by deleting old and inserting new records.

        Performs atomic operation: delete old → insert new.
        If document doesn't exist, performs initial indexing.

        Args:
            document_id: Document UUID
            records: New extracted records to index

        Returns:
            Dict with operation counts:
            {
                "deleted_records": int,
                "deleted_vectors": int,
                "indexed_records": int,
                "indexed_vectors": int
            }

        Raises:
            IndexingError: If reindexing fails
        """
        try:
            # Delete old records/vectors (if they exist)
            deletion_result = await self.delete_document(document_id)

            self.logger.info(
                "reindex_deletion_complete",
                document_id=str(document_id),
                deleted_records=deletion_result["records"],
                deleted_vectors=deletion_result["vectors"],
            )

            # Index new records in PostgreSQL
            records_count = await self.structured_store.store_records(records, document_id)

            # Query the record IDs that were just inserted
            from sqlalchemy import select
            from src.db.models import ExtractedRecord as ExtractedRecordModel

            result = await self.db_session.execute(
                select(ExtractedRecordModel.id)
                .where(ExtractedRecordModel.document_id == document_id)
                .order_by(ExtractedRecordModel.row_number)
            )
            record_ids = [row[0] for row in result.fetchall()]

            # Index new vectors in Milvus
            vector_store = VectorStore()
            await vector_store.create_collection()  # Idempotent
            vector_count = await vector_store.index_records(
                records, document_id, record_ids
            )

            self.logger.info(
                "reindex_complete",
                document_id=str(document_id),
                deleted_records=deletion_result["records"],
                deleted_vectors=deletion_result["vectors"],
                indexed_records=len(record_ids),
                indexed_vectors=vector_count,
            )

            return {
                "deleted_records": deletion_result["records"],
                "deleted_vectors": deletion_result["vectors"],
                "indexed_records": len(record_ids),
                "indexed_vectors": vector_count,
            }

        except Exception as e:
            self.logger.error(
                "reindex_failed",
                document_id=str(document_id),
                error=str(e),
            )
            raise IndexingError(f"Failed to reindex document: {e}") from e