"""Catalog for managing sets of domain-specific models."""
import warnings
from typing import Dict, Type

from pydantic import BaseModel

from epeira.helpers.marking import get_subcomponent_fields
from epeira.helpers.parsing import get_design_indicator, get_model_input_key

CatalogOutDict = Dict[str, Dict[str, Type[BaseModel]]]
CatalogInDict = Dict[str, Dict[str, list[Type[BaseModel]]]]


class BaseCatalog(BaseModel):
    """Base class for creating and managing catalogs of models.

    Create a catalog from a dictionary. The catalog can then have new models added or can
    combine with another catalog. To use the catalog, it can be converted back into a dictionary.
    There are also some convenience e.g. functions for adding null models where appropriate.
    """

    def __init__(self, data: CatalogInDict = None):
        """Create a catalog with or without a starting dictionary."""
        super().__init__()
        self._data = data or {}

    @classmethod
    def from_dict(cls, data: CatalogInDict) -> "BaseCatalog":
        """Create a BaseCatalog instance from a dictionary.

        Expects a dictionary of the form:
        {
        'design': {
            'problemspec': [ProblemSpec1],
            'component1': [Component1a, Component1b],
            'component2': [Component2],
        },
        'cost': {
            'component1': [Component1Cost],
            'component2': [Component2Cost,Component2Costb]
        },
        'performance': {
            'component1': [Component1Performance]
            }
        }

        The 'cost' and 'performance' categories can be any category of model that should be layered onto
        (mixed into) the design model. The only part of this dictionary that is required is the
        design section. Any mixin models should correspond to one of the design models, or they
        will not be used.
        """
        validate_catalog_input_dict(data)
        return cls(data)

    def add_model(
        self, category: str, component_name: str, model: Type[BaseModel], override_previous: bool = False
    ) -> None:
        """Add a model to the catalog.

        Requires the category (e.g. 'design', 'cost') and the component name the model belongs to.
        If the category or component_name do not yet exist in the catalog, they will be created.
        If a model with the same class name already exists at that location, an exception is raised
        unless override_previous is True.
        """
        if category not in self._data:
            self._data[category] = {}
        if component_name not in self._data[category]:
            self._data[category][component_name] = []
        existing = self._data[category][component_name]
        existing_names = [m.__name__ for m in existing]
        if model.__name__ in existing_names and not override_previous:
            raise Exception(
                f"Catalog already contains '{model.__name__}' under [{category}][{component_name}]. "
                f"To overwrite, set override_previous=True."
            )
        if model.__name__ in existing_names:
            existing[existing_names.index(model.__name__)] = model
        else:
            existing.append(model)

    def combine(self, other: "BaseCatalog", override_previous: bool = False) -> None:
        """Merge another BaseCatalog into this one.

        Iterates all categories and components in the other catalog and adds each
        model via add_model, which handles conflict detection and override logic.
        """
        for category, component_dict in other._data.items():
            for component_name, models in component_dict.items():
                for model in models:
                    self.add_model(category, component_name, model, override_previous)

    @classmethod
    def combine_catalogs(cls, catalogs: list["BaseCatalog"], override_previous: bool = False) -> "BaseCatalog":
        """Return a new BaseCatalog combining all catalogs in the list.

        Unlike combine(), this does not mutate any of the input catalogs.
        Raises ValueError if the list is empty.
        """
        if not catalogs:
            raise ValueError("combine_catalogs requires at least one catalog.")
        result = cls()
        for catalog in catalogs:
            result.combine(catalog, override_previous)
        return result

    @property
    def component_names(self) -> list[str]:
        """Return the names of all components registered under the 'design' category."""
        design_str = get_design_indicator()
        return list(self._data.get(design_str, {}).keys())

    @property
    def all_mixin_types(self) -> list[str]:
        """Return a list of all mixin type categories present in the catalog.

        Returns all top-level category keys except 'design'. For example,
        a catalog with 'design', 'cost', and 'thermal' categories returns ['cost', 'thermal'].
        """
        design_str = get_design_indicator()
        return [category for category in self._data if category != design_str]

    def get_component_mixins(self, component_name: str) -> list[str]:
        """Return a list of mixin type categories that contain the given component.

        For example, if component_name='component1' and the catalog has 'cost' and 'performance'
        categories that both contain 'component1', returns ['cost', 'performance'].
        Does not include 'design'.
        """
        design_str = get_design_indicator()
        return [
            category
            for category, component_dict in self._data.items()
            if category != design_str and component_name in component_dict
        ]

    def get_models(self, component_name: str, model_type: str) -> Dict[str, Type[BaseModel]]:
        """Return a dict of available models for a given component name and model type.

        For example, given component_name='component1' and model_type='design',
        returns {'Component1a': Component1a, 'Component1b': Component1b}.
        Raises KeyError if the combination does not exist in the catalog.
        """
        return {m.__name__: m for m in self._data[model_type][component_name]}

    def subcomponent_options(
        self, component_name: str, model_type: str, klass_name: str
    ) -> dict[str, list[str]]:
        """Return available subcomponent options for a given class in the catalog.

        Retrieves the class by component_name, model_type, and klass_name, then
        inspects its fields for subcomponent markers. For each subcomponent field,
        looks up available models in the 'design' category and returns their names.

        Returns a dict of {field_name: [class_name, ...]} for each subcomponent field.
        """
        klass = self.get_models(component_name, model_type)[klass_name]
        subcomp_fields = get_subcomponent_fields(klass)
        design_str = get_design_indicator()
        design_data = self._data.get(design_str, {})
        return {
            field_name: [m.__name__ for m in design_data.get(field_name, [])]
            for field_name in subcomp_fields
        }

    def to_dict(self) -> CatalogOutDict:
        """Convert to the output expected by the model framework.

        This will convert the data structure:
        {'design': {
            'component1': [Component1a, Component1b],
            'component2': [Component2],
            'problemspec': [ProblemSpec1]},
        'cost': {
            'component1': [Component1Cost],
            'component2': [Component2Cost, Component2Costb]},
        'performance': {
            'component1: [Component1Performance]}}

        to the format:

        {
        'component1_design_model': {'Component1a': Component1a, 'Component1b': Component1b},
        'component2_design_model': {'Component2': Component2},
        'problemspec_design_model': {'ProblemSpec1': ProblemSpec1},
        'component1_cost_model': {'Component1Cost': Component1Cost},
        'component2_cost_model': {'Component2Cost': Component2Cost, 'Component2Costb': Component2Costb},
        'component1_performance_model': {'Component1Performance': Component1Performance}
        }
        """
        result = {}
        for category, component_dict in self._data.items():
            for component_name, models in component_dict.items():
                key = get_model_input_key(component_name, category)
                result[key] = {m.__name__: m for m in models}
        return {k: result[k] for k in sorted(result)}


def validate_catalog_input_dict(data: Dict[str, Dict[str, list[Type[BaseModel]]]]) -> None:
    """Validate structure of a model catalog dictionary."""
    # ---- Required design section ----
    design_str = get_design_indicator()
    if design_str not in data:
        warnings.warn(
            "Catalog does not contain a 'design' section. "
            "This is fine if it will be combined with a catalog that "
            "contains designs."
        )
        return

    design = data[design_str]

    if not isinstance(design, dict) or not design:
        warnings.warn(f"'{design_str}' is currently an empty dictionary.")

    # ---- design keys used for validation ----
    design_keys = set(design.keys())

    # ---- check other categories ----
    for category, subdict in data.items():

        if category == design_str:
            continue

        if not isinstance(subdict, dict):
            warnings.warn(
                f"Category '{category}' should contain a dictionary.",
                stacklevel=2,
            )
            continue

        for comp_name in subdict.keys():
            if comp_name not in design_keys:
                warnings.warn(
                    f"'{category}' contains '{comp_name}', "
                    f"but this component is not defined in '{design_str}'.",
                    stacklevel=2,
                )
