"""
Query executor configuration utilities for standalone idxr usage.

This module provides helper functions to build QueryExecutor configurations
from CollectionEnvironment objects.
"""

from __future__ import annotations

from typing import Any, Callable, Dict, Mapping, Optional, Sequence

from idxr.load_model_registry import load_model_registry as load_registry_target
from idxr.models import CollectionEnvironment, ModelSpec


class QueryExecutorConfigError(ValueError):
    """Raised when a CollectionEnvironment is missing required settings."""


def create_openai_embed_function(
    model: str,
) -> Callable[[str, Optional[str]], Sequence[float]]:
    """
    Create an OpenAI embedding function for ChromaSQL TEXT queries.

    Args:
        model: OpenAI embedding model name (e.g., "text-embedding-3-small")

    Returns:
        Embedding function that takes (text, model_override) and returns embeddings
    """
    from openai import OpenAI

    client = OpenAI()

    def embed_fn(text: str, model_override: Optional[str] = None) -> Sequence[float]:
        effective_model = model_override or model
        response = client.embeddings.create(input=text, model=effective_model)
        return response.data[0].embedding

    return embed_fn


def build_query_executor_kwargs(env: CollectionEnvironment) -> Dict[str, Any]:
    """
    Convert a CollectionEnvironment into QueryExecutor keyword arguments.

    Args:
        env: CollectionEnvironment with query configuration

    Returns:
        Dictionary of kwargs for QueryExecutor initialization

    Raises:
        QueryExecutorConfigError: If required configuration is missing
    """
    client_type = env.chroma_client_type.lower()

    # Create embedding function for TEXT queries (OpenAI only for standalone)
    embed_fn = create_openai_embed_function(env.embedding_model)

    kwargs: Dict[str, Any] = {
        "query_config_path": env.query_config_path,
        "discriminator_field": env.discriminator_field,
        "client_type": client_type,
        "fallback_to_all": True,
        "embed_fn": embed_fn,
    }

    if client_type == "cloud":
        if (
            env.cloud_api_key is None
            or env.cloud_tenant is None
            or env.cloud_database is None
        ):
            raise QueryExecutorConfigError(
                "cloud_api_key, cloud_tenant, and cloud_database must be provided for "
                "the 'cloud' Chroma client."
            )

        kwargs.update(
            {
                "cloud_api_key": env.cloud_api_key,
                "cloud_tenant": env.cloud_tenant,
                "cloud_database": env.cloud_database,
            }
        )
    elif client_type == "http":
        if env.http_host is None or env.http_port is None:
            raise QueryExecutorConfigError(
                "http_host and http_port must be provided for the 'http' Chroma client."
            )

        kwargs.update(
            {
                "http_host": env.http_host,
                "http_port": env.http_port,
                "http_ssl": env.http_ssl,
            }
        )
    elif client_type == "persistent":
        if env.local_persist_dir is None:
            raise QueryExecutorConfigError(
                "local_persist_dir must be provided for the 'persistent' Chroma client."
            )

        kwargs.update(
            {
                "persist_directory": str(env.local_persist_dir),
            }
        )
    elif client_type == "memory":
        # No extra configuration required
        pass
    else:
        raise QueryExecutorConfigError(
            "Unsupported chroma_client_type. Expected one of "
            "{'cloud', 'http', 'persistent', 'memory'}."
        )

    return kwargs


def load_model_registry(env: CollectionEnvironment) -> Mapping[str, ModelSpec]:
    """
    Load the model registry referenced by the supplied environment.

    Args:
        env: CollectionEnvironment with model_registry_target

    Returns:
        Dictionary mapping model names to ModelSpec objects
    """
    return load_registry_target(env.model_registry_target)
