# app.py

import os
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
os.environ['KMP_DUPLICATE_LIB_OK']='TRUE'

import streamlit as st
import spacy
from spacy import displacy
import re
from transformers import BlenderbotTokenizer, BlenderbotForConditionalGeneration

from modules.auth import register_user, authenticate_user, get_user_role
from modules.morpho_analysis import get_repeated_words_colors, highlight_repeated_words, POS_COLORS, POS_TRANSLATIONS
from modules.syntax_analysis import visualize_syntax

# Configure the page to use the full width
st.set_page_config(
    page_title="AIdeaText",
    layout="wide",
    page_icon="random"
)

@st.cache_resource

def load_chatbot_model():
    tokenizer = BlenderbotTokenizer.from_pretrained("facebook/blenderbot-400M-distill")
    model = BlenderbotForConditionalGeneration.from_pretrained("facebook/blenderbot-400M-distill")
    return tokenizer, model

# Cargar el modelo del chatbot
chatbot_tokenizer, chatbot_model = load_chatbot_model()

def get_chatbot_response(input_text):
    inputs = chatbot_tokenizer(input_text, return_tensors="pt")
    reply_ids = chatbot_model.generate(**inputs)
    response = chatbot_tokenizer.batch_decode(reply_ids, skip_special_tokens=True)[0]
    return response

def load_spacy_models():
    return {
        'es': spacy.load("es_core_news_lg"),
        'en': spacy.load("en_core_web_lg"),
        'fr': spacy.load("fr_core_news_lg")
    }

def login_page():
    st.title("Iniciar Sesión")
    username = st.text_input("Usuario")
    password = st.text_input("Contraseña", type='password')
    if st.button("Iniciar Sesión"):
        if authenticate_user(username, password):
            st.success(f"Bienvenido, {username}!")
            st.session_state.logged_in = True
            st.session_state.username = username
            st.session_state.role = get_user_role(username)
            st.experimental_rerun()
        else:
            st.error("Usuario o contraseña incorrectos")

def register_page():
    st.title("Registrarse")
    new_username = st.text_input("Nuevo Usuario")
    new_password = st.text_input("Nueva Contraseña", type='password')
    role = st.selectbox("Rol", ["Estudiante", "Profesor"])
    
    # Aquí puedes añadir campos adicionales según el rol si lo deseas
    additional_info = {}
    if role == "Estudiante":
        # Por ejemplo: additional_info['carrera'] = st.text_input("Carrera")
        pass
    elif role == "Profesor":
        # Por ejemplo: additional_info['departamento'] = st.text_input("Departamento")
        pass

    if st.button("Registrarse"):
        if register_user(new_username, new_password, role, additional_info):
            st.success("Registro exitoso. Por favor, inicia sesión.")
        else:
            st.error("El usuario ya existe o ocurrió un error durante el registro")

