from __future__ import annotations

from dataclasses import dataclass
from typing import Callable, List, Optional, Sequence

import torch
from torch import Tensor

from hippotorch.core.episode import Episode
from hippotorch.encoders.dual import DualEncoder
from hippotorch.memory.replay import HippocampalReplayBuffer
from hippotorch.memory.store import MemoryStore


@dataclass
class QueryResult:
    """Container for read-only retrieval results."""

    query: Tensor
    episodes: List[List[Episode]]
    indices: Tensor
    episode_ids: List[List[int]]
    scores: Optional[Tensor] = None
    attention: Optional[Tensor] = None


def _resolve_backend(
    buffer: Optional[HippocampalReplayBuffer],
    memory: Optional[MemoryStore],
    encoder: Optional[DualEncoder],
) -> tuple[MemoryStore, DualEncoder]:
    if buffer is not None:
        return buffer.memory, buffer.encoder
    if memory is None or encoder is None:
        raise ValueError("Provide either a replay buffer or a (memory, encoder) pair.")
    return memory, encoder


def _prepare_query_tensor(state: Tensor) -> Tensor:
    if state.dim() == 1:
        return state.view(1, 1, -1)
    if state.dim() == 2:
        return state.unsqueeze(0)
    return state


@torch.no_grad()
def query(
    observation: Tensor | Sequence[float],
    *,
    buffer: Optional[HippocampalReplayBuffer] = None,
    memory: Optional[MemoryStore] = None,
    encoder: Optional[DualEncoder] = None,
    query_state: Optional[Tensor] = None,
    query_state_fn: Optional[Callable[[Tensor], Tensor]] = None,
    top_k: int = 5,
    return_scores: bool = True,
    return_attention: bool = False,
) -> QueryResult:
    """Read-only retrieval helper for inference code paths.

    Args:
        observation: Raw observation or already-built query tensor.
        buffer: Convenience handle providing both memory and encoder.
        memory: Explicit memory store. Must be paired with ``encoder`` when no buffer
            is provided.
        encoder: Encoder used to embed the query.
        query_state: Optional tensor that overrides ``observation``.
        query_state_fn: Optional callable to turn the observation into the encoder
            input (e.g., concat [obs, 0, 0]).
        top_k: Number of nearest episodes to fetch.
        return_scores: Include cosine similarity scores in the result.
        return_attention: Include a softmax-normalized attention vector for logging.
    """

    target_memory, target_encoder = _resolve_backend(buffer, memory, encoder)
    device = next(target_encoder.parameters()).device
    obs_tensor = torch.as_tensor(observation, dtype=torch.float32, device=device)
    if query_state is not None:
        qs = torch.as_tensor(query_state, dtype=torch.float32, device=device)
    elif query_state_fn is not None:
        built = query_state_fn(obs_tensor)
        qs = torch.as_tensor(built, dtype=torch.float32, device=device)
    else:
        qs = obs_tensor

    query_tensor = _prepare_query_tensor(qs)
    embeds = target_encoder.encode_query(query_tensor)

    scores, episode_lists, indices = target_memory.retrieve(
        embeds.detach(),
        top_k=top_k,
        encoder=target_encoder,
        current_step=None,
        lazy_refresh=False,
        track_access=False,
        allow_refresh=False,
    )

    full_ids: List[List[int]] = []
    live_ids = target_memory.episode_ids
    for row in indices.cpu():
        batch_ids: List[int] = []
        for idx in row.tolist():
            if 0 <= idx < len(live_ids):
                batch_ids.append(live_ids[idx])
        full_ids.append(batch_ids)

    attn = None
    if return_attention:
        if scores.numel():
            attn = torch.softmax(scores, dim=-1)
        else:
            attn = torch.empty_like(scores)
    result_scores = scores if return_scores else None
    return QueryResult(
        query=embeds.cpu(),
        episodes=episode_lists,
        indices=indices.cpu(),
        episode_ids=full_ids,
        scores=result_scores.cpu() if result_scores is not None else None,
        attention=attn.cpu() if attn is not None else None,
    )


__all__ = ["query", "QueryResult"]
