#modules/morphosyntax/morphosyntax_interface.py

#Importaciones generales
import streamlit as st
from streamlit_float import *
from streamlit_antd_components import *
from streamlit.components.v1 import html
import spacy
from spacy import displacy
import spacy_streamlit
import pandas as pd
import base64
import re

#Importaciones locales
from .morphosyntax_process import (
    process_morphosyntactic_input,
    format_analysis_results,
    perform_advanced_morphosyntactic_analysis,
    get_repeated_words_colors,
    highlight_repeated_words,
    POS_COLORS,
    POS_TRANSLATIONS
)

from ..utils.widget_utils import generate_unique_key
from ..database.morphosintax_mongo_db import store_student_morphosyntax_result
from ..database.chat_mongo_db import store_chat_history, get_chat_history

import logging
logger = logging.getLogger(__name__)

###########################################################################
def initialize_arc_analysis_state():
    """Inicializa el estado del análisis de arcos y el caché si no existen"""
    # Inicializar estado de análisis
    if 'arc_analysis_state' not in st.session_state:
        st.session_state.arc_analysis_state = {
            'original_text': '',
            'original_analysis': None,
            'iteration_text': '',
            'iteration_analysis': None,
            'analysis_count': 0
        }
        logger.info("Estado de análisis de arcos inicializado")

    # Inicializar caché de análisis
    if 'analysis_cache' not in st.session_state:
        st.session_state.analysis_cache = {}
        logger.info("Caché de análisis inicializado")

###########################################################################

def reset_morpho_state():
    """Resetea el estado del análisis morfosintáctico"""
    if 'arc_analysis_state' in st.session_state:
        st.session_state.arc_analysis_state = {
            'original_text': '',
            'original_analysis': None,
            'iteration_text': '',
            'iteration_analysis': None,
            'analysis_count': 0
        }
############################################################################

def display_original_analysis(container, analysis, lang_code, morpho_t):
    """Muestra el análisis original en el contenedor especificado"""
    with container:
        st.subheader("Análisis Original")
        display_morphosyntax_results(analysis, lang_code, morpho_t)

def display_iteration_analysis(container, analysis, lang_code, morpho_t):
    """Muestra el análisis de cambios en el contenedor especificado"""
    with container:
        st.subheader("Análisis de Cambios")
        display_morphosyntax_results(analysis, lang_code, morpho_t)

############################################################################

def cache_analysis_results(key, result):
    """Almacena resultados de análisis en caché"""
    if not hasattr(st.session_state, 'analysis_cache'):
        initialize_arc_analysis_state()
    st.session_state.analysis_cache[key] = result
    logger.info(f"Resultado almacenado en caché con clave: {key}")

def get_cached_analysis(key):
    """Recupera resultados de análisis del caché"""
    if not hasattr(st.session_state, 'analysis_cache'):
        initialize_arc_analysis_state()
        return None
    return st.session_state.analysis_cache.get(key)


