"""Tests for M2: trigger gates, train state, trainer, and orchestrator."""

from __future__ import annotations

from datetime import UTC, datetime, timedelta

import numpy as np
import pytest

from opensquilla.gateway.config import RouterSelfLearningConfig, SquillaRouterConfig
from opensquilla.squilla_router.self_learning import encode_features, write_sample
from opensquilla.squilla_router.self_learning.dataset import TrainingDataset
from opensquilla.squilla_router.self_learning.feedback import (
    load_feedback_map,
    write_feedback,
)
from opensquilla.squilla_router.self_learning.gates import (
    AGENT_ACTIVE,
    COOLDOWN,
    DISABLED,
    INSUFFICIENT_CLASS_DIVERSITY,
    INSUFFICIENT_DATA,
    NO_DATA,
    READY,
    evaluate_training_gates,
)
from opensquilla.squilla_router.self_learning.orchestrator import (
    in_process_trainer,
    maybe_run_update_router,
)
from opensquilla.squilla_router.self_learning.schema import RouterTrainSample
from opensquilla.squilla_router.self_learning.state import (
    EventStoreStats,
    TrainState,
    load_train_state,
    save_train_state,
    scan_event_store,
)
from opensquilla.squilla_router.self_learning.store import (
    ENV_DISABLE,
    agent_data_dir,
    prune_expired_samples,
)

NOW = datetime(2026, 6, 6, 12, 0, 0, tzinfo=UTC)


def _cfg(**kw) -> RouterSelfLearningConfig:
    base = dict(enabled=True, train_min_samples=5, idle_hours=2.0, cooldown_hours=72.0)
    base.update(kw)
    return RouterSelfLearningConfig(**base)


def _stats(**kw) -> EventStoreStats:
    base = dict(total=100, high_value=50, distinct_classes=3, last_ts="2026-06-06T00:00:00Z")
    base.update(kw)
    return EventStoreStats(**base)


def mk(session, turn, route_class, *, final=None, complaint=False, conf_gate=False, ts=None):
    return RouterTrainSample(
        session_key=session,
        turn_index=turn,
        ts=ts or f"2026-06-06T00:00:{turn:02d}Z",
        feature_schema_version="v1",
        features_390_b64=encode_features(np.random.RandomState(turn).randn(390)),
        route_class=route_class,
        final_route_class=final or route_class,
        complaint_detected=complaint,
        confidence_gate_applied=conf_gate,
    )


# --------------------------------------------------------------------------- #
# Gates (pure)
# --------------------------------------------------------------------------- #


def test_gate_ready_when_all_pass() -> None:
    res = evaluate_training_gates(config=_cfg(), state=TrainState(), stats=_stats(), now=NOW)
    assert res.should_train and res.reason == READY


def test_gate_disabled_master_off() -> None:
    res = evaluate_training_gates(
        config=_cfg(enabled=False), state=TrainState(), stats=_stats(), now=NOW
    )
    assert not res.should_train and res.reason == DISABLED


def test_gate_disabled_by_env(monkeypatch) -> None:
    monkeypatch.setenv(ENV_DISABLE, "1")
    res = evaluate_training_gates(config=_cfg(), state=TrainState(), stats=_stats(), now=NOW)
    assert res.reason == DISABLED


def test_gate_no_data() -> None:
    res = evaluate_training_gates(
        config=_cfg(), state=TrainState(), stats=_stats(total=0), now=NOW
    )
    assert res.reason == NO_DATA


def test_gate_agent_active_blocks() -> None:
    # last activity 30 min ago < idle_hours=2 -> defer
    recent = (NOW - timedelta(minutes=30)).strftime("%Y-%m-%dT%H:%M:%SZ")
    res = evaluate_training_gates(
        config=_cfg(), state=TrainState(), stats=_stats(last_ts=recent), now=NOW
    )
    assert res.reason == AGENT_ACTIVE


def test_gate_cooldown_blocks() -> None:
    recent_train = (NOW - timedelta(hours=10)).strftime("%Y-%m-%dT%H:%M:%SZ")
    res = evaluate_training_gates(
        config=_cfg(cooldown_hours=72.0),
        state=TrainState(last_train_ts=recent_train),
        stats=_stats(),
        now=NOW,
    )
    assert res.reason == COOLDOWN


