from __future__ import annotations

import random
from collections import Counter
from dataclasses import dataclass

from multi_swarm.genome.hypothesis import HypothesisAgentGenome, ModelTier
from multi_swarm.genome.mutation import weighted_random_mutate

_PROMPT = (
    "Strategia mean-reversion 1h BTC. Entry long quando RSI(14) < 30 e "
    "close > SMA(50). Exit short quando RSI(14) > 70."
)


def _make_genome() -> HypothesisAgentGenome:
    return HypothesisAgentGenome(
        system_prompt=_PROMPT,
        feature_access=["close"],
        temperature=0.9,
        top_p=0.95,
        model_tier=ModelTier.C,
        lookback_window=200,
        cognitive_style="physicist",
    )


@dataclass
class _R:
    text: str


class _AlwaysMutateLLM:
    """Mock LLM che ritorna sempre un prompt mutato valido."""

    def __init__(self) -> None:
        self.calls = 0

    def complete(self, genome, system, user, max_tokens: int = 2000) -> _R:
        self.calls += 1
        return _R(
            text=(
                "<prompt>Strategia momentum 1h BTC. Entry long quando close > "
                f"SMA(70) e ATR(14) crescente. Exit con stop loss 3% (call #{self.calls}).</prompt>"
            )
        )


def test_weighted_dispatcher_zero_weight_never_calls_llm() -> None:
    llm = _AlwaysMutateLLM()
    rng = random.Random(0)
    parent = _make_genome()

    for _ in range(50):
        weighted_random_mutate(parent, rng, llm=llm, prompt_mutation_weight=0.0)

    assert llm.calls == 0


def test_weighted_dispatcher_full_weight_always_calls_llm() -> None:
    llm = _AlwaysMutateLLM()
    rng = random.Random(0)
    parent = _make_genome()

    for _ in range(20):
        child = weighted_random_mutate(
            parent, rng, llm=llm, prompt_mutation_weight=1.0
        )
        assert child.system_prompt != parent.system_prompt

    assert llm.calls == 20


def test_weighted_dispatcher_none_llm_falls_back_to_scalar() -> None:
    """Senza llm passato (backward compat) → solo mutazione scalare."""
    rng = random.Random(0)
    parent = _make_genome()

    for _ in range(50):
        child = weighted_random_mutate(parent, rng, llm=None, prompt_mutation_weight=0.5)
        assert child.system_prompt == parent.system_prompt


def test_weighted_dispatcher_distribution_30_70() -> None:
    """Su 1000 estrazioni con weight=0.3 il prompt mutator deve essere chiamato ~300 volte."""
    llm = _AlwaysMutateLLM()
    rng = random.Random(123)
    parent = _make_genome()

    counter: Counter[str] = Counter()
    for _ in range(1000):
        child = weighted_random_mutate(
            parent, rng, llm=llm, prompt_mutation_weight=0.3
        )
        if child.system_prompt != parent.system_prompt:
            counter["prompt"] += 1
        else:
            counter["scalar"] += 1

    # 30% ± 5% tolerance
    assert 250 <= counter["prompt"] <= 350, f"prompt mutations: {counter['prompt']}"
    assert 650 <= counter["scalar"] <= 750, f"scalar mutations: {counter['scalar']}"