import streamlit as st
import json
import os
import sys
import traceback
import logging
import time
from datetime import datetime

# =====================================================
# PATH SETUP
# =====================================================

sys.path.append(os.getcwd())

# =====================================================
# LOGGING CONFIG
# =====================================================

LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO").upper()

logging.basicConfig(
    level=getattr(logging, LOG_LEVEL, logging.INFO),
    format="%(asctime)s | %(levelname)s | %(message)s",
    handlers=[
        logging.StreamHandler(sys.stdout)
    ]
)

logger = logging.getLogger("TLC_AGENT")

print("\n")
print("=" * 80)
print("🚀 TLC AGENT SCIENTIFIQUE - STARTUP")
print("=" * 80)
print(f"📍 Working directory : {os.getcwd()}")
print(f"📍 Python version    : {sys.version}")
print(f"📍 Time              : {datetime.now()}")
print("=" * 80)
print("\n")

# =====================================================
# SAFE IMPORTS
# =====================================================

try:
    print("📦 Import extraction.py ...")
    from extraction import extract_from_file, extract_from_text
    print("✅ extraction.py loaded")
except Exception as e:
    print("❌ FAILED loading extraction.py")
    print(traceback.format_exc())
    raise e

try:
    print("📦 Import ir_builder.py ...")
    from ir_builder import build_ir_variants
    print("✅ ir_builder.py loaded")
except Exception as e:
    print("❌ FAILED loading ir_builder.py")
    print(traceback.format_exc())
    raise e

try:
    print("📦 Import pattern_detector.py ...")
    from pattern_detector import detect_patterns
    print("✅ pattern_detector.py loaded")
except Exception as e:
    print("❌ FAILED loading pattern_detector.py")
    print(traceback.format_exc())
    raise e

try:
    print("📦 Import optimizer.py ...")
    from optimizer import optimize_ir
    print("✅ optimizer.py loaded")
except Exception as e:
    print("❌ FAILED loading optimizer.py")
    print(traceback.format_exc())
    raise e

# =====================================================
# PAGE CONFIG
# =====================================================

st.set_page_config(
    page_title="TLC Agent Scientifique",
    page_icon="🧠",
    layout="wide"
)

# =====================================================
# CUSTOM CSS
# =====================================================

st.markdown("""
<style>
.block-container {
    padding-top: 2rem;
}

.stCodeBlock {
    border-radius: 12px;
}

.log-box {
    background: #111;
    color: #0f0;
    padding: 12px;
    border-radius: 10px;
    font-family: monospace;
    font-size: 12px;
    overflow-x: auto;
    white-space: pre-wrap;
}

.debug-title {
    font-weight: bold;
    margin-bottom: 10px;
}
</style>
""", unsafe_allow_html=True)

# =====================================================
# TITLE
# =====================================================

st.title("🧠 TLC Agent – Cercle Scientifique")
st.caption("Extraction → IR → Patterns → Optimisation")

# =====================================================
# DEBUG HELPERS
# =====================================================

def debug_log(message, level="INFO"):
    timestamp = datetime.now().strftime("%H:%M:%S")
    formatted = f"[{timestamp}] [{level}] {message}"
    print(formatted)

    if "debug_logs" not in st.session_state:
        st.session_state.debug_logs = []

    st.session_state.debug_logs.append(formatted)

    if level == "ERROR":
        logger.error(message)
    elif level == "WARNING":
        logger.warning(message)
    else:
        logger.info(message)


def debug_exception(e):
    err = traceback.format_exc()

    print("\n")
    print("=" * 80)
    print("❌ EXCEPTION")
    print("=" * 80)
    print(err)
    print("=" * 80)
    print("\n")

    logger.exception(str(e))
    return err


def timed_call(label, fn, *args, **kwargs):
    debug_log(f"START => {label}")
    start = time.time()
    result = fn(*args, **kwargs)
    duration = round(time.time() - start, 2)
    debug_log(f"END => {label} ({duration}s)")
    return result


def normalize_provider(provider: str) -> str:
    """
    UI provider -> provider attendu par les modules internes.
    OpenAI UI = Azure OpenAI interne.
    """
    provider = (provider or "").lower().strip()
    if provider == "openai":
        return "azure"
    return provider


