from __future__ import annotations

from typing import Any

import numpy as np

from zen_fronts import ZenFronts
from zen_fronts.selection import SelectionConfig
from zen_fronts.selection.schema import SCHEMA_NAME, SCHEMA_VERSION, validate_selection_stats


def _make_zf(*, n_samples: int = 50, percentile: float = 0.2, seed: int = 123) -> ZenFronts:
    def sampler(rng: np.random.Generator) -> dict[str, Any]:
        return {"p": float(rng.uniform(-1.0, 1.0))}

    def mutator(parent: dict[str, Any], rng: np.random.Generator, **kwargs: Any):
        child = dict(parent)
        child["p"] = float(child["p"] + rng.normal(scale=0.1))
        return child, {}

    return ZenFronts(
        crits={"cr1": "max", "cr2": "min"},
        selection=SelectionConfig(n_samples=n_samples, percentile=percentile, seed=seed),
        sampler=sampler,
        mutator=mutator,
    )


def test_selection_schema_present_and_validatable_after_refresh() -> None:
    z = _make_zf(n_samples=40, percentile=0.25, seed=7)
    (pid0,) = z.add_random_point(1)
    pid1 = z.add_point({"p": 0.0})

    # Make both ready.
    z.update_crits(pid0, {"cr1": 1.0, "cr2": 0.0}, t=1.0)
    z.update_crits(pid1, {"cr1": 2.0, "cr2": -1.0}, t=1.0)

    z.refresh(now=1.0)

    st = z.info(pid0)["selection"]
    assert st is not None
    assert st["schema_name"] == SCHEMA_NAME
    assert st["schema_version"] == SCHEMA_VERSION
    assert st["point_id"] == pid0
    assert st["n_active"] == 2
    assert st["k"] == 1

    parsed = validate_selection_stats(st)
    assert parsed["schema_name"] == SCHEMA_NAME
    assert parsed["point_id"] == pid0


def test_refresh_with_single_active_point_has_no_losers() -> None:
    z = _make_zf(n_samples=30, percentile=0.2, seed=1)
    (pid,) = z.add_random_point(1)
    z.update_crits(pid, {"cr1": 1.0, "cr2": 0.0}, t=1.0)

    losers = z.refresh(now=1.0)
    assert losers == ()

    st = z.info(pid)["selection"]
    assert st is not None
    assert st["is_winner"] is True
    assert st["is_loser"] is False
    assert st["win_count"] == 30
    assert st["lose_count"] == 0
    assert st["rank_sum"] == 0


def test_validate_selection_stats_allows_new_fields_and_newer_minor_version() -> None:
    z = _make_zf(n_samples=20, percentile=0.3, seed=9)
    (pid0,) = z.add_random_point(1)
    pid1 = z.add_point({"p": 0.25})

    z.update_crits(pid0, {"cr1": 1.0, "cr2": 0.0}, t=1.0)
    z.update_crits(pid1, {"cr1": 2.0, "cr2": -1.0}, t=1.0)
    z.refresh(now=1.0)

    st = dict(z.info(pid0)["selection"])
    assert st is not None

    # Simulate a future minor schema extension: extra fields, but same MAJOR.
    st["schema_version"] = "1.1.0"
    st["future_field"] = {"anything": [1, 2, 3]}
    if st.get("place") is not None:
        st["place"] = dict(st["place"])
        st["place"]["future_nested"] = 123.0

    parsed = validate_selection_stats(st)
    assert parsed["schema_name"] == SCHEMA_NAME
    assert parsed["schema_version"] == "1.1.0"
    assert "future_field" not in parsed


def test_validate_selection_stats_rejects_new_major_version() -> None:
    z = _make_zf(n_samples=20, percentile=0.3, seed=9)
    (pid0,) = z.add_random_point(1)
    pid1 = z.add_point({"p": 0.25})

    z.update_crits(pid0, {"cr1": 1.0, "cr2": 0.0}, t=1.0)
    z.update_crits(pid1, {"cr1": 2.0, "cr2": -1.0}, t=1.0)
    z.refresh(now=1.0)

    st = dict(z.info(pid0)["selection"])
    assert st is not None
    st["schema_version"] = "2.0.0"

    try:
        validate_selection_stats(st)
        assert False, "expected ValueError"
    except ValueError:
        pass