############################################################################
def display_morphosyntax_interface(lang_code, nlp_models, morpho_t):
    try:
        # CSS para layout estable
        st.markdown("""
            <style>
            .stTextArea textarea {
                font-size: 1rem;
                line-height: 1.5;
                min-height: 100px !important;
                height: 100px !important;
            }
            .arc-diagram-container {
                width: 100%;
                padding: 0.5rem;
                margin: 0.5rem 0;
            }
            </style>
        """, unsafe_allow_html=True)

        # Inicializar estado y mantener tab activo
        st.session_state.tab_states['morpho_active'] = True
        st.session_state.selected_tab = 1
        initialize_arc_analysis_state()

        # Crear subtabs
        subtabs = st.tabs([
            "Análisis de Diagramas de Arco",
            "Análisis de Categorías",
            "Análisis Morfológico"
        ])

        # Tab de Diagramas de Arco
        with subtabs[0]:
            # Botón de reset
            col1, col2, col3 = st.columns([2,1,2])
            with col1:
                if st.button("Nuevo Análisis", type="secondary", use_container_width=True):
                    reset_morpho_state()
                    st.rerun()

            # Contenedores separados para cada tipo de análisis
            original_input_container = st.container()
            original_diagram_container = st.container()
            iteration_input_container = st.container()
            iteration_diagram_container = st.container()

            # Sección de análisis original
            with original_input_container:
                text_key = f"original_text_{st.session_state.arc_analysis_state['analysis_count']}"
                text_input = st.text_area(
                    "",
                    value=st.session_state.arc_analysis_state['original_text'],
                    key=text_key,
                    height=100
                )

                col1, col2, col3 = st.columns([2,1,2])
                with col1:
                    analyze_button = st.button(
                        "Analizar Texto Original",
                        type="primary",
                        key=f"analyze_original_{st.session_state.arc_analysis_state['analysis_count']}",
                        use_container_width=True
                    )

            # Procesar texto original
            if analyze_button and text_input.strip():
                try:
                    doc = nlp_models[lang_code](text_input)
                    analysis = perform_advanced_morphosyntactic_analysis(
                        text_input,
                        nlp_models[lang_code]
                    )
                    
                    st.session_state.arc_analysis_state.update({
                        'original_text': text_input,
                        'original_analysis': {'doc': doc, 'analysis': analysis},
                        'iteration_text': text_input,
                        'analysis_count': st.session_state.arc_analysis_state['analysis_count'] + 1
                    })

                    if store_student_morphosyntax_result(
                        username=st.session_state.username,
                        text=text_input,
                        arc_diagrams=analysis['arc_diagrams']
                    ):
                        display_original_analysis(
                            original_diagram_container,
                            st.session_state.arc_analysis_state['original_analysis'],
                            lang_code,
                            morpho_t
                        )

                except Exception as e:
                    st.error("Error al procesar texto original")
                    logger.error(f"Error al procesar texto original: {str(e)}")

            # Sección de iteración
            if st.session_state.arc_analysis_state.get('original_analysis'):
                with iteration_input_container:
                    st.markdown("---")
                    
                    # Formulario de iteración
                    with st.form(key="iteration_form"):
                        iteration_text = st.text_area(
                            "",
                            value=st.session_state.arc_analysis_state.get('iteration_text', ''),
                            key=f"iteration_{st.session_state.arc_analysis_state['analysis_count']}",
                            height=100
                        )
                        
                        # Botón de análisis dentro del formulario
                        submitted = st.form_submit_button(
                            "Analizar Cambios",
                            type="primary",
                            use_container_width=True
                        )
                        
                        # Cuando se envía el formulario
                        if submitted and iteration_text.strip():
                            try:
                                # Verificar si ya existe en caché
                                cache_key = f"iter_{iteration_text}"
                                cached_result = get_cached_analysis(cache_key)
                                
                                if cached_result:
                                    logger.info("Usando resultado cacheado")
                                    doc_iter = cached_result['doc']
                                    analysis_iter = cached_result['analysis']
                                else:
                                    logger.info("Generando nuevo análisis")
                                    doc_iter = nlp_models[lang_code](iteration_text)
                                    analysis_iter = perform_advanced_morphosyntactic_analysis(
                                        iteration_text,
                                        nlp_models[lang_code]
                                    )
                                    # Guardar en caché
                                    cache_analysis_results(cache_key, {
                                        'doc': doc_iter,
                                        'analysis': analysis_iter
                                    })

                                # Actualizar estado
                                st.session_state.arc_analysis_state.update({
                                    'iteration_text': iteration_text,
                                    'iteration_analysis': {'doc': doc_iter, 'analysis': analysis_iter}
                                })

                                if store_student_morphosyntax_result(
                                    username=st.session_state.username,
                                    text=iteration_text,
                                    arc_diagrams=analysis_iter['arc_diagrams']
                                ):
                                    # Mostrar análisis en columnas
                                    col1, col2 = st.columns(2)
                                    
                                    # Mostrar ambos análisis
                                    with col1:
                                        display_original_analysis(
                                            original_diagram_container,
                                            st.session_state.arc_analysis_state['original_analysis'],
                                            lang_code,
                                            morpho_t
                                        )
                                    
                                    with col2:
                                        display_iteration_analysis(
                                            iteration_diagram_container,
                                            {'doc': doc_iter, 'analysis': analysis_iter},
                                            lang_code,
                                            morpho_t
                                        )

                            except Exception as e:
                                st.error("Error al procesar iteración")
                                logger.error(f"Error al procesar iteración: {str(e)}")

        # Otros subtabs...
        with subtabs[1]:
            st.info("Análisis de Categorías en desarrollo...")
            
        with subtabs[2]:
            st.info("Análisis Morfológico en desarrollo...")

    except Exception as e:
        st.error("Error general en la interfaz")
        logger.error(f"Error general en la interfaz: {str(e)}")

############################################################################

def display_morphosyntax_results(result, lang_code, morpho_t):
    """
    Muestra solo el diagrama de arco.
    Args:
        result: Diccionario con el documento procesado y su análisis
        lang_code: Código del idioma
        morpho_t: Diccionario de traducciones
    """
    if result is None:
        return
        
    try:
        doc = result['doc']
        sentences = list(doc.sents)
        for i, sent in enumerate(sentences):
            try:
                st.subheader(f"{morpho_t.get('sentence', 'Sentence')} {i+1}")
                html = displacy.render(sent, style="dep", options={
                    "distance": 100,
                    "arrow_spacing": 20,
                    "word_spacing": 30
                })
                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)
                html = f'<div class="arc-diagram-container">{html}</div>'
                st.write(html, unsafe_allow_html=True)
            except Exception as e:
                logger.error(f"Error en diagrama {i}: {str(e)}")
                continue
    except Exception as e:
        logger.error(f"Error en display_morphosyntax_results: {str(e)}")