from __future__ import annotations

import random
from dataclasses import dataclass
from typing import Any

from ..genome.crossover import uniform_crossover
from ..genome.hypothesis import HypothesisAgentGenome
from ..genome.mutation import weighted_random_mutate
from .selection import elite_select, tournament_select


@dataclass(frozen=True)
class GAConfig:
    population_size: int
    elite_k: int
    tournament_k: int
    p_crossover: float
    prompt_mutation_weight: float = 0.0  # Phase 2.5: opt-in LLM mutator


def next_generation(
    population: list[HypothesisAgentGenome],
    fitnesses: dict[str, float],
    cfg: GAConfig,
    rng: random.Random,
    llm: Any | None = None,
) -> list[HypothesisAgentGenome]:
    """Costruisce la prossima generazione: elitismo + tournament + crossover/mutate.

    Quando ``cfg.prompt_mutation_weight > 0`` e ``llm`` è fornito, la mutazione
    invoca ``weighted_random_mutate`` che con quella probabilità usa
    ``mutate_prompt_llm`` (Phase 2.5).
    """
    new_pop: list[HypothesisAgentGenome] = list(
        elite_select(population, fitnesses, cfg.elite_k)
    )

    while len(new_pop) < cfg.population_size:
        if rng.random() < cfg.p_crossover and len(population) >= 2:
            p1 = tournament_select(population, fitnesses, cfg.tournament_k, rng)
            p2 = tournament_select(population, fitnesses, cfg.tournament_k, rng)
            child = uniform_crossover(p1, p2, rng)
        else:
            parent = tournament_select(population, fitnesses, cfg.tournament_k, rng)
            child = weighted_random_mutate(
                parent, rng, llm=llm, prompt_mutation_weight=cfg.prompt_mutation_weight
            )
        new_pop.append(child)

    return new_pop[: cfg.population_size]