# =====================================================
# SESSION STATE
# =====================================================

DEFAULT_STATE = {
    "step": "input",
    "equations": [],
    "ir_variants": [],
    "chosen_ir": None,
    "patterns": [],
    "optimization_variants": [],
    "chosen_optimization": None,
    "provider": "openai",
    "debug_logs": [],
    "debug_mode": True,
}

for key, value in DEFAULT_STATE.items():
    if key not in st.session_state:
        st.session_state[key] = value

# =====================================================
# START LOG
# =====================================================

debug_log(f"APP RERUN | CURRENT STEP = {st.session_state.step}")

# =====================================================
# RESET
# =====================================================

def reset():
    debug_log("RESET SESSION")
    for key, value in DEFAULT_STATE.items():
        st.session_state[key] = value

# =====================================================
# SIDEBAR
# =====================================================

with st.sidebar:
    st.header("📋 Pipeline")
    st.info(f"Étape actuelle : {st.session_state.step}")
    st.divider()

    debug_mode = st.toggle("🐞 Debug mode", value=st.session_state.debug_mode)
    st.session_state.debug_mode = debug_mode

    st.divider()

    if st.button("🔄 Nouvelle session", use_container_width=True):
        debug_log("NEW SESSION CLICKED")
        reset()
        st.rerun()

    st.divider()

    st.markdown("""
### Workflow

1. Input scientifique  
2. Extraction  
3. Construction IR  
4. Détection patterns  
5. Optimisation  
6. Export JSON
""")

    st.divider()

    st.subheader("🔑 API Keys")

    deepseek_ok = bool(os.getenv("DEEPSEEK_API_KEY"))
    groq_ok = bool(os.getenv("GROQ_API_KEY"))
    openai_ok = all([
        os.getenv("AZUREOPENAI_API_KEY"),
        os.getenv("AZUREOPENAI_API_ENDPOINT"),
        os.getenv("AZUREOPENAI_API_VERSION"),
        os.getenv("OPENAI_MODEL"),
    ])

    st.write(f"OpenAI : {'✅' if openai_ok else '❌'}")
    st.write(f"DeepSeek : {'✅' if deepseek_ok else '❌'}")
    st.write(f"Groq : {'✅' if groq_ok else '❌'}")

    debug_log(
        f"API STATUS => OPENAI={openai_ok} | DEEPSEEK={deepseek_ok} | GROQ={groq_ok}"
    )

# =====================================================
# DEBUG PANEL
# =====================================================

if st.session_state.get("debug_mode", False):
    with st.expander("🐞 Debug Console", expanded=False):
        logs = "\n".join(st.session_state.debug_logs[-200:])
        st.markdown(f"""
<div class="log-box">
{logs}
</div>
""", unsafe_allow_html=True)

# =====================================================
# STEP 1
# =====================================================

