import os
import requests
from transformers import AutoModel

def model_file_exists_and_valid(model_file_path):
    # Check if the model file exists and has a size greater than 0
    return os.path.exists(model_file_path) and os.path.getsize(model_file_path) > 0

def write_model_path_to_txt_file(model_file_path):
    # Write the model path to model_path.txt
    with open('/home/user/data/models/model_path.txt', 'w') as f:
        f.write(model_file_path)

def download_hf_model():
    '''
    Model File Path for HF Models: The download_hf_model function now includes a default model file path (pytorch_model.bin) check. 
    Adjust this path based on the expected model file type (e.g., TensorFlow or Flax models might have different names).
    '''
    model_name = os.getenv("HF_MODEL_NAME")
    model_dir = f"/home/user/data/models/{model_name}"
    model_file_path = os.path.join(model_dir, "pytorch_model.bin") # Assuming PyTorch model for simplicity
    
    if model_file_exists_and_valid(model_file_path):
        print(f"Model {model_name} already downloaded.")
        write_model_path_to_txt_file(model_file_path)
        return
    
    # Authenticate with Hugging Face using the token, if available
    hf_token = os.getenv("HF_TOKEN")
    if hf_token:
        from huggingface_hub import HfFolder
        HfFolder.save_token(hf_token)
    
    print(f"Downloading model: {model_name}...")
    model = AutoModel.from_pretrained(model_name)
    model.save_pretrained(model_dir)
    print(f"Model {model_name} downloaded and saved to {model_dir}")
    write_model_path_to_txt_file(model_file_path)

def download_gguf_model():
    model_name = os.getenv("HF_MODEL_NAME")
    model_dir = f"/home/user/data/models/{model_name}"
    os.makedirs(model_dir, exist_ok=True)
    
    model_url = os.getenv("GGUF_MODEL_URL")  # Assuming URL is provided as an env variable
    model_file_path = os.path.join(model_dir, os.path.basename(model_url))
    
    if model_file_exists_and_valid(model_file_path):
        print(f"Model {model_name} already downloaded.")
        write_model_path_to_txt_file(model_file_path)
        return
    
    print(f"Downloading model from {model_url}...")
    response = requests.get(model_url, stream=True)
    if response.status_code == 200:
        with open(model_file_path, 'wb') as f:
            f.write(response.content)
        print(f"Model downloaded and saved to {model_file_path}")
    else:
        print(f"Failed to download the model. Status code: {response.status_code}")
    write_model_path_to_txt_file(model_file_path)

def download_model():
    model_class = os.getenv("MODEL_CLASS")
    if model_class == 'gguf':
        download_gguf_model()
    elif model_class == 'hf':
        download_hf_model()
    else:
        print(f"Unsupported model class: {model_class}")

if __name__ == "__main__":
    download_model()