"""Query helpers for structured store JSONB operations."""

from typing import Any
from uuid import UUID

import structlog
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from src.db.models import ExtractedRecord

logger = structlog.get_logger()


class StructuredQueryHelper:
    """Helper class for querying stored records with JSONB operations.

    Provides optimized query methods leveraging PostgreSQL GIN indexes
    for fast JSONB field lookups.
    """

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

        Args:
            db_session: Async SQLAlchemy session for database operations
        """
        self.db = db_session
        self.logger = logger.bind(component="structured_query_helper")

    async def find_by_exact_field(
        self,
        document_id: UUID,
        field_name: str,
        field_value: str,
    ) -> list[ExtractedRecord]:
        """Find records by exact field value match.

        Uses GIN index on content column for fast lookups.
        Example: Find all records where content->>'Item Code' = '2024-4_0019'

        Args:
            document_id: UUID of the document
            field_name: JSON field name (e.g., "Item Code")
            field_value: Exact value to match

        Returns:
            List of matching ExtractedRecord objects
        """
        self.logger.debug(
            "exact_field_query",
            document_id=str(document_id),
            field=field_name,
            value=field_value,
        )

        stmt = select(ExtractedRecord).where(
            ExtractedRecord.document_id == document_id,
            ExtractedRecord.content[field_name].astext == field_value,
        )

        result = await self.db.execute(stmt)
        records = result.scalars().all()

        self.logger.info(
            "exact_field_query_completed",
            document_id=str(document_id),
            field=field_name,
            matches=len(records),
        )

        return records

    async def find_by_containment(
        self,
        document_id: UUID,
        search_dict: dict[str, Any],
    ) -> list[ExtractedRecord]:
        """Find records where content contains the search dictionary.

        Uses JSONB @> containment operator with GIN index.
        Example: Find all records where content @> '{"Screen": "○"}'

        Args:
            document_id: UUID of the document
            search_dict: Dictionary to match (e.g., {"Screen": "○"})

        Returns:
            List of matching ExtractedRecord objects
        """
        self.logger.debug(
            "containment_query",
            document_id=str(document_id),
            search_dict=search_dict,
        )

        stmt = select(ExtractedRecord).where(
            ExtractedRecord.document_id == document_id,
            ExtractedRecord.content.contains(search_dict),
        )

        result = await self.db.execute(stmt)
        records = result.scalars().all()

        self.logger.info(
            "containment_query_completed",
            document_id=str(document_id),
            matches=len(records),
        )

        return records

    async def find_by_sheet(
        self,
        document_id: UUID,
        sheet_name: str,
    ) -> list[ExtractedRecord]:
        """Find all records from a specific sheet.

        Uses B-tree index on sheet_name for fast filtering.

        Args:
            document_id: UUID of the document
            sheet_name: Name of the sheet

        Returns:
            List of ExtractedRecord objects from the sheet
        """
        self.logger.debug(
            "sheet_query",
            document_id=str(document_id),
            sheet=sheet_name,
        )

        stmt = select(ExtractedRecord).where(
            ExtractedRecord.document_id == document_id,
            ExtractedRecord.sheet_name == sheet_name,
        )

        result = await self.db.execute(stmt)
        records = result.scalars().all()

        self.logger.info(
            "sheet_query_completed",
            document_id=str(document_id),
            sheet=sheet_name,
            matches=len(records),
        )

        return records

    async def find_by_field_exists(
        self,
        document_id: UUID,
        field_name: str,
    ) -> list[ExtractedRecord]:
        """Find records where a specific field exists.

        Uses JSONB ? operator to check field existence.

        Args:
            document_id: UUID of the document
            field_name: JSON field name to check

        Returns:
            List of ExtractedRecord objects containing the field
        """
        self.logger.debug(
            "field_exists_query",
            document_id=str(document_id),
            field=field_name,
        )

        stmt = select(ExtractedRecord).where(
            ExtractedRecord.document_id == document_id,
            ExtractedRecord.content.has_key(field_name),
        )

        result = await self.db.execute(stmt)
        records = result.scalars().all()

        self.logger.info(
            "field_exists_query_completed",
            document_id=str(document_id),
            field=field_name,
            matches=len(records),
        )

        return records

    async def find_with_resolved_content(
        self,
        document_id: UUID,
    ) -> list[ExtractedRecord]:
        """Find records that have resolved content (after symbol resolution).

        Args:
            document_id: UUID of the document

        Returns:
            List of ExtractedRecord objects with resolved_content
        """
        stmt = select(ExtractedRecord).where(
            ExtractedRecord.document_id == document_id,
            ExtractedRecord.resolved_content.isnot(None),
        )

        result = await self.db.execute(stmt)
        return result.scalars().all()

    async def count_records(
        self,
        document_id: UUID,
        sheet_name: str | None = None,
    ) -> int:
        """Count records for a document, optionally filtered by sheet.

        Args:
            document_id: UUID of the document
            sheet_name: Optional sheet name filter

        Returns:
            Count of matching records
        """
        from sqlalchemy import func

        stmt = select(func.count()).select_from(ExtractedRecord).where(
            ExtractedRecord.document_id == document_id
        )

        if sheet_name:
            stmt = stmt.where(ExtractedRecord.sheet_name == sheet_name)

        result = await self.db.execute(stmt)
        return result.scalar_one()