# modules/morphosyntax/morphosyntax_interface.py

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.morphosyntax_iterative_mongo_db import (
    store_student_morphosyntax_base,
    store_student_morphosyntax_iteration,
    get_student_morphosyntax_analysis,
    update_student_morphosyntax_analysis,
    delete_student_morphosyntax_analysis,
    get_student_morphosyntax_data
)

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."""
    if 'arc_analysis_state' not in st.session_state:
        st.session_state.arc_analysis_state = {
            'base_id': None,           # ID del análisis base
            'original_text': '',       # Texto original
            'original_analysis': None, # Resultado análisis original
            'iteration_text': '',      # Texto de iteración
            'iteration_analysis': None,# Resultado análisis iteración
            '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 en sesión."""
    if 'arc_analysis_state' in st.session_state:
        st.session_state.arc_analysis_state = {
            'base_id': None,
            '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 display_arc_diagram(doc, analysis):
    """Muestra un diagrama de arco sin título."""
    try:
        for sent in doc.sents:
            svg_html = displacy.render(
                sent,
                style="dep",
                options={
                    "distance": 100,
                    "arrow_spacing": 20,
                    "word_spacing": 30
                }
            )
            # Ajustar tamaño y posición
            svg_html = svg_html.replace('height="375"', 'height="200"')
            svg_html = re.sub(
                r'<svg[^>]*>',
                lambda m: m.group(0).replace('height="450"', 'height="300"'),
                svg_html
            )
            svg_html = re.sub(
                r'<g [^>]*transform="translate\((\d+),(\d+)\)"',
                lambda m: f'<g transform="translate({m.group(1)},50)"',
                svg_html
            )
            
            # Envolver en contenedor con estilo
            svg_html = f'<div class="arc-diagram-container">{svg_html}</div>'
            st.write(svg_html, unsafe_allow_html=True)

    except Exception as e:
        logger.error(f"Error en display_arc_diagram: {str(e)}")


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):
    """
    Interfaz principal para el análisis morfosintáctico.
    Evita resets indebidos y conserva la pestaña activa.
    """
    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;
            }
            .divider {
                height: 3px;
                border: none;
                background-color: #333;
                margin: 2rem 0;
            }
            </style>
        """, unsafe_allow_html=True)

        # Inicializar estados de análisis si no existen
        initialize_arc_analysis_state()

        # ------------------------------------------------------------------
        # Si tuvieras un control de tabs global, puedes comentarlo:
        #
        # st.session_state.tab_states['morpho_active'] = True
        # st.session_state.selected_tab = 1
        # ------------------------------------------------------------------

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

        # -------------------- Subtab 0: 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()
                    # Forzar el recálculo limpio
                    st.rerun()

            # Container principal para análisis base
            analysis_container = st.container()
            with analysis_container:
                # Entrada de texto original
                text_input_key = f"original_text_{st.session_state.arc_analysis_state['analysis_count']}"
                text_input = st.text_area(
                    "Texto original",
                    value=st.session_state.arc_analysis_state.get('original_text', ''),
                    key=text_input_key,
                    height=100
                )

                # Botón de análisis
                col1, col2, col3 = st.columns([2, 1, 2])
                with col1:
                    analyze_button = st.button(
                        "Analizar Texto",
                        type="primary",
                        use_container_width=True
                    )

                # Procesar texto original
                if analyze_button and text_input.strip():
                    try:
                        # Realizar análisis base (SpaCy)
                        doc = nlp_models[lang_code](text_input)
                        analysis = perform_advanced_morphosyntactic_analysis(
                            text_input,
                            nlp_models[lang_code]
                        )

                        # Guardar análisis base en BD y obtener ID
                        base_id = store_student_morphosyntax_base(
                            st.session_state.username,
                            text_input,
                            analysis['arc_diagrams']
                        )

                        if base_id:
                            # Actualizar el estado en session_state
                            st.session_state.arc_analysis_state.update({
                                'base_id': base_id,
                                'original_text': text_input,
                                'original_analysis': analysis,  # guardamos el dict
                                'analysis_count': st.session_state.arc_analysis_state['analysis_count'] + 1
                            })

                            # Mostrar diagrama base
                            display_arc_diagram(doc, analysis)

                            # Línea divisora
                            st.markdown('<hr class="divider">', unsafe_allow_html=True)

                            # Área de iteración: usar un formulario
                            with st.form("iteration_form"):
                                # Separamos la key para la iteración para que no se sobreescriba
                                iteration_text_key = f"iteration_text_{st.session_state.arc_analysis_state['analysis_count']}"
                                
                                # Mostrar el texto de iteración que tengamos en session_state
                                iteration_text = st.text_area(
                                    "Texto de iteración",
                                    value=st.session_state.arc_analysis_state.get('iteration_text', text_input),
                                    key=iteration_text_key,
                                    height=100
                                )

                                # Botón de submit en el formulario
                                col1, col2, col3 = st.columns([2,1,2])
                                with col1:
                                    submitted = st.form_submit_button(
                                        "Analizar Cambios",
                                        type="primary",
                                        use_container_width=True
                                    )

                            # Procesar iteración
                            if submitted and iteration_text.strip():
                                try:
                                    doc_iter = nlp_models[lang_code](iteration_text)
                                    analysis_iter = perform_advanced_morphosyntactic_analysis(
                                        iteration_text,
                                        nlp_models[lang_code]
                                    )

                                    # Guardar iteración
                                    iteration_id = store_student_morphosyntax_iteration(
                                        st.session_state.username,
                                        base_id,
                                        text_input,      # Texto original
                                        iteration_text,  # Texto de iteración
                                        analysis_iter['arc_diagrams']
                                    )

                                    if iteration_id:
                                        # Actualizar el estado de iteración en session_state
                                        st.session_state.arc_analysis_state.update({
                                            'iteration_text': iteration_text,
                                            'iteration_analysis': analysis_iter
                                        })
                                        # Mostrar diagrama de iteración
                                        display_arc_diagram(doc_iter, analysis_iter)

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

                    except Exception as e:
                        st.error("Error procesando análisis base")
                        logger.error(f"Error base: {str(e)}")

        # -------------------- Subtab 1: Análisis de Categorías ----------------
        with subtabs[1]:
            st.info("Análisis de Categorías en desarrollo...")

        # -------------------- Subtab 2: Análisis Morfológico ------------------
        with subtabs[2]:
            st.info("Análisis Morfológico en desarrollo...")

    except Exception as e:
        st.error("Error en la interfaz de morfosintaxis")
        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 (opcional)
    """
    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}")
                svg_html = displacy.render(
                    sent,
                    style="dep",
                    options={
                        "distance": 100,
                        "arrow_spacing": 20,
                        "word_spacing": 30
                    }
                )
                svg_html = svg_html.replace('height="375"', 'height="200"')
                svg_html = re.sub(
                    r'<svg[^>]*>',
                    lambda m: m.group(0).replace('height="450"', 'height="300"'),
                    svg_html
                )
                svg_html = re.sub(
                    r'<g [^>]*transform="translate\((\d+),(\d+)\)"',
                    lambda m: f'<g transform="translate({m.group(1)},50)"',
                    svg_html
                )
                svg_html = f'<div class="arc-diagram-container">{svg_html}</div>'
                st.write(svg_html, unsafe_allow_html=True)
            except Exception as exc:
                logger.error(f"Error mostrando diagrama de la oración {i}: {str(exc)}")
                continue
    except Exception as e:
        logger.error(f"Error en display_morphosyntax_results: {str(e)}")