import os
from typing import Literal

from transformers import AutoConfig

from src.config import huggingface_eval_real_data_name, huggingface_eval_synthetic_data_name, huggingface_train_data_name


patterns_list = ["InternVL2"]

def check_hf_model_exists(model_name: str) -> bool:
    for pattern in patterns_list:
        if pattern in model_name:
            return True
    
    try:
        # Try to load the configuration of the model
        _ = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
        return True
    except Exception as e:
        return False


def get_hf_dataset_name(split: Literal["train", "eval"]) -> str:
    hf_account = os.getenv("HF_ACCOUNT")
    if hf_account is None:
        raise ValueError("Please set your HuggingFace to HF_ACCOUNT environment variable: export HF_ACCOUNT=your_username")
    
    if split == "train":
        return f"{hf_account}/{huggingface_train_data_name}"
    elif split == "eval_real":
        return f"{hf_account}/{huggingface_eval_real_data_name}"
    elif split == "eval_synthetic":
        return f"{hf_account}/{huggingface_eval_synthetic_data_name}"
    else:
        raise ValueError(f"Invalid split: {split} not in ['train', 'eval_real', 'eval_synthetic']")
