# app.py

import gradio as gr
from models import load_embedding_model, load_yi_coder_model
from pinecone_utils import connect_to_pinecone, vector_search
from ui import build_interface
from config import SIMILARITY_THRESHOLD_DEFAULT, SYSTEM_PROMPT, MAX_LENGTH_DEFAULT
from decorators import gpu_decorator
import torch
import spaces

# Cargar modelos
embedding_model = load_embedding_model()
tokenizer, yi_coder_model, yi_coder_device = load_yi_coder_model()

# Conectar a Pinecone
index = connect_to_pinecone()

# Función para generar código utilizando Yi-Coder
@gpu_decorator(duration=100)
def generate_code(system_prompt, user_prompt, max_length):
    device = yi_coder_device
    model = yi_coder_model
    tokenizer_ = tokenizer  # Ya lo tenemos cargado

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt}
    ]

    # Preparamos el input para el modelo
    prompt = system_prompt + "\n" + user_prompt
    model_inputs = tokenizer_(prompt, return_tensors="pt").to(device)

    with torch.no_grad():
        generated_ids = model.generate(
            model_inputs.input_ids,
            max_new_tokens=max_length,
            eos_token_id=tokenizer_.eos_token_id  
        )

    # Extraer solo la parte generada
    generated_text = tokenizer_.batch_decode(generated_ids, skip_special_tokens=True)[0]
    response = generated_text[len(prompt):].strip()

    return response

# Función para combinar búsqueda vectorial y Yi-Coder
@gpu_decorator(duration=100)
def combined_function(user_prompt, similarity_threshold, selected_option, system_prompt, max_length):
    if selected_option == "Solo Búsqueda Vectorial":
        # Realizar búsqueda vectorial
        search_results = vector_search(user_prompt, embedding_model, index)
        if search_results:
            # Usar el primer resultado
            content = search_results[0]['content']
            return content, None
        else:
            return "No se encontraron resultados en Pinecone.", None
    elif selected_option == "Solo Yi-Coder":
        # Generar respuesta usando Yi-Coder
        yi_coder_response = generate_code(system_prompt, user_prompt, max_length)
        return yi_coder_response, None
    elif selected_option == "Ambos (basado en umbral de similitud)":
        # Realizar búsqueda vectorial
        search_results = vector_search(user_prompt, embedding_model, index)
        if search_results:
            top_result = search_results[0]
            if top_result['score'] >= similarity_threshold:
                content = top_result['content']
                return content, None
            else:
                yi_coder_response = generate_code(system_prompt, user_prompt, max_length)
                return yi_coder_response, None
        else:
            yi_coder_response = generate_code(system_prompt, user_prompt, max_length)
            return yi_coder_response, None
    else:
        return "Opción no válida.", None

# Funciones para el procesamiento de entradas y actualización de imágenes
def process_input(message, history, selected_option, similarity_threshold, system_prompt, max_length):
    response, image = combined_function(message, similarity_threshold, selected_option, system_prompt, max_length)
    history.append((message, response))
    return history, history, image

def update_image(image_url):
    if image_url:
        return image_url
    else:
        return None

def send_preset_question(question, history, selected_option, similarity_threshold, system_prompt, max_length):
    return process_input(question, history, selected_option, similarity_threshold, system_prompt, max_length)

# Construir y lanzar la interfaz
demo = build_interface(process_input, send_preset_question, update_image)

if __name__ == "__main__":
    demo.launch()