import json
from typing import Any
from pydantic import BaseModel

from .wandb_config import WandbConfig
from .dataset_config import DatasetConfig


__all__ = ["Config", "ComponentConfig"]


class ComponentConfig(BaseModel):
    name: str
    params: BaseModel

    def dump(self) -> dict[str, Any]:
        """
        Recursively serializes the component into a dictionary.

        Returns:
            dict: A dictionary containing the component name and its serialized parameters.
        """
        if isinstance(self.params, BaseModel):
            params_dump = self.params.model_dump()
        else:
            params_dump = self.params
        return {"name": self.name, "params": params_dump}


class Config(BaseModel):
    model: ComponentConfig
    dataset_config: DatasetConfig
    wandb_config: WandbConfig
    criterion: ComponentConfig | None = None
    optimizer: ComponentConfig | None = None
    scheduler: ComponentConfig | None = None

    def asdict(self) -> dict[str, Any]:
        """
        Produce a JSON‐serializable dict of this config, including nested
        ComponentConfig and DatasetConfig entries. Useful for saving to file
        or passing to experiment loggers (e.g. wandb.init(config=...)).

        Returns:
            A dict with keys 'model', 'dataset_config', and (if set)
            'criterion', 'optimizer', 'scheduler'.
        """
        data: dict[str, Any] = {
            "model": self.model.dump(),
            "dataset_config": self.dataset_config.model_dump(),
        }
        if self.criterion is not None:
            data["criterion"] = self.criterion.dump()
        if self.optimizer is not None:
            data["optimizer"] = self.optimizer.dump()
        if self.scheduler is not None:
            data["scheduler"] = self.scheduler.dump()
        data["wandb"] = self.wandb_config.model_dump()
        return data

    def save_json(self, file_path: str, indent: int = 4) -> None:
        """
        Save this config to a JSON file.

        Args:
            file_path: Path to write the JSON file.
            indent: JSON indent level.
        """
        config_dict = self.asdict()
        with open(file_path, "w", encoding="utf-8") as f:
            f.write(json.dumps(config_dict, indent=indent))

    @classmethod
    def load_json(cls, file_path: str) -> "Config":
        """
        Loads a configuration from a JSON file and re-instantiates each section using
        the registry keys to recover the original parameter class(es).

        Args:
            file_path (str): Path to the JSON file.

        Returns:
            Config: An instance of Config with the proper parameter classes.
        """
        with open(file_path, "r", encoding="utf-8") as f:
            data = json.load(f)

        # Parse dataset_config and wandb_config using its Pydantic model.
        dataset_config = DatasetConfig(**data.get("dataset_config", {}))
        wandb_config = WandbConfig(**data.get("wandb", {}))

        # Helper function to parse registry fields.
        def parse_field(
            component_data: dict[str, Any], registry_getter
        ) -> ComponentConfig | None:
            name = component_data.get("name")
            params_data = component_data.get("params", {})

            if name is not None:
                expected = registry_getter(name)
                params = expected(**params_data)
                return ComponentConfig(name=name, params=params)
            return None

        from core import (
            ModelRegistry,
            CriterionRegistry,
            OptimizerRegistry,
            SchedulerRegistry,
        )

        parsed_model = parse_field(
            data.get("model", {}),
            lambda key: ModelRegistry.get_model_params(key),
        )
        parsed_criterion = parse_field(
            data.get("criterion", {}),
            lambda key: CriterionRegistry.get_criterion_params(key),
        )
        parsed_optimizer = parse_field(
            data.get("optimizer", {}),
            lambda key: OptimizerRegistry.get_optimizer_params(key),
        )
        parsed_scheduler = parse_field(
            data.get("scheduler", {}),
            lambda key: SchedulerRegistry.get_scheduler_params(key),
        )

        if parsed_model is None:
            raise ValueError("Failed to load model information")

        return cls(
            model=parsed_model,
            dataset_config=dataset_config,
            criterion=parsed_criterion,
            optimizer=parsed_optimizer,
            scheduler=parsed_scheduler,
            wandb_config=wandb_config,
        )