# modules/morphosyntax/morphosyntax_interface.py

import streamlit as st
import re
import logging
from spacy import displacy

# Importa tu pipeline de spacy, p. ej.:
# nlp_models = {"es": spacy.load("es_core_news_sm")}

from ..database.morphosyntax_iterative_mongo_db import (
    store_student_morphosyntax_base,
    store_student_morphosyntax_iteration,
)

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,
            "original_text": "",
            "iteration_text": "",
            "analysis_count": 0,
        }

def reset_morpho_state():
    """Resetea el estado del análisis morfosintáctico en sesión."""
    st.session_state.arc_analysis_state = {
        "base_id": None,
        "original_text": "",
        "iteration_text": "",
        "analysis_count": 0,
    }

def display_arc_diagram(doc):
    """Muestra un diagrama de arco."""
    try:
        for sent in doc.sents:
            rendered = displacy.render(
                sent, 
                style="dep",
                options={"distance": 100, "arrow_spacing": 20, "word_spacing": 30}
            )
            # Ajustes del SVG
            rendered = rendered.replace('height="375"', 'height="200"')
            rendered = re.sub(
                r'<svg[^>]*>',
                lambda m: m.group(0).replace('height="450"', 'height="300"'),
                rendered
            )
            rendered = re.sub(
                r'<g [^>]*transform="translate\((\d+),(\d+)\)"',
                lambda m: f'<g transform="translate({m.group(1)},50)"',
                rendered
            )
            st.write(f'<div class="arc-diagram-container">{rendered}</div>', unsafe_allow_html=True)

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

def display_morphosyntax_interface(lang_code, nlp_models):
    st.title("Análisis Morfosintáctico")
    
    # 1) Inicializar estados
    initialize_arc_analysis_state()
    
    # 2) Tabs
    tabs = st.tabs(["Arco", "Categorías", "Morfológico"])
    
    # ==================== TAB 0: Análisis de Arco =========================
    with tabs[0]:
        st.write("Texto base e iteraciones...")

        # Botón Nuevo Análisis
        if st.button("Nuevo Análisis", key="new_analysis"):
            reset_morpho_state()
            st.experimental_rerun()
        
        # Texto base
        st.session_state.arc_analysis_state["original_text"] = st.text_area(
            "Texto Original",
            value=st.session_state.arc_analysis_state["original_text"],
            key="base_text_input"
        )
        
        # Botón para analizar texto base
        if st.button("Analizar Texto Base", key="analyze_base"):
            # Lógica de Spacy
            if not st.session_state.arc_analysis_state["original_text"].strip():
                st.warning("Ingrese texto original para analizar.")
            else:
                doc = nlp_models[lang_code](
                    st.session_state.arc_analysis_state["original_text"]
                )
                # Realizar análisis (displacy, etc.)
                display_arc_diagram(doc)
                
                # Guardar en DB
                base_id = store_student_morphosyntax_base(
                    st.session_state.username,
                    st.session_state.arc_analysis_state["original_text"],
                    arc_diagrams=None  # Pon el dict real con arcos
                )
                if base_id:
                    st.session_state.arc_analysis_state["base_id"] = base_id
                    st.success(f"Análisis base guardado. base_id={base_id}")
        
        st.markdown("---")

        # Texto de iteración
        st.session_state.arc_analysis_state["iteration_text"] = st.text_area(
            "Texto de Iteración",
            value=st.session_state.arc_analysis_state["iteration_text"],
            key="iteration_text_input"
        )

        # Botón para analizar iteración
        if st.button("Analizar Cambios", key="analyze_iteration"):
            if not st.session_state.arc_analysis_state["iteration_text"].strip():
                st.warning("Ingrese texto de iteración para analizar.")
            elif not st.session_state.arc_analysis_state["base_id"]:
                st.error("No existe un base_id aún (analice el texto base primero).")
            else:
                # Realizar el análisis con Spacy
                doc_iter = nlp_models[lang_code](
                    st.session_state.arc_analysis_state["iteration_text"]
                )
                display_arc_diagram(doc_iter)

                # Guardar iteración en DB
                iteration_id = store_student_morphosyntax_iteration(
                    st.session_state.username,
                    st.session_state.arc_analysis_state["base_id"],
                    st.session_state.arc_analysis_state["original_text"],
                    st.session_state.arc_analysis_state["iteration_text"],
                    arc_diagrams=None  # Pon el dict real con arcos
                )
                if iteration_id:
                    st.success(f"Iteración guardada. iteration_id={iteration_id}")
    
    # ==================== TAB 1: Análisis de Categorías ====================
    with tabs[1]:
        st.info("En desarrollo...")

    # ==================== TAB 2: Análisis Morfológico ======================
    with tabs[2]:
        st.info("En desarrollo...")