# models.py (VERSIÓN A PRUEBA DE ENTORNO ROTO)

import torch
import os
from huggingface_hub import snapshot_download # Importación clave
from transformers import AutoTokenizer, AutoModelForCausalLM
from sentence_transformers import SentenceTransformer
from config import EMBEDDING_MODEL_NAME

# Cargar el modelo de embeddings
def load_embedding_model():
    print("--- Iniciando carga manual del modelo de embedding ---")
    
    # 1. Descargar los archivos del modelo a una caché local y obtener la ruta
    model_folder = snapshot_download(repo_id=EMBEDDING_MODEL_NAME)
    print(f"Modelo descargado en: {model_folder}")

    # 2. Construir la ruta al archivo problemático
    problematic_file_path = os.path.join(model_folder, "pytorch_model.bin")

    # 3. Eliminar el archivo .bin si existe, para forzar el uso de .safetensors
    if os.path.exists(problematic_file_path):
        print(f"Eliminando archivo problemático: {problematic_file_path}")
        os.remove(problematic_file_path)
    else:
        print("El archivo pytorch_model.bin no existe, se procederá con safetensors.")
        
    # 4. Cargar el modelo desde la carpeta local ya "limpia"
    # Se quita el argumento 'use_safetensors' porque ya no es necesario.
    device_str = 'cuda' if torch.cuda.is_available() else 'cpu'
    device = torch.device(device_str)
    
    embedding_model = SentenceTransformer(
        model_folder,  # Cargar desde la ruta local
        device=device
    )
    
    print(f"Modelo de embedding cargado exitosamente desde la ruta local en el dispositivo: {embedding_model.device}")
    return embedding_model

# Cargar el modelo Yi-Coder (se simplifica para consistencia)
def load_yi_coder_model():
    device_str = 'cuda' if torch.cuda.is_available() else 'cpu'
    device = torch.device(device_str)
    
    model_path = "01-ai/Yi-Coder-9B-Chat"
    
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    
    # Se quita 'use_safetensors' para evitar cualquier posible conflicto.
    yi_coder_model = AutoModelForCausalLM.from_pretrained(
        model_path,
        torch_dtype=torch.float16,
        low_cpu_mem_usage=True
    ).to(device).eval()
    
    print(f"Modelo Yi-Coder cargado en el dispositivo: {yi_coder_model.device}")
    return tokenizer, yi_coder_model, device