"""Generate and manage query configuration for multi-collection ChromaDB setups."""

from __future__ import annotations

import json
import logging
from collections import defaultdict
from pathlib import Path
from typing import Any, Dict, List, Mapping, Optional, Set

logger = logging.getLogger(__name__)


def generate_query_config(
    partition_out_dir: Path,
    *,
    output_path: Optional[Path] = None,
    collection_prefix: Optional[str] = None,
    generate_environment: bool = False,
    environment_output_path: Optional[Path] = None,
    embedding_model: str = "text-embedding-3-small",
    model_registry_target: Optional[str] = None,
    discriminator_field: str = "model",
    chroma_client_type: str = "persistent",
    cloud_api_key: Optional[str] = None,
    cloud_tenant: Optional[str] = None,
    cloud_database: Optional[str] = None,
    http_host: Optional[str] = None,
    http_port: Optional[int] = None,
    http_ssl: bool = False,
) -> Dict[str, Any]:
    """Generate query configuration by scanning resume state files.

    Args:
        partition_out_dir: Directory containing partition subdirectories
        output_path: Optional path to write the query config JSON file
        collection_prefix: Optional prefix used for collection names
        generate_environment: If True, generate environment.json alongside query config
        environment_output_path: Optional path for environment.json
        (default: partition_out_dir/environment.json)
        embedding_model: Embedding model used during indexing
        model_registry_target: Import path to model registry
        (e.g., "kb.std.ecc_6_0_ehp_7.registry:MODEL_REGISTRY")
        discriminator_field: Metadata field used as discriminator for model routing
        chroma_client_type: ChromaDB client type ("persistent", "cloud", "http", "memory")
        cloud_api_key: ChromaDB cloud API key (or "env:VAR_NAME" format)
        cloud_tenant: ChromaDB cloud tenant ID
        cloud_database: ChromaDB cloud database name
        http_host: HTTP client host
        http_port: HTTP client port
        http_ssl: Whether to use SSL for HTTP client

    Returns:
        Dictionary mapping model names to their collections with metadata

    The generated config has the structure:
    {
        "model_to_collections": {
            "ModelName": {
                "collections": ["collection_1", "collection_2"],
                "total_documents": 12345,
                "partitions": ["partition_00001", "partition_00002"]
            }
        },
        "collection_to_models": {
            "collection_1": ["ModelName1", "ModelName2"]
        },
        "metadata": {
            "total_collections": 10,
            "total_models": 5,
            "generated_at": "2025-10-31T12:00:00"
        }
    }
    """
    from datetime import datetime, timezone

    if not partition_out_dir.exists() or not partition_out_dir.is_dir():
        raise ValueError(
            f"Partition output directory {partition_out_dir} does not exist "
            "or is not a directory"
        )

    # Map model_name -> {collections: Set[str], documents: int, partitions: Set[str]}
    model_info: Dict[str, Dict[str, Any]] = defaultdict(
        lambda: {"collections": set(), "documents": 0, "partitions": set()}
    )

    # Map collection_name -> Set[model_name]
    collection_to_models: Dict[str, Set[str]] = defaultdict(set)

    # Track all collections we've seen
    all_collections: Set[str] = set()

    # Scan each partition directory
    partition_dirs = sorted(
        [entry for entry in partition_out_dir.iterdir() if entry.is_dir()],
        key=lambda p: p.name,
    )

    if not partition_dirs:
        logger.warning("No partition subdirectories found in %s", partition_out_dir)

    for partition_dir in partition_dirs:
        partition_name = partition_dir.name

        # Find all resume state files in this partition directory
        resume_files = list(partition_dir.glob("*_resume_state.json"))

        if not resume_files:
            logger.debug("No resume state files found in partition %s", partition_name)
            continue

        for resume_file in resume_files:
            try:
                resume_data = json.loads(resume_file.read_text(encoding="utf-8"))
            except (OSError, json.JSONDecodeError) as exc:
                logger.warning(
                    "Failed to read resume state file %s: %s", resume_file, exc
                )
                continue

            if not isinstance(resume_data, dict):
                logger.warning(
                    "Resume state file %s does not contain a JSON object", resume_file
                )
                continue

            # Extract collection name from filename
            # Format: <collection_name>_resume_state.json
            filename = resume_file.stem  # Remove .json
            if filename.endswith("_resume_state"):
                collection_name = filename[: -len("_resume_state")]
            else:
                logger.warning(
                    "Unexpected resume state filename format: %s", resume_file.name
                )
                continue

            all_collections.add(collection_name)

            # Process each model in the resume state
            for model_name, model_state in resume_data.items():
                if not isinstance(model_state, Mapping):
                    continue

                # Check if this model was completed (not in error state)
                if not model_state.get("started"):
                    continue

                # Check if model has been indexed (collection_count > 0)
                collection_count = model_state.get("collection_count", 0)
                if not isinstance(collection_count, int) or collection_count <= 0:
                    continue

                # Add this collection to the model's collection list
                model_info[model_name]["collections"].add(collection_name)
                model_info[model_name]["documents"] += collection_count
                model_info[model_name]["partitions"].add(partition_name)

                # Add this model to the collection's model list
                collection_to_models[collection_name].add(model_name)

    # Convert sets to sorted lists for JSON serialization
    model_to_collections: Dict[str, Dict[str, Any]] = {}
    for model_name, info in model_info.items():
        model_to_collections[model_name] = {
            "collections": sorted(info["collections"]),
            "total_documents": info["documents"],
            "partitions": sorted(info["partitions"]),
        }

    collection_to_models_serialized: Dict[str, List[str]] = {
        coll: sorted(models) for coll, models in collection_to_models.items()
    }

    # Build metadata
    metadata = {
        "total_collections": len(all_collections),
        "total_models": len(model_to_collections),
        "generated_at": datetime.now(timezone.utc)
        .replace(tzinfo=None)
        .isoformat(timespec="seconds"),
        "partition_out_dir": str(partition_out_dir.resolve()),
    }

    if collection_prefix:
        metadata["collection_prefix"] = collection_prefix

    query_config = {
        "model_to_collections": model_to_collections,
        "collection_to_models": collection_to_models_serialized,
        "metadata": metadata,
    }

    # Write to file if output_path is specified
    if output_path is not None:
        try:
            output_path.parent.mkdir(parents=True, exist_ok=True)
            output_path.write_text(
                json.dumps(query_config, indent=2, sort_keys=False),
                encoding="utf-8",
            )
            logger.info("Query config written to %s", output_path)
        except OSError as exc:
            logger.error("Failed to write query config to %s: %s", output_path, exc)
            raise

    # Generate environment config if requested
    if generate_environment:
        from idxr.models import CollectionEnvironment

        # Determine query_config path for environment
        if output_path:
            # Use relative path from environment file to query config
            env_path = environment_output_path or (
                partition_out_dir / "environment.json"
            )
            try:
                query_config_rel = output_path.relative_to(env_path.parent)
            except ValueError:
                # Not in same directory tree, use absolute path
                query_config_rel = output_path.resolve()
        else:
            # No query config path, use placeholder
            query_config_rel = Path("query_config.json")

        # Determine collection_name for the environment
        # If there's a single collection, use it; otherwise use partition_out_dir name
        if len(all_collections) == 1:
            collection_name = list(all_collections)[0]
        else:
            # Use partition_out_dir name as base for collection name
            collection_name = partition_out_dir.name

        # Build environment kwargs
        env_kwargs: Dict[str, Any] = {
            "collection_name": collection_name,
            "query_config_path": query_config_rel,
            "discriminator_field": discriminator_field,
            "model_registry_target": model_registry_target or "",
            "embedding_model": embedding_model,
            "chroma_client_type": chroma_client_type,
        }

        # Add client-specific settings
        if chroma_client_type == "cloud":
            env_kwargs.update(
                {
                    "cloud_api_key": cloud_api_key or "env:CHROMA_API_KEY",
                    "cloud_tenant": cloud_tenant,
                    "cloud_database": cloud_database,
                }
            )
        elif chroma_client_type == "http":
            env_kwargs.update(
                {
                    "http_host": http_host,
                    "http_port": http_port,
                    "http_ssl": http_ssl,
                }
            )
        elif chroma_client_type == "persistent":
            env_kwargs["local_persist_dir"] = partition_out_dir

        # Create and save environment
        env = CollectionEnvironment(**env_kwargs)
        env_output_path = environment_output_path or (
            partition_out_dir / "environment.json"
        )

        try:
            env.to_file(env_output_path)
            logger.info("Environment config written to %s", env_output_path)
        except OSError as exc:
            logger.error(
                "Failed to write environment config to %s: %s", env_output_path, exc
            )
            raise

    return query_config


