"""
LLM pool — weighted sampling over one or more LLM backends.
"""

import asyncio
import logging
import random
from typing import Any, Dict, List

from skydiscover.llm.base import LLMResponse
from skydiscover.llm.openai import OpenAILLM
from skydiscover.config import LLMModelConfig

logger = logging.getLogger("skydiscover.llm")


class LLMPool:
    """
    Pool of one or more LLMs.

    Holds a weighted collection of LLM backends and samples one per
    ``generate()`` call.  When the pool contains a single model it
    behaves like a simple LLM client.
    """

    def __init__(self, models_cfg: List[LLMModelConfig]):
        if not models_cfg:
            raise ValueError(
                "LLMPool requires at least one model. Check that your config YAML "
                "exists and specifies a 'models' list under 'llm'."
            )

        self.models_cfg = models_cfg

        # Initialize models from the configuration
        self.models = [
            model_cfg.init_client(model_cfg) if model_cfg.init_client else OpenAILLM(model_cfg)
            for model_cfg in models_cfg
        ]

        # Extract and normalize model weights
        self.weights = [model.weight for model in models_cfg]
        total = sum(self.weights)
        self.weights = [w / total for w in self.weights]

        # Set up random state for deterministic model selection
        self.random_state = random.Random()
        if (
            models_cfg
            and hasattr(models_cfg[0], "random_seed")
            and models_cfg[0].random_seed is not None
        ):
            self.random_state.seed(models_cfg[0].random_seed)
            logger.debug(
                f"LLMPool: Set random seed to {models_cfg[0].random_seed} for deterministic model selection"
            )

        logger.info(
            "Initialized LLM pool with models: "
            + ", ".join(
                f"{model.name} (weight: {weight:.2f})"
                for model, weight in zip(models_cfg, self.weights)
            )
        )

    def _sample_model(self):
        """Sample a model from the pool based on weights."""
        index = self.random_state.choices(
            range(len(self.models)), weights=self.weights, k=1
        )[0]
        return self.models[index]

    async def generate(
        self, system_message: str, messages: List[Dict[str, Any]], **kwargs
    ) -> str:
        """Sample a model from the pool based on weights and generate text."""
        model = self._sample_model()
        return await model.generate(system_message, messages, **kwargs)

    async def generate_image(
        self, system_message: str, messages: List[Dict[str, Any]], **kwargs
    ) -> LLMResponse:
        """Sample a model from the pool and generate text + image."""
        model = self._sample_model()
        return await model.generate_image(system_message, messages, **kwargs)

    async def generate_all(
        self, system_message: str, messages: List[Dict[str, Any]], **kwargs
    ) -> List[str]:
        """Generate text using all available models concurrently."""
        return await asyncio.gather(
            *(model.generate(system_message, messages, **kwargs) for model in self.models)
        )