def main_app():
    # Load spaCy models
    nlp_models = load_spacy_models()

    # Language selection
    languages = {
        'Español': 'es',
        'English': 'en',
        'Français': 'fr'
    }
    selected_lang = st.sidebar.selectbox("Select Language / Seleccione el idioma / Choisissez la langue", list(languages.keys()))
    lang_code = languages[selected_lang]

    # Translations
    translations = {
        'es': {
            'title': "AIdeaText - Análisis morfológico y sintáctico",
            'input_label': "Ingrese un texto para analizar (máx. 5,000 palabras):",
            'input_placeholder': "El objetivo de esta aplicación es que mejore sus habilidades de redacción. Para ello, después de ingresar su texto y presionar el botón obtendrá tres vistas horizontales. La primera, le indicará las palabras que se repiten por categoría gramátical; la segunda, un diagrama de arco le indicara las conexiones sintácticas en cada oración; y la tercera, es un grafo en el cual visualizara la configuración de su texto.",  
            'analyze_button': "Analizar texto",
            'repeated_words': "Palabras repetidas",
            'legend': "Leyenda: Categorías gramaticales",
            'arc_diagram': "Análisis sintáctico: Diagrama de arco",
            'network_diagram': "Análisis sintáctico: Diagrama de red",
            'sentence': "Oración"
        },
        'en': {
            'title': "AIdeaText - Morphological and Syntactic Analysis",
            'input_label': "Enter a text to analyze (max 5,000 words):",
            'input_placeholder': "The goal of this app is for you to improve your writing skills. To do this, after entering your text and pressing the button you will get three horizontal views. The first will indicate the words that are repeated by grammatical category; second, an arc diagram will indicate the syntactic connections in each sentence; and the third is a graph in which you will visualize the configuration of your text.",
            'analyze_button': "Analyze text",
            'repeated_words': "Repeated words",
            'legend': "Legend: Grammatical categories",
            'arc_diagram': "Syntactic analysis: Arc diagram",
            'network_diagram': "Syntactic analysis: Network diagram",
            'sentence': "Sentence"
        },
        'fr': {
            'title': "AIdeaText - Analyse morphologique et syntaxique",
            'input_label': "Entrez un texte à analyser (max 5 000 mots) :",
            'input_placeholder': "Le but de cette application est d'améliorer vos compétences en rédaction. Pour ce faire, après avoir saisi votre texte et appuyé sur le bouton vous obtiendrez trois vues horizontales. Le premier indiquera les mots répétés par catégorie grammaticale; deuxièmement, un diagramme en arcs indiquera les connexions syntaxiques dans chaque phrase; et le troisième est un graphique dans lequel vous visualiserez la configuration de votre texte.",
            'analyze_button': "Analyser le texte",
            'repeated_words': "Mots répétés",
            'legend': "Légende : Catégories grammaticales",
            'arc_diagram': "Analyse syntaxique : Diagramme en arc",
            'network_diagram': "Analyse syntaxique : Diagramme de réseau",
            'sentence': "Phrase"
        }
    }

    # Use translations
    t = translations[lang_code]

    # Crear dos columnas: una para el chat y otra para el análisis
    col1, col2 = st.columns([1, 2])

    with col1:
        st.markdown(f"### Chat con AIdeaText")

        # Inicializar el historial de chat si no existe
        if 'chat_history' not in st.session_state:
            st.session_state.chat_history = []

        # Mostrar el historial de chat
        for i, (role, text) in enumerate(st.session_state.chat_history):
            if role == "user":
                st.text_area(f"Tú:", value=text, height=50, key=f"user_message_{i}", disabled=True)
            else:
                st.text_area(f"AIdeaText:", value=text, height=50, key=f"bot_message_{i}", disabled=True)

        # Campo de entrada para el usuario
        user_input = st.text_input("Escribe tu mensaje aquí:")

        if st.button("Enviar"):
            if user_input:
                # Añadir mensaje del usuario al historial
                st.session_state.chat_history.append(("user", user_input))
                
                # Obtener respuesta del chatbot
                response = get_chatbot_response(user_input)
                
                # Añadir respuesta del chatbot al historial
                st.session_state.chat_history.append(("bot", response))
                
                # Limpiar el campo de entrada
                st.experimental_rerun()

    with col2: 
        st.markdown(f"### {t['title']}")
    
        if st.session_state.role == "Estudiante":
            # Código para la interfaz del estudiante
            if 'input_text' not in st.session_state:
                st.session_state.input_text = ""
    
            sentence_input = st.text_area(t['input_label'], height=150, placeholder=t['input_placeholder'], value=st.session_state.input_text)
            st.session_state.input_text = sentence_input
    
            if st.button(t['analyze_button']):
                if sentence_input:
                    doc = nlp_models[lang_code](sentence_input)
    
                    # Highlighted Repeated Words
                    with st.expander(t['repeated_words'], expanded=True):
                        word_colors = get_repeated_words_colors(doc)
                        highlighted_text = highlight_repeated_words(doc, word_colors)
                        st.markdown(highlighted_text, unsafe_allow_html=True)
    
                    # Legend for grammatical categories
                    st.markdown(f"##### {t['legend']}")
                    legend_html = "<div style='display: flex; flex-wrap: wrap;'>"
                    for pos, color in POS_COLORS.items():
                        if pos in POS_TRANSLATIONS:
                            legend_html += f"<div style='margin-right: 10px;'><span style='background-color: {color}; padding: 2px 5px;'>{POS_TRANSLATIONS[pos]}</span></div>"
                    legend_html += "</div>"
                    st.markdown(legend_html, unsafe_allow_html=True)
    
                    # Arc Diagram
                    with st.expander(t['arc_diagram'], expanded=True):
                        sentences = list(doc.sents)
                        for i, sent in enumerate(sentences):
                            st.subheader(f"{t['sentence']} {i+1}")
                            html = displacy.render(sent, style="dep", options={"distance": 100})
                            html = html.replace('height="375"', 'height="200"')
                            html = re.sub(r'<svg[^>]*>', lambda m: m.group(0).replace('height="450"', 'height="300"'), html)
                            html = re.sub(r'<g [^>]*transform="translate\((\d+),(\d+)\)"', lambda m: f'<g transform="translate({m.group(1)},50)"', html)
                            st.write(html, unsafe_allow_html=True)
    
                    # Network graph
                    with st.expander(t['network_diagram'], expanded=True):
                        fig = visualize_syntax(sentence_input, nlp_models[lang_code], lang_code)
                        st.pyplot(fig)

        elif st.session_state.role == "Profesor":
            # Código para la interfaz del profesor
            st.write("Bienvenido, profesor. Aquí podrás ver el progreso de tus estudiantes.")
            # Añade aquí la lógica para mostrar el progreso de los estudiantes

def main():
    if 'logged_in' not in st.session_state:
        st.session_state.logged_in = False

    if not st.session_state.logged_in:
        menu = ["Iniciar Sesión", "Registrarse"]
        choice = st.sidebar.selectbox("Menu", menu)
        if choice == "Iniciar Sesión":
            login_page()
        elif choice == "Registrarse":
            register_page()
    else:
        if st.sidebar.button("Cerrar Sesión"):
            st.session_state.logged_in = False
            st.experimental_rerun()
        main_app()

if __name__ == "__main__":
    main()