from __future__ import annotations

import pytest

torch = pytest.importorskip("torch")

from hippotorch.core.episode import Episode
from hippotorch.encoders.dual import DualEncoder
from hippotorch.memory.async_worker import AsyncConsolidationWorker
from hippotorch.memory.consolidator import Consolidator
from hippotorch.memory.store import MemoryStore


def _episode(value: float = 0.0, length: int = 4) -> Episode:
    states = torch.full((length, 2), value)
    actions = torch.zeros(length, 1)
    rewards = torch.ones(length) * value
    dones = torch.zeros(length, dtype=torch.bool)
    return Episode(states=states, actions=actions, rewards=rewards, dones=dones)


def test_async_worker_metrics_reflect_stale_snapshot_application():
    # Small encoder and store
    input_dim = 4  # 2 state + 1 action + 1 reward
    dual = DualEncoder(input_dim=input_dim, embed_dim=8, momentum=0.0)
    store = MemoryStore(embed_dim=8, capacity=8)

    for v in [0.0, 1.0]:
        ep = _episode(value=v, length=4)
        x = torch.cat([ep.states, ep.actions, ep.rewards.unsqueeze(-1)], dim=-1).unsqueeze(0)
        key = dual.encode_key(x)
        store.write(key, [ep], encoder_step=0)

    cons = Consolidator(encoder=dual, temperature=0.1, reward_weight=0.0)
    worker = AsyncConsolidationWorker(cons, store, start_method="fork")
    worker.start()

    # Take a snapshot and then stale one slot on the live memory
    snapshot = store.clone(device="cpu")
    metadata = snapshot.snapshot_metadata()
    timestamps = list(snapshot.key_timestamps)
    # Mutate live slot version so one entry becomes stale
    store.slot_versions[0] += 1

    # Simulate worker payload with original metadata and keys
    payload = {
        "job_id": 123,
        "metrics": {"loss": 0.0},
        "encoder_state": cons.encoder.state_dict(),
        "optimizer_state": cons.optimizer.state_dict(),
        "steps": cons.total_steps + 1,
        "keys": snapshot.keys.clone(),
        "key_timestamps": timestamps,
        "metadata": metadata,
    }

    # Apply and check metrics
    worker._apply_result(payload)
    stats = worker.stats()
    assert stats.get("dropped_results", 0.0) >= 1.0
    # Swap duration should be present and non-negative
    assert stats.get("last_swap_duration", 0.0) >= 0.0
    worker.shutdown()