if st.session_state.step == "input":
    debug_log("ENTER STEP INPUT")

    st.header("1. Fournir le contenu scientifique")

    providers = ["openai", "deepseek", "groq"]

    provider = st.selectbox(
        "LLM Provider",
        providers,
        index=providers.index(st.session_state.provider)
    )
    st.session_state.provider = provider

    effective_provider = normalize_provider(provider)
    debug_log(f"PROVIDER SELECTED => {provider} (effective: {effective_provider})")

    mode = st.radio(
        "Mode d'entrée",
        ["Fichier (PDF/DOCX/TXT)", "Texte long"],
        horizontal=True
    )
    debug_log(f"INPUT MODE => {mode}")

    uploaded_file = None
    content = ""

    if mode == "Fichier (PDF/DOCX/TXT)":
        uploaded_file = st.file_uploader(
            "Choisir un fichier",
            type=["pdf", "docx", "txt"]
        )
        st.caption("Formats acceptés : PDF, DOCX, TXT")

        if uploaded_file is not None:
            debug_log(f"FILE UPLOADED => {uploaded_file.name}")
            debug_log(f"FILE SIZE => {uploaded_file.size} bytes")

    else:
        content = st.text_area(
            "Texte scientifique",
            height=350,
            placeholder="""
Exemple :

$$
E = mc^2
$$

ou texte scientifique brut.
"""
        )
        debug_log(f"TEXT LENGTH => {len(content)}")

    st.divider()

    if st.button("🚀 Lancer l'extraction", type="primary"):
        debug_log("EXTRACTION BUTTON CLICKED")

        try:
            equations = []

            if uploaded_file is not None:
                debug_log("START FILE EXTRACTION")
                with st.spinner("📄 Extraction du fichier..."):
                    equations = timed_call(
                        "extract_from_file",
                        extract_from_file,
                        uploaded_file
                    )
                debug_log(f"FILE EXTRACTION DONE => {len(equations)} equations")

            elif content.strip():
                debug_log("START TEXT EXTRACTION")
                with st.spinner("🧠 Analyse scientifique..."):
                    equations = timed_call(
                        "extract_from_text",
                        extract_from_text,
                        content
                    )
                debug_log(f"TEXT EXTRACTION DONE => {len(equations)} equations")

            else:
                debug_log("NO INPUT PROVIDED", level="WARNING")
                st.warning("Veuillez fournir du contenu.")
                st.stop()

            debug_log(f"EQUATIONS TYPE => {type(equations)}")

            if equations:
                debug_log(f"FIRST EQUATION => {str(equations[0])[:300]}")

            if len(equations) > 0:
                st.session_state.equations = equations
                st.success(f"✅ {len(equations)} équation(s) détectée(s)")
                debug_log("GO TO STEP EXTRACTION")
                st.session_state.step = "extraction"
                st.rerun()
            else:
                debug_log("NO EQUATIONS DETECTED", level="ERROR")
                st.error("❌ Aucune équation détectée.")

        except Exception as e:
            err = debug_exception(e)
            st.error(str(e))
            st.code(err)

# =====================================================
# STEP 2
# =====================================================

elif st.session_state.step == "extraction":
    debug_log("ENTER STEP EXTRACTION")

    st.header("2. Équations extraites")

    equations = st.session_state.equations

    debug_log(f"DISPLAYING {len(equations)} EQUATIONS")

    for i, eq in enumerate(equations):
        debug_log(f"RENDER EQUATION {i+1}")

        with st.expander(f"Équation {i+1}", expanded=True):
            st.code(eq.get("latex", ""), language="latex")

            context = eq.get("context", "")
            if context:
                st.caption(context[:300])

    st.divider()

    col1, col2 = st.columns(2)

    with col1:
        if st.button("✅ Valider", use_container_width=True):
            debug_log("VALIDATE EXTRACTION")
            st.session_state.step = "ir"
            st.rerun()

    with col2:
        if st.button("❌ Recommencer", use_container_width=True):
            debug_log("RESTART FROM EXTRACTION")
            reset()
            st.rerun()

# =====================================================
# STEP 3
# =====================================================

elif st.session_state.step == "ir":
    debug_log("ENTER STEP IR")

    st.header("3. Construction IR")

    if len(st.session_state.ir_variants) == 0:
        debug_log("NO IR IN CACHE => BUILDING")

        try:
            with st.spinner("⚙️ Génération des IR..."):
                st.session_state.ir_variants = timed_call(
                    "build_ir_variants",
                    build_ir_variants,
                    st.session_state.equations,
                    num_variants=3,
                    provider=normalize_provider(st.session_state.provider)
                )

            debug_log(f"IR VARIANTS GENERATED => {len(st.session_state.ir_variants)}")

        except Exception as e:
            err = debug_exception(e)
            st.error(str(e))
            st.code(err)
            st.stop()

    ir_variants = st.session_state.ir_variants

    for i, ir in enumerate(ir_variants):
        debug_log(f"DISPLAY IR VARIANT {i+1}")
        with st.expander(f"IR Variante {i+1}", expanded=(i == 0)):
            st.json(ir)

    choice = st.radio(
        "Choisir une IR",
        range(len(ir_variants)),
        format_func=lambda x: f"Variante {x+1}"
    )

    debug_log(f"IR CHOICE => {choice}")

    if st.button("✅ Valider cette IR", type="primary"):
        debug_log(f"IR VALIDATED => VARIANT {choice+1}")
        st.session_state.chosen_ir = ir_variants[choice]
        st.session_state.step = "patterns"
        st.rerun()

