import os
import time
import shutil
from pathlib import Path
from functools import partial
from typing import Union, Dict, List

import torch
from torch.utils.data import DataLoader
import datasets
from datasets import load_dataset, Dataset
from transformers import AutoTokenizer, PreTrainedTokenizer, DataCollatorWithPadding
from huggingface_hub import Repository, create_repo, HfApi
from optimum.onnxruntime import (
    AutoOptimizationConfig,
    ORTModelForFeatureExtraction,
    ORTOptimizer,
)

os.environ["TOKENIZERS_PARALLELISM"] = "false"

opt_configs = {
    "O2": AutoOptimizationConfig.O2(),
    "O3": AutoOptimizationConfig.O3(),
    "O4": AutoOptimizationConfig.O4(),
}


def get_batch_size(device_name: str, model_name: str, opt_level: str):
    """
    TODO: run actual tests

    T4 has 16GB
    A10 has 24GB

    Args:
        device_name (`str`):
            The name of the GPU device in use.
        model_name (`str`):
            The name of the model in use.
        opt_level (`str`):
            The optimization level in use.

    Returns:
        `int`:
            The batch size to use.
    """

    if "small" in model_name:
        bs = 192
    elif "base" in model_name:
        bs = 128
    elif "large" in model_name:
        bs = 64
    else:
        bs = 32

    if "A10" in device_name:
        bs *= 2

    if opt_level == "O4":
        bs *= 2

    return bs


def mean_pooling(last_hidden_state: torch.Tensor, attention_mask: torch.Tensor):
    """
    Mean pool the token embeddings.

    Args:
        last_hidden_state (`tuple`):
            The output of the model.
        attention_mask (`torch.Tensor`):
            The attention mask.

    Returns:
        `torch.Tensor`:
            The mean pooled embeddings.
    """
    input_mask_expanded = (
        attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()
    )
    return torch.sum(last_hidden_state * input_mask_expanded, 1) / torch.clamp(
        input_mask_expanded.sum(1), min=1e-9
    )


def get_model_and_tokenizer(model_name: str, optimization_level: str, progress):
    """
    Load the model and tokenizer from the HuggingFace Hub.

    If the model is not already optimized, optimize it and save it to the local directory.

    Args:
        model_name (`str`):
            The name of the model to load.
        optimization_level (`str`):
            The optimization level to use. Should be one of `"O2"`, `"O3"`, or `"O4"`.

    Returns:
        model (`ORTModelForFeatureExtraction`):
            The optimized model.
        tokenizer (`PreTrainedTokenizer`):
            The tokenizer.
    """
    optimized_model_name = f"model_optimized_{optimization_level}.onnx"

    model_dir = Path(model_name.replace("/", "_"))
    if not (model_dir / optimized_model_name).exists():
        if progress is not None:
            progress(0.2, "Downloading tokenizer...")

        tokenizer = AutoTokenizer.from_pretrained(model_name)
        tokenizer.save_pretrained(model_dir)

        if progress is not None:
            progress(0.4, "Downloading model...")

        model = ORTModelForFeatureExtraction.from_pretrained(model_name, export=True)
        model.save_pretrained(model_dir)

        optimizer = ORTOptimizer.from_pretrained(model)
        optimization_config = opt_configs[optimization_level]

        if progress is not None:
            progress(0.6, "Optimizing model...")

        optimizer.optimize(save_dir=model_dir, optimization_config=optimization_config)
        Path(model_dir / "model_optimized.onnx").rename(
            model_dir / optimized_model_name
        )

    else:
        tokenizer = AutoTokenizer.from_pretrained(model_dir)

    if progress is not None:
        progress(0.8, "Loading optimized model and tokenizer...")

    return (
        ORTModelForFeatureExtraction.from_pretrained(
            model_dir,
            file_name=optimized_model_name,
            provider="CUDAExecutionProvider",
        ),
        tokenizer,
    )


def collate_fn(examples, column_name, tokenizer):
    feature_cols = ["input_ids", "attention_mask"]
    features = [{k: x[k] for k in feature_cols} for x in examples]

    tokenized = tokenizer.pad(
        features,
        padding=True,
        max_length=512,
        return_tensors="pt",
        pad_to_multiple_of=16,
        )

    tokenized[column_name] = [x[column_name] for x in examples]

    return tokenized


