import streamlit as st
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
import pyperclip

# Cargar modelo y tokenizer
model_name = "microsoft/codebert-base-mlm"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Configurar pipeline de generación de texto
text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Crear la interfaz de usuario
st.set_page_config(page_title="CodeBERT Autocompletado y Corrección de Errores")
st.title("CodeBERT Autocompletado y Corrección de Errores")
st.sidebar.title("Configuración de la Generación de Texto")

# Agregar opciones de configuración
num_suggestions = st.sidebar.slider("Número de Sugerencias", 1, 10, 3)
max_length = st.sidebar.slider("Longitud Máxima de la Secuencia de Salida", 10, 1000, 100)
languages = ['python', 'java', 'javascript', 'c#', 'ruby', 'php', 'go', 'swift', 'kotlin', 'pine']
language = st.sidebar.selectbox('Lenguaje de Programación', options=languages)
code_type = st.sidebar.selectbox('Tipo de Código', options=['Indicador', 'Estrategia'])

# Agregar sección de información
st.info("Ingrese su código en el cuadro de texto a continuación y haga clic en el botón para generar sugerencias de autocompletado o corrección de errores utilizando el modelo CodeBERT.")

# Agregar cuadro de texto para ingresar el código
input_code = st.text_area("Ingresa tu código aquí:", height=250)

# Agregar botón para generar sugerencias
if st.button("Generar Sugerencias"):
    # Agregar indicador de carga mientras se generan las sugerencias
    with st.spinner('Generando sugerencias...'):
        # Generar sugerencias de autocompletado o corrección de errores utilizando la API de transformers
        if code_type == 'Indicador':
            prefix = f"{language}: "
        else:
            prefix = f"{language} strategy: "
        generated_text = text_generator(
            input_code,
            max_length=max_length,
            num_return_sequences=num_suggestions,
            do_sample=True,
            top_p=0.9,
            temperature=0.8,
            repetition_penalty=2.0,
            no_repeat_ngram_size=2,
            pad_token_id=tokenizer.eos_token_id,
            prefix=prefix
        )
        
        # Mostrar el botón "Limpiar" para borrar el cuadro de texto de entrada
        st.button("Limpiar")
        
        # Mostrar las sugerencias generadas una a una en respuesta a un botón "Mostrar siguiente sugerencia"
        if len(generated_text) > 0:
            st.subheader("Sugerencias Generadas:")
            suggestion_index = 0
            st.code(generated_text[suggestion_index]['generated_text'].strip(), language=language)
            while suggestion_index < len(generated_text)-1:
                if st.button("Mostrar siguiente sugerencia"):
                    suggestion_index += 1
                    st.code(generated_text[suggestion_index]['generated_text'].strip(), language=language)
                else:
                    break
            
            # Agregar botón "Copiar" para copiar la sugerencia seleccionada al portapapeles
        if len(generated_text) == 1:
            if st.button("Copiar"):
                pyperclip.copy(generated_text[0]['generated_text'].strip())
                st.success("La sugerencia se ha copiado al portapapeles")