import json
import logging
import os
import subprocess
import time

import pandas as pd
from huggingface_hub import snapshot_download

from src.envs import DATA_ARENA_PATH, HF_HOME

# Configure logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")


def time_diff_wrapper(func):
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        diff = end_time - start_time
        logging.info("Time taken for %s: %s seconds", func.__name__, diff)
        return result

    return wrapper


@time_diff_wrapper
def download_dataset(repo_id, local_dir, repo_type="dataset", max_attempts=3, backoff_factor=1.5):
    """Download dataset with exponential backoff retries."""
    attempt = 0
    while attempt < max_attempts:
        try:
            logging.info("Downloading %s to %s", repo_id, local_dir)
            snapshot_download(
                repo_id=repo_id,
                local_dir=local_dir,
                repo_type=repo_type,
                tqdm_class=None,
                token=os.environ.get("HF_TOKEN_PRIVATE"),
                etag_timeout=30,
                max_workers=8,
            )
            logging.info("Download successful")
            return
        except Exception as e:
            wait_time = backoff_factor**attempt
            logging.error("Error downloading %s: %s, retrying in %ss", repo_id, e, wait_time)
            time.sleep(wait_time)
            attempt += 1
    logging.error("Failed to download %s after %s attempts", repo_id, max_attempts)


def download_openbench():
    """Downloads pre generated data"""
    os.makedirs(DATA_ARENA_PATH, exist_ok=True)

    # download answers of different models that we trust
    download_dataset("Vikhrmodels/openbench-eval", DATA_ARENA_PATH)

    print("\nInternal models in openbench-eval:")
    subprocess.run(["ls", f"{DATA_ARENA_PATH}/model_answers/internal/"], check=False)

    print("\nExternal models in openbench-eval:")
    subprocess.run(["ls", f"{DATA_ARENA_PATH}/model_answers/external/"], check=False)

    print("\nJudgement in openbench-eval")
    subprocess.run(["ls", f"{DATA_ARENA_PATH}/model_judgment/gpt-4-1106-preview"], check=False)


def build_leadearboard_df():
    # Retrieve the leaderboard DataFrame
    with open(f"{HF_HOME}/data/leaderboard.json", "r", encoding="utf-8") as eval_file:
        leaderboard_df = pd.DataFrame.from_records(json.load(eval_file))
    return leaderboard_df.copy()