def test_gate_insufficient_data() -> None:
    res = evaluate_training_gates(
        config=_cfg(train_min_samples=200), state=TrainState(), stats=_stats(high_value=50), now=NOW
    )
    assert res.reason == INSUFFICIENT_DATA


def test_gate_failure_backoff_doubles_threshold() -> None:
    # base 5, 3 failures -> 5 * 2^3 = 40; high_value=30 fails, 50 passes
    state = TrainState(consecutive_failures=3)
    res = evaluate_training_gates(
        config=_cfg(train_min_samples=5), state=state, stats=_stats(high_value=30), now=NOW
    )
    assert res.reason == INSUFFICIENT_DATA
    assert res.effective_min_samples == 40
    res2 = evaluate_training_gates(
        config=_cfg(train_min_samples=5), state=state, stats=_stats(high_value=50), now=NOW
    )
    assert res2.should_train


def test_gate_class_diversity_floor() -> None:
    res = evaluate_training_gates(
        config=_cfg(), state=TrainState(), stats=_stats(distinct_classes=1), now=NOW
    )
    assert res.reason == INSUFFICIENT_CLASS_DIVERSITY


# --------------------------------------------------------------------------- #
# State + stats scan
# --------------------------------------------------------------------------- #


def test_train_state_roundtrip(tmp_path) -> None:
    assert load_train_state("a", tmp_path) == TrainState()
    st = TrainState(
        last_train_ts="2026-06-06T00:00:00Z", consecutive_failures=2, last_version="v1-x"
    )
    save_train_state(st, "a", tmp_path)
    assert load_train_state("a", tmp_path) == st


def test_scan_event_store_counts(tmp_path) -> None:
    write_sample(mk("s", 0, "R1"), "ag", home=tmp_path)
    write_sample(mk("s", 1, "R1", complaint=True), "ag", home=tmp_path)
    write_sample(mk("s", 2, "R2", conf_gate=True), "ag", home=tmp_path)
    stats = scan_event_store("ag", home=tmp_path)
    assert stats.total == 3
    assert stats.high_value == 2  # complaint + conf_gate
    assert stats.distinct_classes == 2  # R1, R2
    assert stats.dominant_schema_version == "v1"


# --------------------------------------------------------------------------- #
# Trainer (real LightGBM on small synthetic data)
# --------------------------------------------------------------------------- #


def _synthetic_dataset(n=80) -> TrainingDataset:
    rng = np.random.RandomState(0)
    feats = rng.randn(n, 390).astype(np.float32)
    y = rng.randint(0, 4, size=n).astype(np.int64)
    # make features weakly separable so training has signal
    for c in range(4):
        feats[y == c, c] += 3.0
    w = np.ones(n, dtype=np.float32)
    return TrainingDataset(X=feats, y=y, w=w, feature_schema_version="vTEST", n_sessions=5)


def test_train_candidate_fresh_builds_loadable_bundle(tmp_path) -> None:
    pytest.importorskip("lightgbm")
    from types import SimpleNamespace

    import lightgbm as lgb

    from opensquilla.squilla_router.self_learning.train import build_candidate_bundle

    base = tmp_path / "base"
    base.mkdir()
    (base / "router.runtime.yaml").write_text("k: v\n", encoding="utf-8")
    (base / "features").mkdir()
    (base / "features" / "meta.json").write_text("{}", encoding="utf-8")
    # no lgbm_main.bin in base -> trains fresh

    learned = tmp_path / "learned"
    info = build_candidate_bundle(
        _synthetic_dataset(),
        base_dir=base,
        learned_root=learned,
        config=SimpleNamespace(num_boost_round=20),
    )
    assert info.used_init_model is False
    bundle = learned / info.version
    assert (bundle / "lgbm_main.bin").is_file()  # real file, not symlink
    assert (bundle / "router.runtime.yaml").exists()  # reused artifact
    assert (bundle / "learned_manifest.json").exists()
    booster = lgb.Booster(model_file=str(bundle / "lgbm_main.bin"))
    pred = booster.predict(_synthetic_dataset(4).X.astype(np.float64))
    assert pred.shape == (4, 4)  # 4 samples x 4 classes


