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 = 6) -> 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_consolidation_worker_completes_and_applies_results():
    # Use IdentityDualEncoder for stability; wrap it into a DualEncoder-like interface
    input_dim = 4  # 2 state + 1 action + 1 reward
    # Create a minimal DualEncoder by reusing identity on both branches
    # DualEncoder requires (input_dim, embed_dim, ...); align dims
    dual = DualEncoder(input_dim=input_dim, embed_dim=8, momentum=0.0)
    # Use the default initialized backbones; determinism is not required here

    store = MemoryStore(embed_dim=8, capacity=8)
    for v in [0.0, 1.0]:
        ep = _episode(value=v, length=6)
        tensor = store.episodes_to_tensor([ep])
        key = dual.encode_key(tensor)
        store.write(key, [ep], encoder_step=0)

    cons = Consolidator(encoder=dual, temperature=0.1, reward_weight=0.0)
    # Use 'fork' to avoid pickling issues with internal locks on Linux
    worker = AsyncConsolidationWorker(cons, store, start_method="fork")
    # Exercise lifecycle without sending unpicklable objects via queues
    worker.start()
    # Simulate a payload coming back from the worker loop
    payload = {
        "job_id": 1,
        "metrics": {"loss": 0.0},
        "encoder_state": cons.encoder.state_dict(),
        "optimizer_state": cons.optimizer.state_dict(),
        "steps": cons.total_steps + 1,
        "keys": store.keys.clone(),
        "key_timestamps": list(store.key_timestamps),
        "metadata": store.snapshot_metadata(),
    }
    worker._apply_result(payload)
    # Poll should be empty since we didn't enqueue; stats should be updated
    assert worker.poll(timeout=0.01) is None
    stats = worker.stats()
    assert "last_applied" in stats
    worker.shutdown()
