from __future__ import annotations

import math
from typing import Iterable, Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from hypothesis import given, settings
from hypothesis import strategies as st

from hippotorch.core.episode import Episode
from hippotorch.memory.consolidator import Consolidator
from hippotorch.memory.store import MemoryStore

# ---------------------
# Strategies
# ---------------------


def _tensor_strategy(
    shape_st: st.SearchStrategy[tuple[int, ...]], *, dtype: torch.dtype
) -> st.SearchStrategy[torch.Tensor]:
    return shape_st.flatmap(
        lambda shape: st.builds(
            lambda data: torch.tensor(data, dtype=dtype).reshape(shape),
            st.lists(
                st.floats(min_value=-5.0, max_value=5.0, allow_infinity=False, allow_nan=False),
                min_size=max(1, math.prod(shape)),
                max_size=max(1, math.prod(shape)),
            ),
        )
    )


@st.composite
def episode_strategy(draw: st.DrawFn) -> Episode:
    T = draw(st.integers(min_value=1, max_value=10))
    Ds = draw(st.integers(min_value=1, max_value=6))
    Da = draw(st.integers(min_value=1, max_value=6))
    shape_states = (T, Ds)
    shape_actions = (T, Da)
    states = draw(_tensor_strategy(st.just(shape_states), dtype=torch.float32))
    actions = draw(_tensor_strategy(st.just(shape_actions), dtype=torch.float32))
    rewards = draw(_tensor_strategy(st.just((T,)), dtype=torch.float32))

    # Optionally include dones and infos with matching lengths
    with_dones = draw(st.booleans())
    dones = None
    if with_dones:
        dones_list = draw(st.lists(st.booleans(), min_size=T, max_size=T))
        dones = torch.tensor(dones_list, dtype=torch.bool)

    with_infos = draw(st.booleans())
    infos: Optional[list[dict[str, float]]] = None
    if with_infos:
        infos = [{"t": float(t), "r": float(rewards[t].item())} for t in range(T)]

    return Episode(states=states, actions=actions, rewards=rewards, dones=dones, infos=infos)


# ---------------------
# Episode properties
# ---------------------


@given(episode_strategy())
def test_episode_len_and_total_reward(episode: Episode) -> None:
    assert len(episode) == episode.states.shape[0]
    assert math.isclose(
        episode.total_reward,
        float(episode.rewards.sum().item()),
        rel_tol=1e-6,
        abs_tol=1e-6,
    )