def load_query_config(config_path: Path) -> Dict[str, Any]:
    """Load query configuration from a JSON file.

    Args:
        config_path: Path to the query config JSON file

    Returns:
        Dictionary containing the query configuration

    Raises:
        ValueError: If the config file doesn't exist or is invalid
    """
    if not config_path.exists():
        raise ValueError(f"Query config file {config_path} does not exist")

    try:
        config = json.loads(config_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        raise ValueError(
            f"Failed to read query config from {config_path}: {exc}"
        ) from exc

    if not isinstance(config, dict):
        raise ValueError(
            f"Query config file {config_path} does not contain a JSON object"
        )

    # Validate required keys
    required_keys = ["model_to_collections", "collection_to_models", "metadata"]
    missing_keys = [key for key in required_keys if key not in config]
    if missing_keys:
        raise ValueError(
            f"Query config missing required keys: {', '.join(missing_keys)}"
        )

    return config


def get_collections_for_models(
    query_config: Dict[str, Any],
    model_names: Optional[List[str]] = None,
) -> List[str]:
    """Get list of collections that need to be queried for given models.

    Args:
        query_config: Query configuration dictionary
        model_names: List of model names to query, or None for all collections

    Returns:
        Sorted list of collection names to query
    """
    if model_names is None or len(model_names) == 0:
        # Query all collections
        return sorted(query_config["collection_to_models"].keys())

    collections: Set[str] = set()
    model_to_collections = query_config["model_to_collections"]

    for model_name in model_names:
        if model_name not in model_to_collections:
            logger.warning("Model %s not found in query config; skipping", model_name)
            continue

        model_colls = model_to_collections[model_name]["collections"]
        collections.update(model_colls)

    return sorted(collections)