@torch.inference_mode()
def batch_embed(
    ds: datasets.IterableDataset,
    model: ORTModelForFeatureExtraction,
    tokenizer: PreTrainedTokenizer,
    model_name: str,
    column_name: str,
    new_dataset_id: str,
    opt_level: str,
    upload_batch_size: int = 10_000,
    map_batch_size: int = 2000,
    num2skip: int = 0,
    num2embed: int = -1,
    progress=None,
):
    """
    Run the model on the dataset and upload the embeddings to the hub.

    Args:
        ds (`datasets.Dataset`):
            dataset to embed. From `load_hf_dataset`
        model (`ORTModelForFeatureExtraction`):
            model to use for embedding. From `get_model_and_tokenizer`
        tokenizer (`AutoTokenizer`):
            tokenizer to use for embedding. From `get_model_and_tokenizer`
        model_name (`str`):
            name of the model to use. Used to determine batch size.
        column_name (`str`):
            column name to use for embedding. Default option in gradio app is `text`
        new_dataset_id (`str`):
            id of the new dataset to create. Should include username or organization.
            e.g. nbroad/new-embeddings
        opt_level (`str`):
            optimization level to use. Should be one of `O2`, `O3`, `O4`
            See here for more details on optimization levels:
            https://huggingface.co/docs/optimum/onnxruntime/usage_guides/optimization#optimization-configuration
        upload_batch_size (`int`, *optional*, defaults to `10_000`):
            number of embeddings to upload at once. Defaults to 10,000.
        map_batch_size (`int`, *optional*, defaults to `2000`):
            number of examples to tokenize at once. Defaults to 2000.
        num2skip (`int`, *optional*, defaults to `0`):
            number of examples to skip. Defaults to 0.
        num2embed (`int`, *optional*, defaults to `-1`):
            number of examples to embed. Defaults to -1, which means all examples.

    Returns:
        current_count (`int`):
            number of examples embedded so far
        time_taken (`float`):
            time taken to embed the examples in seconds

    """

    api = HfApi(
        token=os.environ["HF_TOKEN"],
    )

    username = api.whoami()["name"]

    if "/" not in new_dataset_id:
        new_dataset_id = username + "/" + new_dataset_id

    repo = init_git_repo(new_dataset_id)

    embeds = []
    texts = []

    # current count keeps track of how many have been embedded in total
    current_count = num2skip

    # last_count keeps track of how many had been embedded since last push
    last_count = current_count

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    inference_bs = get_batch_size(torch.cuda.get_device_name(0), model_name, opt_level)

    start_time = time.time()

    collator = partial(
        collate_fn, column_name=column_name, tokenizer=tokenizer
    )

    dl = DataLoader(
        ds,
        batch_size=inference_bs,
        shuffle=False,
        num_workers=2,
        pin_memory=True,
        drop_last=False,
        collate_fn=collator,
    )

    for batch in dl:
        ids = batch["input_ids"].to(device)
        mask = batch["attention_mask"].to(device)

        t_ids = torch.zeros_like(ids)

        outputs = model(input_ids=ids, attention_mask=mask, token_type_ids=t_ids)

        embeds.extend(mean_pooling(outputs[0], mask).cpu().tolist())
        texts.extend(batch[column_name])

        current_count += ids.shape[0]

        # Periodically upload to the hub
        if len(embeds) > upload_batch_size:
            push_to_repo(new_dataset_id, last_count, current_count, embeds, texts, api)
            embeds = []
            texts = []
            last_count = current_count

        # Provide updates
        if progress is not None:
            progress(
                (current_count, None),
                "Embedding docs...",
                total=None,
                unit="Docs Embedded",
            )

    time_taken = time.time() - start_time

    # If there are any remaining embeddings, upload them
    if len(embeds) > 0:
        push_to_repo(new_dataset_id, last_count, current_count, embeds, texts, api)

    return current_count - num2skip, time_taken


def init_git_repo(repo_id: str):
    """
    Initialize a git repo for the new dataset.

    ***Removes existing local folder if exists***

    Args:
        repo_id (`str`):
            id of the new dataset to create. Should include username or organization.
            e.g. nbroad/new-embeddings
    """
    local_dir = repo_id.replace("/", "_")

    create_repo(
        repo_id,
        repo_type="dataset",
        token=os.environ["HF_TOKEN"],
        private=True,
        exist_ok=True,
    )
    try:
        repo = Repository(
            local_dir=local_dir,
            clone_from=repo_id,
            repo_type="dataset",
            token=os.environ["HF_TOKEN"],
            skip_lfs_files=True,
        )
    except EnvironmentError:
        shutil.rmtree(local_dir)
        repo = Repository(
            local_dir=local_dir,
            clone_from=repo_id,
            repo_type="dataset",
            token=os.environ["HF_TOKEN"],
            skip_lfs_files=True,
        )

    if repo is not None:
        repo.git_pull()

    return repo


def push_to_repo(
    repo_id: str,
    last_count: int,
    current_count: int,
    embeds: List[List[float]],
    texts: List[str],
    api: HfApi,
):
    """
    Push embeddings to the repo.

    Args:
        repo_id (`str`):
            id of the new dataset to create. Should include username or organization.
        last_count (`int`):
            last count of embeddings.
            This is the number of embeddings that have already been pushed.
        current_count (`int`):
            current count of embeddings.
            This is the number of embeddings that have been pushed after this batch.
        embeds (`List[List[float]]`):
            list of embeddings to push to the repo
        texts (`List[str]`):
            list of texts to push to the repo
        api (`huggingface_hub.HfApi`):
            api to use to push to the repo
    """

    temp_ds = Dataset.from_dict(
        {
            "embedding": embeds,
            "text": texts,
        }
    )

    local_dir = repo_id.replace("/", "_")

    data_dir = Path(local_dir) / "data"
    data_dir.mkdir(exist_ok=True, parents=True)

    # use zfill so sorting puts the files in order
    filename = f"embeddings_{str(last_count).zfill(8)}_{current_count}.parquet"
    filepath = str(data_dir / filename)

    temp_ds.to_parquet(filepath)

    files = sorted(list(data_dir.glob("*.parquet")))

    api.upload_file(
        path_or_fileobj=filepath,
        path_in_repo=f"data/{filename}",
        repo_id=repo_id,
        repo_type="dataset",
        run_as_future=True,
        token=os.environ["HF_TOKEN"],
        commit_message=f"Embedded examples {last_count} thru {current_count}",
    )

    # Delete old files

    if len(files) > 4:
        for file in files[:2]:
            file.unlink()