@given(episode_strategy(), st.integers(min_value=0, max_value=10))
def test_episode_slice_invariants(ep: Episode, offset: int) -> None:
    T = len(ep)
    if T == 0:
        return
    start = min(offset % (T + 1), T)
    end = min(T, max(start, (start + (offset // 2)) % (T + 1)))
    sub = ep.slice(start, end)
    assert len(sub) == max(0, end - start)
    assert torch.allclose(sub.states, ep.states[start:end])
    assert torch.allclose(sub.actions, ep.actions[start:end])
    assert torch.allclose(sub.rewards, ep.rewards[start:end])
    if ep.dones is not None:
        assert torch.equal(sub.dones, ep.dones[start:end])
    if ep.infos is not None:
        assert isinstance(sub.infos, list)
        assert len(sub.infos) == len(sub)


@given(episode_strategy())
def test_episode_to_dict_roundtrip(ep: Episode) -> None:
    payload = ep.to_dict()
    ep2 = Episode.from_dict(payload)
    assert torch.allclose(ep2.states, ep.states)
    assert torch.allclose(ep2.actions, ep.actions)
    assert torch.allclose(ep2.rewards, ep.rewards)
    if ep.dones is not None:
        assert torch.equal(ep2.dones, ep.dones)
    if ep.infos is not None:
        assert isinstance(ep2.infos, Iterable)


@given(st.integers(min_value=1, max_value=10))
def test_episode_post_init_mismatch_raises(T: int) -> None:
    # Build mismatched actions/rewards lengths to trigger post-init validation
    states = torch.zeros(T, 3, dtype=torch.float32)
    actions = torch.zeros(T + 1, 2, dtype=torch.float32)
    rewards = torch.zeros(T, dtype=torch.float32)
    try:
        Episode(states=states, actions=actions, rewards=rewards)
        assert False, "Expected ValueError due to mismatched time dimensions"
    except ValueError:
        pass


@given(episode_strategy(), st.floats(min_value=0.0, max_value=1.0))
def test_episode_mc_returns_matches_definition(ep: Episode, gamma: float) -> None:
    # Full Monte Carlo path (no bootstrapping)
    returns = ep.compute_returns(gamma=gamma, n_steps=None)
    T = len(ep)
    dones = ep.dones if ep.dones is not None else torch.zeros(T, dtype=torch.bool)
    expected = torch.zeros(T, dtype=ep.rewards.dtype)
    running = torch.tensor(0.0, dtype=ep.rewards.dtype)
    for t in range(T - 1, -1, -1):
        running = ep.rewards[t] + gamma * running * (1.0 - dones[t].to(ep.rewards.dtype))
        expected[t] = running
    assert torch.allclose(returns, expected, atol=1e-5)


@given(
    episode_strategy().filter(lambda ep: ep.dones is None or not ep.dones.any()),
    st.floats(min_value=0.0, max_value=0.99),
    st.integers(min_value=1, max_value=10),
)
def test_episode_n_step_returns_no_dones(ep: Episode, gamma: float, n_steps: int) -> None:
    # Restrict to no-done episodes to simplify the expected form
    T = len(ep)
    n = max(1, min(n_steps, T))
    out = ep.compute_returns(gamma=gamma, n_steps=n, value_fn=None)
    expected = torch.zeros(T, dtype=ep.rewards.dtype)
    for t in range(T):
        ret = torch.tensor(0.0, dtype=ep.rewards.dtype)
        discount = torch.tensor(1.0, dtype=ep.rewards.dtype)
        for k in range(n):
            idx = t + k
            if idx >= T:
                break
            ret = ret + discount * ep.rewards[idx]
            discount = discount * gamma
        expected[t] = ret
    assert torch.allclose(out, expected, atol=1e-5)


@given(
    episode_strategy(),
    st.floats(min_value=0.0, max_value=0.99),
    st.floats(min_value=0.0, max_value=1.0),
)
def test_episode_gae_returns_equal_adv_plus_values(ep: Episode, gamma: float, lam: float) -> None:
    # Deterministic value function: scaled sum over state features
    def value_fn(states: torch.Tensor) -> torch.Tensor:
        if states.dim() > 2:
            states_ = states.reshape(states.shape[0], -1)
        else:
            states_ = states
        # Return shape [T, 1] to be robust to Episode squeezing logic
        return states_.sum(dim=-1, keepdim=True) * 0.1

    returns = ep.compute_returns(gamma=gamma, value_fn=value_fn, gae_lambda=lam, last_value=0.0)
    assert ep.advantages is not None
    # Recompute values deterministically to compare
    v = value_fn(ep.states)
    v = torch.as_tensor(v, dtype=ep.rewards.dtype).squeeze(-1)
    assert torch.allclose(returns, ep.advantages + v, atol=1e-5)


# ---------------------
# MemoryStore properties
# ---------------------


@st.composite
def store_and_batch_strategy(
    draw: st.DrawFn,
) -> tuple[MemoryStore, torch.Tensor, list[Episode]]:
    embed_dim = draw(st.integers(min_value=2, max_value=8))
    num_eps = draw(st.integers(min_value=1, max_value=12))
    capacity = draw(st.integers(min_value=num_eps, max_value=max(num_eps, 16)))
    store = MemoryStore(embed_dim=embed_dim, capacity=capacity, device="cpu")

    # Simple episodes aligned with keys
    episodes: list[Episode] = []
    keys = F.normalize(torch.randn(num_eps, embed_dim), dim=-1)
    for i in range(num_eps):
        T = draw(st.integers(min_value=2, max_value=6))
        Ds = draw(st.integers(min_value=1, max_value=4))
        Da = draw(st.integers(min_value=1, max_value=4))
        s = torch.randn(T, Ds)
        a = torch.randn(T, Da)
        r = torch.randn(T)
        episodes.append(Episode(states=s, actions=a, rewards=r))

    return store, keys, episodes


@given(store_and_batch_strategy())
def test_store_write_and_retrieve_identity(
    data: tuple[MemoryStore, torch.Tensor, list[Episode]],
) -> None:
    store, keys, episodes = data
    store.write(keys, episodes)
    # Retrieve with the exact stored keys as queries
    scores, batch_eps, idxs = store.retrieve(keys, top_k=1)
    assert scores.shape[0] == keys.shape[0]
    assert idxs.shape == (keys.shape[0], 1)
    # Identity: top-1 should point to the corresponding row
    for i in range(idxs.shape[0]):
        assert int(idxs[i, 0].item()) == i
        assert isinstance(batch_eps[i][0], Episode)


@given(
    st.integers(min_value=2, max_value=8),
    st.integers(min_value=1, max_value=8),
    st.integers(min_value=1, max_value=8),
)
def test_store_capacity_enforced(embed_dim: int, num_eps: int, capacity: int) -> None:
    capacity = max(1, capacity)
    num_eps = max(1, num_eps)
    store = MemoryStore(embed_dim=embed_dim, capacity=capacity, device="cpu")
    episodes: list[Episode] = []
    keys = torch.randn(num_eps, embed_dim)
    for _ in range(num_eps):
        T, Ds, Da = 3, 2, 2
        episodes.append(
            Episode(
                states=torch.randn(T, Ds),
                actions=torch.randn(T, Da),
                rewards=torch.randn(T),
            )
        )
    store.write(keys, episodes)
    n = len(store)
    assert n <= capacity
    assert store.keys.shape[0] == n
    assert len(store.episodes) == n
    assert len(store.episode_ids) == n
    assert len(store.slot_versions) == n
    assert len(store.access_counts) == n
    assert len(store.insertion_steps) == n
    assert len(store.priorities) == n


@given(store_and_batch_strategy())
def test_store_checkpoint_roundtrip(
    data: tuple[MemoryStore, torch.Tensor, list[Episode]],
) -> None:
    store, keys, episodes = data
    store.write(keys, episodes)
    payload = store.to_checkpoint()
    clone = MemoryStore.from_checkpoint(payload)
    assert clone.embed_dim == store.embed_dim
    assert clone.capacity == store.capacity
    assert clone.keys.shape == store.keys.shape
    assert len(clone) == len(store)


# ---------------------
# Consolidator properties
# ---------------------


class TinyDual(nn.Module):
    """Lightweight dual-encoder stub for fast property tests."""

    def __init__(self, input_dim: int, embed_dim: int) -> None:
        super().__init__()
        self.online = nn.Linear(input_dim, embed_dim)
        self.target = nn.Linear(input_dim, embed_dim)
        self.momentum = 0.9
        self.refresh_interval = 1_000_000

    def encode_query(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        # x: [B, T, D]
        return self.online(x.mean(dim=1))

    @torch.no_grad()
    def encode_key(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        return self.target(x.mean(dim=1))

    @torch.no_grad()
    def update_target(self) -> None:
        for t_p, o_p in zip(self.target.parameters(), self.online.parameters()):
            t_p.data = 0.9 * t_p.data + 0.1 * o_p.data

    def should_refresh_keys(self) -> bool:
        return False

    def mark_refreshed(self) -> None:
        pass


@st.composite
def consolidator_setup(draw: st.DrawFn) -> tuple[Consolidator, MemoryStore]:
    embed_dim = draw(st.integers(min_value=2, max_value=8))
    # Episode encoder input dim = states + actions + reward
    state_dim = draw(st.integers(min_value=1, max_value=4))
    action_dim = draw(st.integers(min_value=1, max_value=4))
    input_dim = state_dim + action_dim + 1
    model = TinyDual(input_dim=input_dim, embed_dim=embed_dim)

    store = MemoryStore(embed_dim=embed_dim, capacity=32, device="cpu")
    # Populate with a few short episodes so sampling works
    num_eps = draw(st.integers(min_value=2, max_value=8))
    episodes: list[Episode] = []
    keys = F.normalize(torch.randn(num_eps, embed_dim), dim=-1)
    for _ in range(num_eps):
        # Ensure episodes are long enough for contrastive windows (>= 2 * window_size)
        T = draw(st.integers(min_value=4, max_value=6))
        s = torch.randn(T, state_dim)
        a = torch.randn(T, action_dim)
        r = torch.randn(T)
        episodes.append(Episode(states=s, actions=a, rewards=r))
    store.write(keys, episodes)

    return Consolidator(model, temperature=0.1, reward_weight=0.5), store


@given(consolidator_setup())
@settings(max_examples=20)
def test_consolidator_loss_finite(data: tuple[Consolidator, MemoryStore]) -> None:
    cons, store = data
    metrics = cons.sleep(store, steps=1, batch_size=4, refresh_keys=False)
    assert math.isfinite(metrics["loss"]) and metrics["loss"] >= 0.0
    assert math.isfinite(metrics["temporal_loss"]) and metrics["temporal_loss"] >= 0.0
    assert math.isfinite(metrics["reward_loss"]) and metrics["reward_loss"] >= 0.0


@given(consolidator_setup())
@settings(max_examples=20)
def test_consolidator_quality_report_succeeds(
    data: tuple[Consolidator, MemoryStore],
) -> None:
    cons, store = data
    report = cons.get_representation_quality(store, sample_size=8, top_k=3)
    # Values should be finite numbers
    for v in report.to_dict().values():
        assert isinstance(v, float)
        assert math.isfinite(v)
