"""PostgreSQL storage for extracted records."""

import asyncio
import time
from dataclasses import dataclass
from typing import Any
from uuid import UUID, uuid4

import structlog
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.exc import IntegrityError, OperationalError
from sqlalchemy.ext.asyncio import AsyncSession

from src.core.exceptions import IndexingError
from src.db.models import ExtractedRecord as ExtractedRecordModel
from src.extraction.models import ExtractedRecord


@dataclass
class StoreResult:
    """Result from storing records in PostgreSQL.

    Attributes:
        count: Number of records successfully stored
        record_ids: List of UUIDs for each stored record (in same order as input)
    """

    count: int
    record_ids: list[UUID]


# Constants
MAX_BATCH_SIZE = 500
MAX_RETRIES = 3
RETRY_BACKOFF_SECONDS = [1, 2, 4]  # Exponential backoff: 1s, 2s, 4s

logger = structlog.get_logger()


class StructuredStore:
    """Manages PostgreSQL storage for extracted records.

    Provides optimized batch insertion of extracted records with:
    - JSONB storage for content, resolved_content, and headers
    - Batch processing (500 records per transaction)
    - ON CONFLICT handling for idempotent operations
    - Retry logic for transient database errors
    - Performance logging
    """

    def __init__(self, db_session: AsyncSession, tenant_id: UUID | None = None):
        """Initialize StructuredStore.

        Args:
            db_session: Async SQLAlchemy session for database operations
            tenant_id: Optional tenant UUID for RLS policy
        """
        self.db = db_session
        self.tenant_id = tenant_id
        self.logger = logger.bind(component="structured_store")

    async def _set_rls_context(self) -> None:
        """Set RLS tenant context on the session if tenant_id is available."""
        if self.tenant_id:
            from sqlalchemy import text
            from uuid import UUID as _UUID
            _tid = str(_UUID(str(self.tenant_id)))  # validate UUID format
            await self.db.execute(text(f"SET LOCAL app.tenant_id = '{_tid}'"))

    async def store_records(
        self,
        records: list[ExtractedRecord],
        document_id: UUID,
    ) -> StoreResult:
        """Store extracted records in PostgreSQL.

        Processes records in batches of 500, using ON CONFLICT for idempotency.
        Includes retry logic for transient database errors.

        Args:
            records: List of ExtractedRecord from extraction pipeline
            document_id: UUID of source document

        Returns:
            StoreResult with count and list of record UUIDs (for Milvus linking)

        Raises:
            IndexingError: If storage fails after retries
        """
        if not records:
            self.logger.warning("store_records_called_with_empty_list", document_id=str(document_id))
            return StoreResult(count=0, record_ids=[])

        start_time = time.time()
        total_stored = 0
        all_record_ids: list[UUID] = []

        self.logger.info(
            "store_records_started",
            document_id=str(document_id),
            record_count=len(records),
        )

        try:
            # Split records into batches
            batches = [
                records[i : i + MAX_BATCH_SIZE]
                for i in range(0, len(records), MAX_BATCH_SIZE)
            ]

            for batch_idx, batch in enumerate(batches):
                batch_start = time.time()

                # Store batch with retry logic - returns (count, ids)
                stored_count, batch_ids = await self._store_batch_with_retry(
                    batch, document_id, batch_idx
                )
                total_stored += stored_count
                all_record_ids.extend(batch_ids)

                batch_duration = (time.time() - batch_start) * 1000
                self.logger.debug(
                    "batch_stored",
                    batch_idx=batch_idx,
                    batch_size=len(batch),
                    stored_count=stored_count,
                    duration_ms=int(batch_duration),
                )

            # Log final metrics
            total_duration = (time.time() - start_time) * 1000
            records_per_second = (
                len(records) / (total_duration / 1000) if total_duration > 0 else 0
            )

            self.logger.info(
                "store_records_completed",
                document_id=str(document_id),
                records_stored=total_stored,
                record_ids_count=len(all_record_ids),
                duration_ms=int(total_duration),
                records_per_second=int(records_per_second),
            )

            return StoreResult(count=total_stored, record_ids=all_record_ids)

        except Exception as e:
            duration = (time.time() - start_time) * 1000
            self.logger.error(
                "store_records_failed",
                document_id=str(document_id),
                error=str(e),
                error_type=type(e).__name__,
                duration_ms=int(duration),
            )
            raise IndexingError(
                message=f"Failed to store records for document {document_id}",
                details={
                    "document_id": str(document_id),
                    "error": str(e),
                    "error_type": type(e).__name__,
                },
            ) from e

    async def _store_batch_with_retry(
        self,
        batch: list[ExtractedRecord],
        document_id: UUID,
        batch_idx: int,
    ) -> tuple[int, list[UUID]]:
        """Store a batch of records with retry logic for transient errors.

        Args:
            batch: Batch of records to store
            document_id: UUID of source document
            batch_idx: Index of this batch (for logging)

        Returns:
            Tuple of (count of stored records, list of record UUIDs)

        Raises:
            IndexingError: If storage fails after all retries
        """
        for attempt in range(MAX_RETRIES):
            try:
                return await self._store_batch(batch, document_id)

            except OperationalError as e:
                # Transient database error (connection loss, timeout)
                if attempt < MAX_RETRIES - 1:
                    backoff = RETRY_BACKOFF_SECONDS[attempt]
                    self.logger.warning(
                        "batch_store_retry",
                        batch_idx=batch_idx,
                        attempt=attempt + 1,
                        max_retries=MAX_RETRIES,
                        error=str(e),
                        retry_after_seconds=backoff,
                    )
                    await asyncio.sleep(backoff)
                else:
                    # Final attempt failed
                    self.logger.error(
                        "batch_store_failed_after_retries",
                        batch_idx=batch_idx,
                        attempts=MAX_RETRIES,
                        error=str(e),
                    )
                    raise

            except IntegrityError as e:
                # Permanent error (data integrity issue)
                self.logger.error(
                    "batch_store_integrity_error",
                    batch_idx=batch_idx,
                    error=str(e),
                )
                raise

            except Exception as e:
                # Unexpected error
                self.logger.error(
                    "batch_store_unexpected_error",
                    batch_idx=batch_idx,
                    error=str(e),
                    error_type=type(e).__name__,
                )
                raise

        # Should never reach here (covered by raises above)
        return 0, []

    async def _store_batch(
        self,
        batch: list[ExtractedRecord],
        document_id: UUID,
    ) -> tuple[int, list[UUID]]:
        """Store a single batch of records in a transaction.

        Args:
            batch: Batch of records to store
            document_id: UUID of source document

        Returns:
            Tuple of (count of stored records, list of record UUIDs)

        Raises:
            SQLAlchemy exceptions on database errors
        """
        # Generate UUIDs upfront so we can return them for Milvus linking
        record_ids = [uuid4() for _ in batch]

        # Convert ExtractedRecord objects to database model dicts with pre-generated IDs
        db_records = [
            self._to_db_model(record, document_id, record_id)
            for record, record_id in zip(batch, record_ids)
        ]

        # Build INSERT statement
        stmt = pg_insert(ExtractedRecordModel).values(db_records)

        # Execute within transaction (only start new transaction if not already in one)
        if not self.db.in_transaction():
            async with self.db.begin():
                await self._set_rls_context()
                result = await self.db.execute(stmt)
                return result.rowcount, record_ids
        else:
            # Already in a transaction, just execute
            await self._set_rls_context()
            result = await self.db.execute(stmt)
            return result.rowcount, record_ids

    def _to_db_model(
        self,
        extracted_record: ExtractedRecord,
        document_id: UUID,
        record_id: UUID,
    ) -> dict[str, Any]:
        """Convert ExtractedRecord to database model dict.

        Args:
            extracted_record: Extraction pipeline record
            document_id: UUID of source document
            record_id: Pre-generated UUID for this record (for Milvus linking)

        Returns:
            Dictionary with database model fields
        """
        # Extract source metadata
        source = extracted_record._source

        # Build database record with explicit ID for Milvus linking
        db_record = {
            "id": record_id,
            "document_id": document_id,
            "sheet_name": source["sheet"],
            "row_number": source["row"],
            "col_range": source.get("col_range"),
            "content": extracted_record.content,
            "headers": extracted_record.headers,
        }

        # Add resolved_content if present (from symbol resolution)
        # Note: ExtractedRecord doesn't have resolved_content attribute yet
        # This will be added in Story 3.6 (Symbol Resolution)
        # For now, we check if it exists via getattr
        resolved_content = getattr(extracted_record, "resolved_content", None)
        if resolved_content:
            db_record["resolved_content"] = resolved_content

        # Add tags if present (from SemanticTagger)
        # Tags are TagAssignment objects or dicts, serialize to JSON-compatible format
        tags = getattr(extracted_record, "tags", None)
        if tags:
            serialized_tags = []
            for tag in tags:
                if hasattr(tag, "model_dump"):
                    # Pydantic v2 model
                    serialized_tags.append(tag.model_dump())
                elif hasattr(tag, "dict"):
                    # Pydantic v1 model
                    serialized_tags.append(tag.dict())
                elif isinstance(tag, dict):
                    # Already a dict
                    serialized_tags.append(tag)
                else:
                    # String or other - skip
                    self.logger.warning(
                        "skipping_invalid_tag_type",
                        tag_type=type(tag).__name__,
                    )
            db_record["tags"] = serialized_tags
        else:
            db_record["tags"] = []

        return db_record

    async def delete_records_by_document_id(self, document_id: UUID) -> int:
        """Delete all records for a document using CASCADE.

        Args:
            document_id: Document UUID to filter records

        Returns:
            Count of deleted records

        Raises:
            IndexingError: If deletion fails
        """
        try:
            from sqlalchemy import delete

            # Build DELETE statement
            stmt = delete(ExtractedRecordModel).where(
                ExtractedRecordModel.document_id == document_id
            )

            # Execute deletion (CASCADE will handle dependencies)
            if not self.db.in_transaction():
                async with self.db.begin():
                    await self._set_rls_context()
                    result = await self.db.execute(stmt)
                    deleted_count = result.rowcount
            else:
                # Already in a transaction
                await self._set_rls_context()
                result = await self.db.execute(stmt)
                deleted_count = result.rowcount

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

            return deleted_count

        except Exception as e:
            self.logger.error(
                "records_deletion_failed",
                document_id=str(document_id),
                error=str(e),
                error_type=type(e).__name__,
            )
            raise IndexingError(
                message=f"Failed to delete records for document {document_id}",
                details={
                    "document_id": str(document_id),
                    "error": str(e),
                    "error_type": type(e).__name__,
                },
            ) from e
