import os
from langchain.memory import ConversationBufferMemory
import streamlit as st
import requests
from transformers import pipeline

# Carga las variables de entorno
api_key = os.getenv("HUGGINGFACE_API_TOKEN")

# Streamlit Settings (POner vuestro nombre en la pantalla)
st.set_page_config(
    page_title="Mi asistente personal de IA (By Daniela) 🤖", 
    page_icon="🤖")
st.title("Mi asistente personal de IA (By Daniela) 🤖")

# Inicio seción del chat
if "chat_history" not in st.session_state:
    st.session_state.chat_history = [
        {"role": "assistant", "content": "Hola soy tu asistente virtual, ¿en qué puedo ayudarte hoy?"}
    ]

# Definir la clase del modelo
# Selector de modelo LLM
model_class = st.selectbox(
    "Selecciona el modelo que deseas usar:",
    options=["ollama", "openai", "hf_hub"],
    index=0  # opción por defecto
)

temperature = st.slider("Temperatura", 0.0, 1.0, 0.5, 0.05)
max_tokens = st.slider("Longitud máxima de respuesta (tokens)", 10, 1024, 100, 10)

# estableceremos 3 maneras de conexion a LLM ( 1. con la api deHuggingFace, 2 Con CHATGPT (hay que tener un usuario de pago ) y 3 con Ollama de forma local)
 
"""def model_hf_hub(model="google/shieldgemma-2b", temperature=0.1):
  llm = HuggingFaceEndpoint(
      repo_id=model,
      temperature=temperature,
      max_new_tokens=64,
      return_full_text=False,
      huggingfacehub_api_token=api_key,
      model_kwargs={
        "stop": ["<|endoftext|>", "<|eot_id|>"],  # Tokens de parada
      }
  )
  return llm"""

def model_hf_hub(model="google/shieldgemma-2b", temperature=0.1):
    # Usamos el pipeline para la tarea de generación de texto
    llm = pipeline("text-generation", model=model, tokenizer=model)
    return llm
    
def model_openai(model="gpt-4o-mini", temperature=0.1):
    llm = ChatOpenAI(
        model=model,
        temperature=temperature,
        # other parameters...
        max_tokens=100, #número máximo de tokens a generar en una respuesta.
        frequency_penalty = 0.3 #(penaliza moderadamente la repetición de palabras)
    )
    return llm

def model_ollama(model="phi3:latest", temperature=0.1,):  # probar con 27b
    llm = ChatOllama(
        model=model,
        temperature=temperature,
        base_url="http://127.0.0.1:11434"
    )
    return llm

#Antes de generar la respuesta se comprueba el idioma del usuario que hace la consulta
def detectar_idioma(text):
    try:
        return detect(text)
    except:
        return "error"
        
def model_response(user_query,chat_history, model_class):
    
    # Detecto el idioma
    idioma = detectar_idioma(user_query)
    print(f"Idioma detectado: {idioma}")

    # Traduce códigos de idioma a nombres legibles
    idioma_nombres = {
        "es": "Spanish",
        "en": "English",
        "fr": "French",
        "de": "German",
        "it": "Italian",
        "pt": "Portuguese"
    }
    # Si el idioma detectado no está en el diccionario, por defecto será español
    language = idioma_nombres.get(idioma, "Spanish")
    
    # Mostrar el idioma detectado en la interfaz
    st.write(f"Idioma detectado: {language}")
    
    # hay que poner un if para que llame a las 3 funciones anteriores segun el valor de la variable  model_class
    #llm = model_ollama()
    # Cargar el modelo correspondiente según model_class
    if model_class == "ollama":
        llm = model_ollama()
    elif model_class == "openai":
        llm = model_openai()
    elif model_class == "hf_hub":
        llm = model_hf_hub()
    else:
        raise ValueError("Invalid model class")
    
    # Memoria del chat
    memory = ConversationBufferMemory(return_messages=True)

    # Cargar la memoria con los mensajes previos    
    for message in chat_history:
        if message['role'] == 'user':
            memory.chat_memory.add_message(HumanMessage(content=message['content']))
        elif message["role"] == "assistant":
            memory.chat_memory.add_message(AIMessage(content=message['content']))

    ## definicion del Prompt
    system_prompt = "You are a helpful assistant answering general questions. Please respond in {language}."
    # Definir el idioma de respuesta
    language = "spanish" 
    user_prompt = "{input}"
   
    prompt_template = ChatPromptTemplate.from_messages([
        ("system", system_prompt),
   #     https://python.langchain.com/docs/integrations/memory/streamlit_chat_message_history/
        ("user", user_prompt)
    ])

    # Creamos la cadena de procesamiento (chain)
    chain = prompt_template | llm | StrOutputParser()

    ## Respuesta del modelo
    response = chain.stream({
        "input": user_query,  # tambien hay que mandarle el historial
        "language": language
    })

    # Añadir la nueva respuesta al historial de memoria
    return response

#bucle para visualizar los mensajes del historial
user_query = st.chat_input("Introduce mensaje...")

# Verificar si la consulta no es nula o vacía
# Elimina los espacios vacíos o en blanco al principio o al final de la frase.
if user_query and user_query.strip():
    # Añadir mensaje del usuario al historial
    st.session_state.chat_history.append({"role": "user", "content": user_query})

    # Mostrar el mensaje del usuario en la interfaz
    with st.chat_message("Human"):
            st.write(user_query)

    # Obtener la respuesta del modelo y mostrarla
    response_stream = model_response(user_query, st.session_state.chat_history, model_class)

    # Mostrar la respuesta de la IA en la interfaz
    with st.chat_message("AI"):
        respuesta_completa = st.write_stream(response_stream) # mandar el historial a la funcion model_response

 
    # Añadir la respuesta del asistente al historial
    st.session_state.chat_history.append({"role": "assistant", "content": respuesta_completa})