"""Pricing map utilities for Fluxibly.

Load a pricing YAML file and apply model-specific pricing to AgentConfigs.
This keeps pricing data out of individual agent configs.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any

from fluxibly.config.loader import load_yaml
from fluxibly.llm.base import PricingConfig


def load_pricing_map(path: str | Path) -> dict[str, PricingConfig]:
    """Load a pricing YAML file into a ``{model_name: PricingConfig}`` map.

    The YAML is a flat mapping of model names to pricing dicts::

        gpt-4.1-nano:
          input_price: 0.10
          cached_input_price: 0.025
          output_price: 0.40

    Args:
        path: Path to the pricing YAML file.

    Returns:
        Dict mapping model name to PricingConfig.
    """
    raw: dict[str, Any] = load_yaml(path)
    pricing_map: dict[str, PricingConfig] = {}
    for model_name, entry in raw.items():
        if isinstance(entry, dict):
            pricing_map[model_name] = PricingConfig(**entry)
    return pricing_map


def apply_pricing(
    config: Any,
    pricing_map: dict[str, PricingConfig],
) -> None:
    """Apply pricing from a pricing map to an AgentConfig (in-place).

    Sets ``config.llm.pricing`` if the model is found in the map and
    pricing is not already set.  Also applies to ``config.additional_llms``.

    Args:
        config: An AgentConfig (or any object with an ``llm`` attribute).
        pricing_map: The pricing map from :func:`load_pricing_map`.
    """
    # Primary LLM
    llm = getattr(config, "llm", None)
    if llm is not None and llm.pricing is None:
        pricing = pricing_map.get(llm.model)
        if pricing is not None:
            llm.pricing = pricing

    # Additional LLMs (if any)
    additional = getattr(config, "additional_llms", None)
    if additional:
        for llm_cfg in additional.values():
            if llm_cfg.pricing is None:
                pricing = pricing_map.get(llm_cfg.model)
                if pricing is not None:
                    llm_cfg.pricing = pricing