# =====================================================
# STEP 4
# =====================================================

elif st.session_state.step == "patterns":
    debug_log("ENTER STEP PATTERNS")

    st.header("4. Détection des patterns")

    if len(st.session_state.patterns) == 0:
        debug_log("START PATTERN DETECTION")

        try:
            with st.spinner("🔍 Analyse des patterns..."):
                st.session_state.patterns = timed_call(
                    "detect_patterns",
                    detect_patterns,
                    st.session_state.chosen_ir,
                    provider=normalize_provider(st.session_state.provider)
                )

            debug_log(f"PATTERNS FOUND => {len(st.session_state.patterns)}")

        except Exception as e:
            err = debug_exception(e)
            st.error(str(e))
            st.code(err)
            st.stop()

    patterns = st.session_state.patterns

    for pattern in patterns:
        debug_log(f"PATTERN => {pattern.get('name', 'Unknown')}")
        with st.container(border=True):
            st.subheader(pattern.get("name", "Unknown"))
            st.write(pattern.get("description", ""))

    st.divider()

    col1, col2 = st.columns(2)

    with col1:
        if st.button("✅ Continuer", use_container_width=True):
            debug_log("GO TO OPTIMIZATION")
            st.session_state.step = "optimization"
            st.rerun()

    with col2:
        if st.button("⬅ Retour", use_container_width=True):
            debug_log("BACK TO IR STEP")
            st.session_state.step = "ir"
            st.rerun()

# =====================================================
# STEP 5
# =====================================================

elif st.session_state.step == "optimization":
    debug_log("ENTER STEP OPTIMIZATION")

    st.header("5. Optimisation IR")

    if len(st.session_state.optimization_variants) == 0:
        debug_log("START OPTIMIZATION")

        try:
            with st.spinner("🚀 Optimisation..."):
                st.session_state.optimization_variants = timed_call(
                    "optimize_ir",
                    optimize_ir,
                    st.session_state.chosen_ir,
                    provider=normalize_provider(st.session_state.provider)
                )

            debug_log(f"OPTIMIZATION VARIANTS => {len(st.session_state.optimization_variants)}")

        except Exception as e:
            err = debug_exception(e)
            st.error(str(e))
            st.code(err)
            st.stop()

    optimizations = st.session_state.optimization_variants

    for i, opt in enumerate(optimizations):
        debug_log(f"DISPLAY OPTIMIZATION {i+1}")
        with st.expander(f"Stratégie {i+1}", expanded=(i == 0)):
            st.markdown(f"**Explication :** {opt.get('explanation', '')}")
            st.json(opt.get("optimized_ir", {}))

    choice = st.radio(
        "Choisir une optimisation",
        range(len(optimizations)),
        format_func=lambda x: f"Stratégie {x+1}"
    )

    debug_log(f"OPTIMIZATION CHOICE => {choice}")

    if st.button("✅ Finaliser", type="primary"):
        debug_log(f"FINAL OPTIMIZATION SELECTED => {choice+1}")
        st.session_state.chosen_optimization = optimizations[choice]
        st.session_state.step = "done"
        st.rerun()

# =====================================================
# STEP 6
# =====================================================

elif st.session_state.step == "done":
    debug_log("ENTER STEP DONE")

    st.header("🎉 Pipeline terminé")

    final_ir = (
        st.session_state.chosen_optimization.get(
            "optimized_ir",
            st.session_state.chosen_ir
        )
    )

    debug_log(f"FINAL IR NODES => {len(final_ir.get('nodes', []))}")
    debug_log(f"FINAL IR EDGES => {len(final_ir.get('edges', []))}")

    st.success("IR finale générée avec succès.")
    st.json(final_ir)

    json_data = json.dumps(
        final_ir,
        indent=2,
        ensure_ascii=False
    )

    st.download_button(
        label="📥 Télécharger JSON",
        data=json_data,
        file_name="final_ir.json",
        mime="application/json"
    )

    st.divider()

    if st.button("🔄 Nouveau traitement"):
        debug_log("NEW PROCESS STARTED")
        reset()
        st.rerun()

# =====================================================
# FOOTER DEBUG
# =====================================================

debug_log(f"END RENDER | STEP={st.session_state.step}")