"""
Utilities for generating environment.json and query_config.json files.

This module provides standalone functions for creating ChromaSQL-compatible
environment configurations from ingestion results.
"""

from __future__ import annotations

import json
import logging
from datetime import datetime, timezone
from pathlib import Path

from idxr.models import CollectionEnvironment

logger = logging.getLogger(__name__)


def generate_environment_configs(
    persist_dir: Path | str,
    collection_name: str,
    total_chunks: int,
    embedding_backend: str,
    embedding_model: str = "text-embedding-3-small",
    model_registry_target: str = (
        "adri_agents.app.pipelines.document_ingest.registry:MODEL_REGISTRY"
    ),
    discriminator_field: str = "collection_name",
    query_config_output: Path | str | None = None,
    environment_output: Path | str | None = None,
    chroma_persist_dir: Path | str | None = None,
    chunk_encoder_model: str | None = None,
    chunk_size_tokens: int | None = None,
    chunk_overlap_tokens: int | None = None,
) -> tuple[Path, Path]:
    """
    Generate environment.json and query_config.json files after ingestion.

    Args:
        persist_dir: Directory where config files will be written
        collection_name: Name of the ingested collection
        total_chunks: Total number of chunks ingested
        embedding_backend: Name of the embedding backend used
            (e.g., "openai", "sentence_transformer")
        embedding_model: Embedding model name (default: "text-embedding-3-small")
        model_registry_target: Import path to model registry
        discriminator_field: Metadata field used as discriminator
        query_config_output: Optional custom path for query_config.json
        environment_output: Optional custom path for environment.json
        chroma_persist_dir: Optional separate ChromaDB database directory
            (defaults to persist_dir if not specified)
        chunk_encoder_model: Tokenizer model used for chunking (optional)
        chunk_size_tokens: Maximum tokens per chunk (optional)
        chunk_overlap_tokens: Token overlap between chunks (optional)

    Returns:
        Tuple of (query_config_path, environment_path)

    Raises:
        OSError: If unable to write configuration files
    """
    persist_dir = Path(persist_dir).expanduser().resolve()

    # ChromaDB database location (may be different from config file location)
    chroma_dir = (
        Path(chroma_persist_dir).expanduser().resolve()
        if chroma_persist_dir
        else persist_dir
    )

    # Determine output paths
    query_config_path = (
        Path(query_config_output).resolve()
        if query_config_output
        else persist_dir / "query_config.json"
    )
    environment_path = (
        Path(environment_output).resolve()
        if environment_output
        else persist_dir / "environment.json"
    )

    # Generate query_config.json
    metadata = {
        "total_collections": 1,
        "total_models": 1,
        "generated_at": datetime.now(timezone.utc)
        .replace(tzinfo=None)
        .isoformat(timespec="seconds"),
        "persist_dir": str(persist_dir),
        "collection_name": collection_name,
        "embedding_backend": embedding_backend,
        "embedding_model": embedding_model,
    }

    # Add optional chunking parameters to metadata
    if chunk_encoder_model is not None:
        metadata["chunk_encoder_model"] = chunk_encoder_model
    if chunk_size_tokens is not None:
        metadata["chunk_size_tokens"] = chunk_size_tokens
    if chunk_overlap_tokens is not None:
        metadata["chunk_overlap_tokens"] = chunk_overlap_tokens

    query_config = {
        "model_to_collections": {
            "DocumentChunk": {
                "collections": [collection_name],
                "total_documents": total_chunks,
            }
        },
        "collection_to_models": {collection_name: ["DocumentChunk"]},
        "metadata": metadata,
    }

    # Write query_config.json
    try:
        query_config_path.parent.mkdir(parents=True, exist_ok=True)
        query_config_path.write_text(
            json.dumps(query_config, indent=2), encoding="utf-8"
        )
        logger.info(f"Query config written to {query_config_path}")
    except OSError as exc:
        logger.error(f"Failed to write query config to {query_config_path}: {exc}")
        raise

    # Use absolute path for query_config_path to avoid ambiguity
    # when loading the environment from different working directories
    query_config_abs = query_config_path.resolve()

    # Generate environment.json
    env = CollectionEnvironment(
        collection_name=collection_name,
        query_config_path=query_config_abs,
        discriminator_field=discriminator_field,
        model_registry_target=model_registry_target,
        embedding_model=embedding_model,
        chroma_client_type="persistent",
        local_persist_dir=chroma_dir,  # Use actual ChromaDB location
        local_collection_name=collection_name,
        is_local=True,
    )

    # Write environment.json
    try:
        env.to_file(environment_path)
        logger.info(f"Environment config written to {environment_path}")
    except OSError as exc:
        logger.error(f"Failed to write environment config to {environment_path}: {exc}")
        raise

    return query_config_path, environment_path