def test_train_candidate_uses_init_model_when_base_present(tmp_path) -> None:
    pytest.importorskip("lightgbm")
    from types import SimpleNamespace

    from opensquilla.squilla_router.self_learning.train import (
        build_candidate_bundle,
        train_booster,
    )

    # First produce a real base lgbm model to continue from.
    base = tmp_path / "base"
    base.mkdir()
    booster, _ = train_booster(
        _synthetic_dataset(), base_model_path=None, config=SimpleNamespace(num_boost_round=20)
    )
    booster.save_model(str(base / "lgbm_main.bin"))

    info = build_candidate_bundle(
        _synthetic_dataset(),
        base_dir=base,
        learned_root=tmp_path / "learned",
        config=SimpleNamespace(num_boost_round=10),
    )
    assert info.used_init_model is True


def test_assembled_bundle_manifest_matches_retrained_head(tmp_path) -> None:
    """The copied artifact manifest must describe the *new* lgbm head.

    The base manifest pins size/sha256 of the shipped model; without a rewrite
    the runtime's bundle validation rejects every candidate as incomplete.
    """
    pytest.importorskip("lightgbm")
    import hashlib
    import json
    from types import SimpleNamespace

    from opensquilla.squilla_router.self_learning.train import (
        build_candidate_bundle,
        train_booster,
    )

    base = tmp_path / "base"
    base.mkdir()
    booster, _ = train_booster(
        _synthetic_dataset(), base_model_path=None, config=SimpleNamespace(num_boost_round=20)
    )
    booster.save_model(str(base / "lgbm_main.bin"))
    base_bytes = (base / "lgbm_main.bin").read_bytes()
    (base / "artifact_manifest.json").write_text(
        json.dumps(
            {
                "schema_version": 1,
                "files": [
                    {
                        "path": "lgbm_main.bin",
                        "size_bytes": len(base_bytes),
                        "sha256": hashlib.sha256(base_bytes).hexdigest(),
                    }
                ],
            }
        ),
        encoding="utf-8",
    )

    info = build_candidate_bundle(
        _synthetic_dataset(),
        base_dir=base,
        learned_root=tmp_path / "learned",
        config=SimpleNamespace(num_boost_round=10),
    )
    bundle = tmp_path / "learned" / info.version
    manifest_path = bundle / "artifact_manifest.json"
    assert manifest_path.is_file() and not manifest_path.is_symlink()
    new_bytes = (bundle / "lgbm_main.bin").read_bytes()
    entry = next(
        e
        for e in json.loads(manifest_path.read_text(encoding="utf-8"))["files"]
        if e["path"] == "lgbm_main.bin"
    )
    assert entry["size_bytes"] == len(new_bytes)
    assert entry["sha256"] == hashlib.sha256(new_bytes).hexdigest()
    # The base manifest is untouched.
    base_entry = json.loads((base / "artifact_manifest.json").read_text(encoding="utf-8"))
    assert base_entry["files"][0]["sha256"] == hashlib.sha256(base_bytes).hexdigest()


# --------------------------------------------------------------------------- #
# Orchestrator (injected trainer; no subprocess)
# --------------------------------------------------------------------------- #


def _router_cfg(**kw) -> SquillaRouterConfig:
    return SquillaRouterConfig(self_learning=_cfg(**kw))


def _write_ready_store(tmp_path, agent="ag", n_high=8) -> None:
    # idle: timestamps old enough; >=2 distinct final classes + high-value signals
    for i in range(n_high):
        write_sample(
            mk("s1", i, "R0" if i % 2 else "R1", final="R2" if i % 2 else "R3",
               complaint=True, ts=f"2026-06-01T00:00:{i:02d}Z"),
            agent, home=tmp_path,
        )
    write_sample(mk("s2", 0, "R1", final="R1", ts="2026-06-01T01:00:00Z"), agent, home=tmp_path)


def test_orchestrator_noop_when_gates_fail(tmp_path) -> None:
    # empty store -> NO_DATA, trainer never called
    calls = []
    res = maybe_run_update_router(
        "ag",
        router_cfg=_router_cfg(train_min_samples=5),
        home=tmp_path,
        now=NOW,
        trainer=lambda *a, **k: calls.append(1),
        base_dir=tmp_path / "base",
    )
    assert not res.ran and res.reason == NO_DATA and not calls


