"""Query executor that integrates ChromaSQL with AsyncMultiCollectionQueryClient."""

from __future__ import annotations

import logging
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional

from chromasql import parse
from chromasql.adapters import MetadataFieldRouter
from chromasql.multi_collection import execute_multi_collection
from idxr.vectorize_lib.query_client import AsyncMultiCollectionQueryClient
from idxr.vectorize_lib.query_config import load_query_config

from .async_multi_collection_adapter import AsyncMultiCollectionAdapter

logger = logging.getLogger(__name__)


@dataclass
class QueryResult:
    """Result from executing a ChromaSQL query.

    Attributes
    ----------
    rows:
        List of result rows, each a dictionary mapping field names to values.
    total_rows:
        Total number of rows returned.
    collections_queried:
        Number of collections that were queried.
    collection_names:
        Optional list of collection names that were queried.
        Useful for multi-collection queries to track which collections returned results.
    query:
        The original ChromaSQL query string.
    """

    rows: List[Dict[str, Any]]
    total_rows: int
    collections_queried: int
    query: str
    collection_names: Optional[List[str]] = None


class QueryExecutor:
    """Executor for running ChromaSQL queries against multi-collection ChromaDB.

    This class integrates ChromaSQL's query language with the multi-collection
    infrastructure, enabling SQL-like queries across partitioned data.

    Parameters
    ----------
    query_config_path:
        Path to query_config.json (maps models to collections).
    discriminator_field:
        Metadata field used for routing (e.g., "model_name", "tenant_id").
    client_type:
        Type of ChromaDB client ("http", "cloud", "persistent", or "memory").
    http_host:
        Host for HTTP client.
    http_port:
        Port for HTTP client.
    cloud_api_key:
        API key for Chroma Cloud.
    cloud_tenant:
        Tenant ID for Chroma Cloud.
    cloud_database:
        Database name for Chroma Cloud.
    persist_directory:
        Directory path for persistent local ChromaDB storage.
    embed_fn:
        Optional embedding function for TEXT queries.
    fallback_to_all:
        If True, query all collections when discriminator not specified.

    Example
    -------
    >>> async with QueryExecutor(
    ...     query_config_path=Path("query_config.json"),
    ...     discriminator_field="model_name",
    ...     client_type="memory",
    ... ) as executor:
    ...     result = await executor.execute(
    ...         "SELECT id, distance FROM demo "
    ...         "WHERE metadata.model_name = 'Table' "
    ...         "USING EMBEDDING (TEXT 'search query') "
    ...         "TOPK 10;"
    ...     )
    ...     print(f"Found {result.total_rows} results")
    """

    def __init__(
        self,
        query_config_path: Path,
        discriminator_field: str = "model_name",
        *,
        client_type: str = "http",
        http_host: Optional[str] = None,
        http_port: Optional[int] = None,
        http_ssl: bool = False,
        cloud_api_key: Optional[str] = None,
        cloud_tenant: str = "default_tenant",
        cloud_database: str = "default_database",
        persist_directory: Optional[str] = None,
        embed_fn: Optional[Any] = None,
        fallback_to_all: bool = True,
    ):
        """Initialize the query executor."""
        self.query_config_path = query_config_path
        self.discriminator_field = discriminator_field
        self.client_type = client_type
        self.http_host = http_host
        self.http_port = http_port or 8000
        self.http_ssl = http_ssl
        self.cloud_api_key = cloud_api_key
        self.cloud_tenant = cloud_tenant
        self.cloud_database = cloud_database
        self.persist_directory = persist_directory
        self.embed_fn = embed_fn
        self.fallback_to_all = fallback_to_all

        # Runtime state
        self._client: Optional[AsyncMultiCollectionQueryClient] = None
        self._adapter: Optional[AsyncMultiCollectionAdapter] = None
        self._router: Optional[MetadataFieldRouter] = None
        self._query_config: Optional[Dict[str, Any]] = None

    async def connect(self) -> None:
        """Connect to ChromaDB and initialize routing infrastructure."""
        # Load query config
        self._query_config = load_query_config(self.query_config_path)
        logger.info(
            "Loaded query config: %d models, %d collections",
            self._query_config["metadata"]["total_models"],
            self._query_config["metadata"]["total_collections"],
        )

        # Initialize ChromaDB client
        if self.client_type == "memory":
            # For testing with in-memory ChromaDB
            import chromadb

            self._client = chromadb.Client()  # type: ignore[assignment]
            logger.info("Using in-memory ChromaDB client")
        elif self.client_type == "persistent":
            # For local persistent ChromaDB
            import chromadb

            if not self.persist_directory:
                raise ValueError(
                    "persist_directory is required for persistent client type"
                )

            self._client = chromadb.PersistentClient(
                path=self.persist_directory
            )  # type: ignore[assignment]
            logger.info(
                "Using persistent ChromaDB client at %s", self.persist_directory
            )
        else:
            # Use AsyncMultiCollectionQueryClient for production
            self._client = AsyncMultiCollectionQueryClient(
                config_path=self.query_config_path,
                client_type=self.client_type,
                http_host=self.http_host,
                http_port=self.http_port,
                http_ssl=self.http_ssl,
                cloud_api_key=self.cloud_api_key,
                cloud_tenant=self.cloud_tenant,
                cloud_database=self.cloud_database,
            )
            await self._client.connect()
            logger.info("Connected to ChromaDB (%s)", self.client_type)

        # Create adapters for ChromaSQL
        if self.client_type in ("memory", "persistent"):
            # For sync clients (memory and persistent), create a sync-to-async wrapper
            from .sync_adapter import SyncClientAdapter

            collection_names = list(self._query_config["collection_to_models"].keys())
            self._adapter = SyncClientAdapter(  # type: ignore[assignment]
                client=self._client, collection_names=collection_names
            )
        else:
            self._adapter = AsyncMultiCollectionAdapter(self._client)

        # Create router
        self._router = MetadataFieldRouter(
            query_config=self._query_config,
            field_path=(self.discriminator_field,),
            fallback_to_all=self.fallback_to_all,
        )

        logger.info(
            "Query executor ready (discriminator=%s, fallback=%s)",
            self.discriminator_field,
            self.fallback_to_all,
        )

    async def close(self) -> None:
        """Close ChromaDB connection."""
        if self._client is not None:
            if hasattr(self._client, "close"):
                await self._client.close()
            self._client = None
            self._adapter = None
            self._router = None
            logger.info("Closed ChromaDB connection")

    async def execute(self, cql: str) -> QueryResult:
        """Execute a ChromaSQL query.

        Parameters
        ----------
        cql:
            ChromaSQL query string.

        Returns
        -------
        QueryResult
            Query execution results.

        Raises
        ------
        RuntimeError:
            If executor not connected.
        ChromaSQLError:
            If query parsing or execution fails.
        """
        if (
            self._client is None
            or self._adapter is None
            or self._router is None
            or self._query_config is None
        ):
            raise RuntimeError("Executor not connected. Call await executor.connect()")

        logger.info("Executing ChromaSQL query")
        logger.debug("Query: %s", cql)

        # Preview routing
        query_ast = parse(cql)
        collections = self._router.route(query_ast)

        if collections is None:
            # type: ignore[index]
            num_collections = len(self._query_config["collection_to_models"])
            logger.info("Routing: ALL collections (%d total)", num_collections)
        else:
            logger.info(
                "Routing: %d collection(s): %s", len(collections), list(collections)[:5]
            )

        # Execute query
        result = await execute_multi_collection(
            query_str=cql,
            router=self._router,
            collection_provider=self._adapter,
            embed_fn=self.embed_fn,
        )

        logger.info("Query returned %d row(s)", len(result.rows))

        # Convert collections to list (if it's a set, list, etc.)
        collection_names_list: Optional[List[str]] = None
        if collections is not None:
            collection_names_list = list(collections)
        elif self._query_config is not None:
            # All collections were queried
            collection_names_list = list(
                self._query_config["collection_to_models"].keys()
            )

        return QueryResult(
            rows=result.rows,
            total_rows=len(result.rows),
            collections_queried=(
                num_collections if collections is None else len(collections)
            ),
            query=cql,
            collection_names=collection_names_list,
        )

    async def __aenter__(self) -> QueryExecutor:
        """Async context manager entry."""
        await self.connect()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        """Async context manager exit."""
        await self.close()
