import torch
import os
import numpy as np
from torch.utils.data import Dataset, DataLoader
from typing import Dict, Optional, Any, List, Tuple, Union, Callable
from tqdm.auto import tqdm
from charm_toolkit.Datasets import collator
from charm_toolkit.Data import Data
import asyncio
import nest_asyncio
nest_asyncio.apply()

async def _compute_data_embeddings(client: Any, data: Dict[str, torch.Tensor]) -> torch.Tensor:
    """
    Compute embeddings asynchronously for a batch of data using the given client.

    data:
        {
            "ts_array": torch.Tensor of shape [B, ...]   (will be converted to list)
            "descriptions": list or list[list[str]]
        }
    """
    ts_array = data["ts_array"].cpu().numpy().tolist()
    descriptions = data["descriptions"]

    response = await client.embeddings.async_create(
        descriptions=descriptions,
        ts_array=ts_array,
        return_tensors="torch",
        progress=True,
        max_B_per_request=32,
        concurrency_per_call=4,
    )
    return response.embeds   # torch.Tensor of shape [B, ...]


def _gather_batch(dataset, indices):
    """
    Gather a high-level batch from the dataset.
    """
    return torch.cat([dataset[idx][0] for idx in indices])


def precompute_dataset_embeddings(
    client: Any,
    dataset,
    output_path: str,
    memory_batch_size: int = 512,
):
    """
    Precompute embeddings for an entire dataset in high-level batches and save to disk
    via a memmap, flushing after each batch.

    Parameters
    ----------
    client : Any
        Client object that defines the endpoint to compute embeddings.
    dataset : torch.utils.data.Dataset
        Dataset containing samples to embed.
    output_path : str
        Path to save precomputed embeddings (e.g., 'outputs/train_embeddings.pt').
    memory_batch_size : int, default=512
        Number of samples per high-level batch to keep in memory (gets batched further for async embedding call).
    """
    n_samples = len(dataset)
    if n_samples == 0:
        raise ValueError("Dataset is empty.")

    # ---- 1) Run ONE batch to infer embedding shape and initialize memmap ----
    first_indices = list(range(0, min(memory_batch_size, n_samples)))
    first_batch = _gather_batch(dataset, first_indices)
    first_embeds = asyncio.run(_compute_data_embeddings(client, first_batch)).cpu().numpy()
    b0 = first_embeds.shape[0]

    # Global embedding shape: (N, ...) where N = len(dataset)
    emb_shape = (n_samples,) + first_embeds.shape[1:]

    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    memmap_array = np.memmap(output_path, dtype="float32", mode="w+", shape=emb_shape)

    # Write first batch
    memmap_array[0:b0] = first_embeds
    memmap_array.flush()
    print(f"Saved batch 0 ({b0} samples)")

    # ---- 2) Process remaining batches, saving each batch as it completes ----
    offset = b0
    batch_id = 1

    while offset < n_samples:
        end = min(offset + memory_batch_size, n_samples)
        indices = list(range(offset, end))

        batch = _gather_batch(dataset, indices)
        embeds = asyncio.run(_compute_data_embeddings(client, batch)).cpu().numpy()
        bsz = embeds.shape[0]

        memmap_array[offset:offset + bsz] = embeds
        memmap_array.flush()
        print(f"Saved batch {batch_id} ({bsz} samples): [{offset}, {offset + bsz})")

        offset += bsz
        batch_id += 1

    print("Saved embeddings (memmap):", emb_shape)
    return emb_shape


class PrecomputedEmbeddingsDataset(Dataset):
    """
    Dataset wrapper that lazily loads precomputed embeddings using memory-mapped file.
    
    This dataset wraps an existing dataset and adds precomputed embeddings
    to each batch, reducing the need for real-time embedding computation.
    
    Parameters:
    -----------
    base_dataset : torch.utils.data.Dataset
        The original dataset to wrap.
    embeddings_path : str
        Path to the file containing precomputed embeddings.
    """
    def __init__(self, base_dataset, embeddings_path,  emb_shape):
        self.base_dataset = base_dataset
        self.embeddings_path = embeddings_path
        
        # Load the memmap lazily
        self.embeds = np.memmap(embeddings_path, dtype="float32", mode="r", shape=emb_shape)
      
    def __len__(self):
        return len(self.base_dataset)
    
    def __getitem__(self, idx):
        data, target = self.base_dataset[idx]
        
        # Add precomputed embeddings to the data dictionary
        assert isinstance(data, Data), "data must be a Data object"
        if idx < len(self.embeds):
            emb = torch.from_numpy(self.embeds[idx].copy())
            data["embeds"] = emb.unsqueeze(0)

        return data, target