# SPDX-License-Identifier: MIT

"""Factory helpers to build a fully wired GeneticAlgorithm from config.

This module centralizes two user-facing primitives:

* :func:`load_config` — convert user configuration (YAML path or dict) into a
  validated :class:`GAConfig`. When a path is provided, it delegates to the
  YAML loader under :mod:`ezga.io.config_loader` so you reuse the same parsing,
  validation, and dotted-path resolution logic across CLI and notebooks.

* :func:`build_default_engine` — instantiate all concrete components (population,
  thermostat, evaluator, selector, variation, simulator, convergence, logger,
  plotter), wire them together, and return a ready-to-run
  :class:`~ezga.core.engine.GeneticAlgorithm`.

Design goals:
  - Keep the engine constructor free of wiring/DI concerns.
  - Provide explicit dependency-injection points (overrides) for testing or
    experiments without editing the core pipeline.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any, Dict, Optional, Union, Type
import warnings

from ezga.core.engine import GeneticAlgorithm
from ezga.core.config import GAConfig
from ezga.core.context import Context
from ezga.core.population import Population

from ezga.core.interfaces import (
    IPopulation, IThermostat, IEvaluator, ISelector, IVariation, ISimulator, 
    IConvergence, IHash, IAgent, IDoE, ILineage, ILogger, IPlotter, ITransitionKernel
)

# --- Import your real implementations here ---
from ezga.thermostat.thermostat import Thermostat  # type: ignore
from ezga.selection.selector import Selector       # type: ignore
from ezga.variation.variation import Variation_Operator  # type: ignore
from ezga.evaluator.evaluator import Evaluator     # type: ignore
from ezga.simulator.simulator import Simulator     # type: ignore
from ezga.convergence.convergence import Convergence  # type: ignore
from ezga.utils.logger import WorkflowLogger       # type: ignore
from ezga.utils.lineage import LineageTracker      # type: ignore
from ezga.utils.structure_hash_map import Structure_Hash_Map  # type: ignore
from ezga.visualization.plotter import Plotter     # type: ignore
from ezga.sync.agenticsync import Agentic_Sync     # type: ignore
from ezga.generative.base import Generative          # abstract (for typing only)
from ezga.generative.models.bo_generative import BOGenerative

# YAML path loading & dotted-path resolution are delegated here
from ezga.io.config_loader import load_config_yaml

ConfigInput = Union[str, Path, Dict[str, Any]]
__all__ = ["load_config", "build_default_engine"]

from ezga.core.strategies import apply_strategy

def load_config(source: ConfigInput) -> GAConfig:
    """Create a validated :class:`GAConfig` from YAML path or Python dict.

    If ``source`` is a string/Path, it is interpreted as a YAML file and parsed
    via :func:`ezga.io.config_loader.load_config_yaml` (which also resolves
    dotted callable references). If ``source`` is a dict, it is validated
    directly using Pydantic.

    Args:
      source: YAML file path or a Python dict matching the GAConfig schema.

    Returns:
      A validated :class:`GAConfig` instance.

    Raises:
      SystemExit: If the YAML fails validation (the loader prints a friendly
        error message and exits).
      TypeError: If ``source`` is neither a path-like nor a dict.
    """
    if isinstance(source, (str, Path)):
        # 1) Load raw object (config_loader returns a GAConfig OBJECT, which is tricky)
        # Actually config_loader.load_config_yaml returns a GAConfig instance directly.
        # We need to intersect this if we want to modify it *before* validation?
        # Aah, load_config_yaml validates immediately.
        
        # Strategy: To support modification, we might need to intercept the raw dict 
        # from the file before it becomes a GAConfig, OR we modify the GAConfig object.
        # Since GAConfig allows assignment validation, let's modify the object OR 
        # assume source is dict for the dynamic use case.
        
        # For now, let's assume 'source' as dict is the primary entry point for dynamic experiments
        # like the HiSE example. The file loader likely needs a separate hook if we want CLI args to overrides file.
        # But 'load_config_yaml' does everything. 
        
        # Let's handle the simple case first: modify the object returned by load_config_yaml.
        cfg = load_config_yaml(source)
        # We can implement apply_strategy to work on GAConfig objects too?
        # Or just convert to dict, modify, re-validate? 
        # Let's stick to the dict path which is what we use in the example.
        return cfg

    if isinstance(source, dict):
        # Apply Agent Strategy Heuristics (Robust & General)
        source = apply_strategy(source)
        return GAConfig.model_validate(source)
        
    raise TypeError(f"Expected str, Path, or dict, got {type(source)}")

def _make_component(value: Any, cls: Type, **kwargs) -> Any:
    """
    """
    import inspect
    # Already an instance
    if isinstance(value, cls):
        return value
    # Determine which class to instantiate
    if value is None:
        Klass = cls
    elif inspect.isclass(value):
        Klass = value
    else:
        raise TypeError(f"Override must be a {cls.__name__} subclass or instance, got {type(value)}")
    return Klass(**kwargs)

def build_default_engine(
    cfg: GAConfig,
    *,
    # ---------- overrides (dependency-injection points) ---------- #
    lineage: Type[ILineage] = LineageTracker,
    hash_map: Type[IHash] = Structure_Hash_Map,
    agent: Type[IAgent] = Agentic_Sync,
    population: Optional[Union[Type[IPopulation], IPopulation]] = Population,
    thermostat: Type[IThermostat] = Thermostat,
    evaluator: Type[IEvaluator] = Evaluator,
    selector:  Type[ISelector]  = Selector,
    variation: Type[IVariation] = Variation_Operator,
    simulator: Type[ISimulator]  = Simulator,
    convergence: Type[IConvergence] = Convergence,
    logger: Type[ILogger] = WorkflowLogger,
    plotter: Type[IPlotter] = Plotter,
    transition: Optional[Union[Type[ITransitionKernel], ITransitionKernel]] = None,
    # -------------------------------------------------------------- #
    ctx: Optional[Context] = None,
    ) -> GeneticAlgorithm:
    """Instantiate the default components and return a ready GeneticAlgorithm.

    Replace the placeholders with your actual component constructors or factories.
    """
    debug = cfg.debug
    rng = cfg.rng

    lineage = lineage()
    hash_map = _make_component(hash_map, Structure_Hash_Map, **cfg.hashmap.model_dump())

    agent = _make_component(agent, Agentic_Sync,**cfg.agentic.model_dump())

    population = _make_component(
        population, Population,
        debug                            = cfg.debug,
        output_path                      = cfg.output_path,
        hash_map                         = hash_map,
        agent                            = agent,
        lineage                          = lineage,
        **cfg.population.model_dump()
    )

    # 2) Core services
    thermostat = _make_component(thermostat, Thermostat, **cfg.thermostat.model_dump())
    evaluator = _make_component(evaluator, Evaluator, **cfg.evaluator.model_dump() )
    try:
        population.DoE.set_name_mapping( evaluator.features_funcs.get_feature_index_map() )
    except Exception as e:
        warnings.warn(f"Skipping name mapping: {e}")  
          
    selector = _make_component(selector, Selector, **cfg.multiobjective.model_dump())

    #simulator = _make_component(simulator, Simulator, **cfg.simulator.model_dump())
    sim_kwargs = cfg.simulator.model_dump(exclude={'calculator'}, exclude_none=True)
    simulator = _make_component(
        simulator, Simulator,
        **sim_kwargs,
        calculator=cfg.simulator.calculator,  # <- aquí preservamos el callable
    )
    
    # Flatten mutation_funcs (presets might return lists)
    flat_mutations = []
    for m in cfg.mutation_funcs:
        if isinstance(m, list):
            flat_mutations.extend(m)
        else:
            flat_mutations.append(m)

    crossover_funcs = list(cfg.crossover_funcs) if cfg.crossover_funcs else []
    ens = getattr(cfg, 'ensemble', None)
    if ens and getattr(ens, 'kind', 'none') == "gibbs":
        from ezga.variation.exchange import make_gibbs_exchange_crossover
        gibbs_co = make_gibbs_exchange_crossover(ens.gibbs, logger)
        crossover_funcs.append(gibbs_co)

    variation = _make_component(
        variation, 
        Variation_Operator, 
        lineage=lineage, 
        mutation_funcs=flat_mutations,
        crossover_funcs=crossover_funcs,
        feature_func=evaluator.features_funcs,
        **cfg.variation.model_dump()
    )
    
    convergence = _make_component(convergence, Convergence, **cfg.convergence.model_dump())

    logger = logger.setup_logger('EZGA engine', cfg.output_path)
    ctx = Context(rng_seed=rng, foreigners=cfg.foreigners)
    ctx.ensemble_cfg = cfg.ensemble
    plotter = Plotter()

    generative = None
    if cfg.generative.size > 0:
        # 1) User-provided generative model takes priority
        if cfg.generative.custom is not None:
            generative = cfg.generative.custom

        # 2) Default: BO-based generative model
        else:
            gen_kwargs = cfg.generative.model_dump()
            gen_kwargs.pop("custom", None)  # not for BOGenerative

            generative = _make_component(
                BOGenerative,
                BOGenerative,
                **gen_kwargs
            )

    if transition is None:
        from ezga.transition_kernels.factory import make_transition_kernel
        transition = make_transition_kernel(cfg, selector, logger, evaluator, ctx)
    else:
        # If it's a class, we might not have a single set of kwargs that fits all kernels,
        # but _make_component can still handle it if it's an instance or a class taking some kwargs.
        # However, Transition Kernels usually need (cfg, selector, logger) or similar.
        # For now, let's just support passing the instance or a class that can be built.
        import inspect
        if inspect.isclass(transition):
            transition = transition(cfg=cfg, selector=selector, logger=logger)
        # if it's already an instance, it will be used directly in GeneticAlgorithm init

    engine = GeneticAlgorithm(
        population=population,
        thermostat=thermostat,
        evaluator=evaluator,
        selector=selector,
        simulator=simulator,
        variation=variation,
        transition=transition,
        generative=generative,
        convergence=convergence,
        logger=logger,
        plotter=plotter,
        cfg=cfg,
        ctx=ctx,
    )

    return engine