def test_prune_expired_samples_removes_only_files_past_retention(tmp_path) -> None:
    data_dir = agent_data_dir("ag", tmp_path)
    data_dir.mkdir(parents=True)
    stale = data_dir / "samples-20260401.jsonl"
    fresh = data_dir / "samples-20260605.jsonl"
    stale.write_text("{}\n", encoding="utf-8")
    fresh.write_text("{}\n", encoding="utf-8")

    removed = prune_expired_samples("ag", 7, home=tmp_path, now=NOW)

    assert removed == 1
    assert not stale.exists()
    assert fresh.exists()


def test_orchestrator_enforces_retention_before_gates_and_training(tmp_path) -> None:
    import json

    data_dir = agent_data_dir("ag", tmp_path)
    data_dir.mkdir(parents=True)
    stale_path = data_dir / "samples-20260401.jsonl"
    rows = [
        mk("s1", i, "R0" if i % 2 else "R1", final="R2" if i % 2 else "R3",
           complaint=True, ts=f"2026-04-01T00:00:{i:02d}Z")
        for i in range(8)
    ]
    rows.append(mk("s2", 0, "R1", final="R1", ts="2026-04-01T01:00:00Z"))
    stale_path.write_text(
        "".join(json.dumps(s.to_json_dict(), ensure_ascii=False) + "\n" for s in rows),
        encoding="utf-8",
    )
    calls = []

    res = maybe_run_update_router(
        "ag",
        router_cfg=_router_cfg(train_min_samples=5, retention_days=7),
        home=tmp_path,
        now=NOW,
        trainer=lambda *a, **k: calls.append(1),
        base_dir=tmp_path / "base",
    )

    assert not stale_path.exists()
    assert not res.ran and res.reason == NO_DATA and not calls


def test_orchestrator_applies_sample_retention_to_feedback(tmp_path) -> None:
    write_feedback(
        "ag",
        decision_id="stale-rating",
        session_key="agent:ag:webchat:s1",
        turn_index=0,
        rating="down",
        home=tmp_path,
        now=NOW - timedelta(days=10),
        retention_days=30,
    )
    assert "stale-rating" in load_feedback_map("ag", home=tmp_path)

    res = maybe_run_update_router(
        "ag",
        router_cfg=_router_cfg(train_min_samples=5, retention_days=7),
        home=tmp_path,
        now=NOW,
        trainer=lambda *a, **k: None,
        base_dir=tmp_path / "base",
    )

    assert not res.ran and res.reason == NO_DATA
    assert load_feedback_map("ag", home=tmp_path) == {}


def test_orchestrator_trains_and_records_state(tmp_path) -> None:
    pytest.importorskip("lightgbm")
    base = tmp_path / "base"
    base.mkdir()
    (base / "router.runtime.yaml").write_text("k: v\n", encoding="utf-8")
    _write_ready_store(tmp_path, n_high=8)

    res = maybe_run_update_router(
        "ag",
        router_cfg=_router_cfg(train_min_samples=5, idle_hours=2.0, cooldown_hours=72.0),
        home=tmp_path,
        now=NOW,
        trainer=in_process_trainer,
        base_dir=base,
    )
    # Training ran and a candidate + state were recorded (promotion is gated
    # separately; with only 8 samples the gate rejects on insufficient_eval).
    assert res.ran and res.version
    state = load_train_state("ag", tmp_path)
    assert state.last_version == res.version
    assert state.last_train_ts is not None
    receipts = list((tmp_path / "router" / ".receipts").glob("ag-*-*.json"))
    assert receipts


def test_orchestrator_records_failure_and_backsoff(tmp_path) -> None:
    _write_ready_store(tmp_path, n_high=8)

    def boom(*a, **k):
        raise RuntimeError("training blew up")

    res = maybe_run_update_router(
        "ag",
        router_cfg=_router_cfg(train_min_samples=5),
        home=tmp_path,
        now=NOW,
        trainer=boom,
        base_dir=tmp_path / "base",
    )
    assert not res.ran and res.reason == "train_failed" and res.error
    state = load_train_state("ag", tmp_path)
    assert state.consecutive_failures == 1
    assert list((tmp_path / "router" / ".receipts").glob("ag-*-train_failure.json"))