# 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

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

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

@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"])
    if st.button("Registrarse"):
        if register_user(new_username, new_password, role):
            st.success("Registro exitoso. Por favor, inicia sesión.")
        else:
            st.error("El usuario ya existe")

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...",  
            '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': {
            # ... (mantén las traducciones en inglés)
        },
        'fr': {
            # ... (mantén las traducciones en francés)
        }
    }

    # Use translations
    t = translations[lang_code]

    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

    # Añadir sección de chatbot
    st.header("Chat con AIdeaText")
    user_input = st.text_input("Escribe tu mensaje aquí:")
    if st.button("Enviar"):
        if user_input:
            response = get_chatbot_response(user_input)
            st.text_area("Respuesta del chatbot:", value=response, height=100, max_chars=None, key=None)

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()