import streamlit as st
import json
import os
from datetime import datetime
from dotenv import load_dotenv


# Import custom functions
from pipelines_functions import (
    generate_technical_content, correct_technical_content, refine_technical_content,
    generate_operational_content, correct_operational_content, refine_operational_content
)
from utils_functions import (
    validate_and_sanitize_topic, check_cache, save_to_cache, validate_and_select_urls, 
    get_collections, PipelineMetrics
)
from image_generation_functions import process_images_for_pipeline


load_dotenv()


# ============================================================
# PAGE CONFIGURATION
# ============================================================


st.set_page_config(
    page_title="LearnOnTheGo",
    page_icon="🎓",
    layout="wide",
    initial_sidebar_state="collapsed"
)


# ============================================================
# HELPER FUNCTIONS
# ============================================================


def sanitize_for_html(raw_text: str) -> str:
    """Escape HTML special characters for safe embedding."""
    if not isinstance(raw_text, str):
        raw_text = str(raw_text)
    return raw_text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")


def detect_code_content(text: str) -> bool:
    """Detect if content looks like code (has HTML tags, brackets, etc)."""
    code_indicators = ['<div', '<html', 'class=', 'style=', '{', '}', 'function', 'import', 'def ']
    return any(indicator in text for indicator in code_indicators)


# ============================================================
# CUSTOM CSS - FINAL VERSION WITH SEPARATED PROGRESS BAR
# ============================================================


st.markdown("""
<style>
    /* Root color palette */
    :root {
        --primary-blue: #2563eb;
        --accent-teal: #0891b2;
        --light-blue: #eff6ff;
        --light-teal: #e0f2fe;
        --text-dark: #1e293b;
        --text-light: #64748b;
        --border-color: #bae6fd;
        --shadow: 0 8px 24px rgba(37, 99, 235, 0.12);
        --shadow-lg: 0 20px 50px rgba(37, 99, 235, 0.25);
    }
    
    /* Overall app styling */
    .stApp {
        background: linear-gradient(135deg, #dbeafe 0%, #cffafe 100%);
    }
    
    /* Hide default Streamlit elements */
    #MainMenu {visibility: hidden;}
    footer {visibility: hidden;}
    header {visibility: hidden;}
    
    /* Main container */
    .block-container {
        max-width: 1400px;
        padding-top: 2rem;
        padding-bottom: 2rem;
    }
    
    /* Header styling */
    .header-container {
        background: linear-gradient(110deg, var(--primary-blue) 0%, var(--accent-teal) 100%);
        padding: 60px 30px;
        border-radius: 28px;
        text-align: center;
        margin-bottom: 40px;
        box-shadow: var(--shadow-lg);
    }
    
    .header-container h1 {
        color: white;
        font-size: 62px;
        margin: 0;
        font-weight: 900;
        letter-spacing: 3px;
        text-shadow: 0 4px 12px rgba(0, 0, 0, 0.2);
    }
    
    /* Search box styling */
    .search-container {
        background: linear-gradient(135deg, #ffffff 0%, #f8fafc 100%);
        padding: 35px;
        border-radius: 24px;
        margin-bottom: 30px;
        border: 3px solid var(--border-color);
        box-shadow: var(--shadow);
    }
    
    /* Text input styling */
    .stTextInput > div > div > input {
        border: 3px solid var(--accent-teal) !important;
        border-radius: 14px !important;
        padding: 16px 20px !important;
        font-size: 17px !important;
        background-color: white !important;
        transition: all 0.3s ease !important;
    }
    
    .stTextInput > div > div > input:focus {
        border-color: var(--primary-blue) !important;
        box-shadow: 0 0 0 4px rgba(37, 99, 235, 0.15) !important;
        outline: none !important;
    }
    
    .stTextInput > div > div > input::placeholder {
        color: rgba(100, 116, 139, 0.6) !important;
        font-weight: 500 !important;
    }
    
    /* Radio container */
    .stRadio > div[role="radiogroup"] {
        display: flex !important;
        gap: 0 !important;
        background: #e0e7ff !important;
        border-radius: 14px !important;
        padding: 4px !important;
        border: 3px solid var(--border-color) !important;
        width: fit-content !important;
        margin: 0 auto !important;
    }
    
    /* Individual radio labels */
    .stRadio > div[role="radiogroup"] > label {
        background: transparent !important;
        border: none !important;
        padding: 12px 32px !important;
        border-radius: 10px !important;
        cursor: pointer !important;
        transition: all 0.3s ease !important;
        font-weight: 700 !important;
        color: var(--text-light) !important;
        font-size: 15px !important;
        text-align: center !important;
        min-width: 140px !important;
        margin: 0 !important;
        flex: 1 !important;
    }
    
    .stRadio > div[role="radiogroup"] > label:hover {
        background: rgba(255, 255, 255, 0.5) !important;
    }
    
    .stRadio > div[role="radiogroup"] > label[data-checked="true"],
    .stRadio > div[role="radiogroup"] > label:has(input:checked),
    .stRadio > div[role="radiogroup"] > label[aria-checked="true"] {
        background: linear-gradient(110deg, var(--primary-blue) 0%, var(--accent-teal) 100%) !important;
        color: white !important;
        box-shadow: 0 4px 12px rgba(37, 99, 235, 0.35) !important;
    }
    
    .stRadio input[type="radio"] {
        display: none !important;
    }
    
    /* Button styling */
    .stButton > button {
        background: linear-gradient(100deg, var(--primary-blue) 0%, var(--accent-teal) 100%) !important;
        color: white !important;
        font-weight: 800 !important;
        padding: 18px 40px !important;
        border-radius: 14px !important;
        border: none !important;
        transition: all 0.3s ease !important;
        box-shadow: 0 6px 20px rgba(37, 99, 235, 0.3) !important;
        font-size: 17px !important;
        letter-spacing: 0.5px !important;
    }
    
    .stButton > button:hover {
        background: linear-gradient(100deg, #1d4ed8 0%, #0e7490 100%) !important;
        transform: translateY(-3px) !important;
        box-shadow: 0 10px 30px rgba(37, 99, 235, 0.4) !important;
    }
    
    /* SLIDE BOX - WITHOUT PROGRESS BAR */
    .slide-box-wrapper {
        background: white;
        border-radius: 28px;
        border: 3px solid var(--border-color);
        box-shadow: 0 20px 60px rgba(37, 99, 235, 0.18);
        padding: 50px 45px;
        margin: 40px auto;
        max-width: 1000px;
        animation: slideIn 0.5s ease-out;
    }
    
    @keyframes slideIn {
        from {
            opacity: 0;
            transform: translateY(30px);
        }
        to {
            opacity: 1;
            transform: translateY(0);
        }
    }
    
    /* Slide content */
    .slide-box-wrapper > * {
        display: block !important;
        width: 100% !important;
        box-sizing: border-box !important;
    }
    
    /* Slide title */
    .slide-box-wrapper h2 {
        font-size: 42px;
        font-weight: 900;
        color: var(--primary-blue);
        margin: 0 0 20px 0 !important;
        letter-spacing: 1px;
        line-height: 1.3;
        text-shadow: 0 2px 8px rgba(37, 99, 235, 0.15);
        word-wrap: break-word;
        overflow-wrap: break-word;
        text-align: center;
    }
    
    /* Slide text and paragraphs */
    .slide-box-wrapper p {
        font-size: 20px;
        color: var(--text-dark);
        line-height: 2.2;
        margin: 0 0 24px 0 !important;
        font-weight: 500;
        text-align: left;
        padding: 0 20px;
        box-sizing: border-box;
    }
    
    /* Code block styling */
    .slide-box-wrapper pre {
        background: #f8fafc;
        border-radius: 12px;
        padding: 20px;
        font-family: 'Consolas', 'Monaco', 'Courier New', monospace;
        font-size: 14px;
        color: #0f172a;
        overflow-x: auto;
        border: 2px solid #e0e7ff;
        margin: 16px 20px !important;
        text-align: left;
        line-height: 1.6;
        white-space: pre-wrap;
        word-wrap: break-word;
    }
    
    /* Images inside slide box */
    .slide-box-wrapper img {
        max-width: 90% !important;
        height: auto !important;
        border-radius: 20px !important;
        box-shadow: 0 8px 24px rgba(37, 99, 235, 0.2) !important;
        display: block !important;
        margin: 24px auto !important;
    }
    
    /* Learning Resources */
    .resources-section {
        margin-top: 32px;
        padding-top: 28px;
        border-top: 3px solid var(--border-color);
        text-align: left;
    }
    
    .resources-section h4 {
        color: var(--primary-blue);
        font-size: 22px;
        margin: 0 0 18px 0 !important;
        font-weight: 800;
        text-align: center;
    }
    
    .resources-section a {
        color: var(--accent-teal);
        text-decoration: none;
        font-weight: 600;
        transition: all 0.3s ease;
        display: block;
        padding: 10px 15px;
        font-size: 16px;
        border-radius: 8px;
        margin-bottom: 8px;
    }
    
    .resources-section a:hover {
        color: white;
        background: var(--primary-blue);
        padding-left: 20px;
        box-shadow: 0 4px 12px rgba(37, 99, 235, 0.2);
    }
    
    /* PROGRESS CONTAINER - MOVED OUTSIDE BOX */
    .progress-container {
        display: flex;
        align-items: center;
        justify-content: center;
        gap: 20px;
        margin: 30px auto;
        padding: 20px 0;
        max-width: 1000px;
        width: 100%;
    }
    
    .progress-bar {
        flex: 1;
        max-width: 700px;
        height: 10px;
        background: #e0e7ff;
        border-radius: 12px;
        overflow: hidden;
        box-shadow: inset 0 2px 4px rgba(37, 99, 235, 0.1);
    }
    
    .progress-fill {
        height: 100%;
        background: linear-gradient(90deg, var(--primary-blue) 0%, var(--accent-teal) 100%);
        transition: width 0.4s ease;
        box-shadow: 0 0 10px rgba(37, 99, 235, 0.4);
    }
    
    /* Slide counter badge */
    .slide-counter-badge {
        background: linear-gradient(110deg, var(--primary-blue) 0%, var(--accent-teal) 100%);
        color: white;
        padding: 10px 20px;
        border-radius: 24px;
        font-size: 16px;
        font-weight: 800;
        min-width: 90px;
        text-align: center;
        box-shadow: 0 4px 12px rgba(37, 99, 235, 0.3);
        white-space: nowrap;
    }
    
    /* Footer */
    .footer-bar {
        margin-top: 60px;
        text-align: center;
        color: var(--text-dark);
        font-size: 16px;
        letter-spacing: 0.5px;
        padding: 32px 0;
        border-top: 3px solid var(--border-color);
        background: white;
        border-radius: 20px;
        box-shadow: 0 4px 16px rgba(37, 99, 235, 0.08);
    }
    
    .footer-bar p {
        margin: 10px 0;
        font-weight: 600;
    }
    
    .footer-bar p:last-child {
        font-size: 14px;
        color: var(--text-light);
        font-weight: 500;
    }
    
    /* Alert messages */
    .stSuccess {
        border-radius: 14px !important;
        border-left: 5px solid #10b981 !important;
        background-color: #ecfdf5 !important;
        padding: 16px !important;
        font-weight: 600 !important;
    }
    
    .stInfo {
        border-radius: 14px !important;
        border-left: 5px solid var(--primary-blue) !important;
        background-color: #f0f9ff !important;
        padding: 16px !important;
        font-weight: 600 !important;
    }
    
    .stError {
        border-radius: 14px !important;
        border-left: 5px solid #ef4444 !important;
        background-color: #fef2f2 !important;
        padding: 16px !important;
        font-weight: 600 !important;
    }
    
    /* Mobile responsive */
    @media (max-width: 768px) {
        .header-container h1 {
            font-size: 40px;
        }
        
        .slide-box-wrapper {
            padding: 30px 20px;
        }
        
        .slide-box-wrapper h2 {
            font-size: 30px;
        }
        
        .slide-box-wrapper p {
            font-size: 17px;
            line-height: 1.9;
            padding: 0 10px;
        }
        
        .progress-container {
            flex-direction: row;
            gap: 15px;
            margin: 20px auto;
            padding: 15px 0;
        }
        
        .progress-bar {
            width: 100%;
            max-width: none;
        }
    }
</style>
""", unsafe_allow_html=True)


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


if "current_slide" not in st.session_state:
    st.session_state.current_slide = 0
if "slides_data" not in st.session_state:
    st.session_state.slides_data = None
if "search_query" not in st.session_state:
    st.session_state.search_query = ""
if "mode" not in st.session_state:
    st.session_state.mode = "technical"
if "is_loading" not in st.session_state:
    st.session_state.is_loading = False
if "error_message" not in st.session_state:
    st.session_state.error_message = None
if "metrics" not in st.session_state:
    st.session_state.metrics = None


# ============================================================
# PIPELINE FUNCTION
# ============================================================


def run_pipeline(query, mode):
    """Execute the 5-stage pipeline with metrics tracking."""
    try:
        metrics = PipelineMetrics(query, mode)
        query = validate_and_sanitize_topic(query)
        
        technical_col, operational_col, db = get_collections()
        collection = operational_col if mode == "operational" else technical_col
        
        # Cache check
        metrics.start_stage("Cache Check")
        cached_content, is_cached = check_cache(query, collection)
        if is_cached:
            metrics.set_cache_hit("mongodb")
        metrics.end_stage("Cache Check")
        
        if is_cached:
            st.session_state.slides_data = cached_content
            st.session_state.current_slide = 0
            metrics.end()
            metrics.save_metrics()
            return True, "✅ Retrieved from cache (instant!)"
        
        st.session_state.is_loading = True
        
        with st.spinner(f"🔄 Generating {mode} content with images (5 stages)..."):
            if mode == "technical":
                metrics.start_stage("Generate")
                generated = generate_technical_content(query)
                metrics.end_stage("Generate", f"{len(generated.get('content', []))} slides")
                
                metrics.start_stage("Correct")
                corrected = correct_technical_content(generated)
                metrics.end_stage("Correct", "Content improved")
                
                metrics.start_stage("Validate URLs")
                validated, _ = validate_and_select_urls(corrected)
                metrics.end_stage("Validate URLs", f"{len(validated.get('urls', []))} URLs validated")
                
                metrics.start_stage("Refine")
                refined = refine_technical_content(validated)
                metrics.end_stage("Refine", "Content refined")
                
                metrics.start_stage("Generate Images")
                final_result = process_images_for_pipeline(refined, mode="technical")
                metrics.end_stage("Generate Images", "Images generated")
            else:
                metrics.start_stage("Generate")
                generated = generate_operational_content(query)
                metrics.end_stage("Generate", f"{len(generated.get('content', []))} slides")
                
                metrics.start_stage("Correct")
                corrected = correct_operational_content(generated)
                metrics.end_stage("Correct", "Content improved")
                
                metrics.start_stage("Validate URLs")
                validated, _ = validate_and_select_urls(corrected)
                metrics.end_stage("Validate URLs", f"{len(validated.get('urls', []))} URLs validated")
                
                metrics.start_stage("Refine")
                refined = refine_operational_content(validated)
                metrics.end_stage("Refine", "Content refined")
                
                metrics.start_stage("Generate Images")
                final_result = process_images_for_pipeline(refined, mode="operational")
                metrics.end_stage("Generate Images", "Images generated")
            
            save_to_cache(query, final_result, collection)
            st.session_state.slides_data = final_result
            st.session_state.current_slide = 0
            st.session_state.is_loading = False
        
        pipeline_metrics = metrics.end()
        metrics.save_metrics()
        st.session_state.metrics = pipeline_metrics
        
        total_time = pipeline_metrics.get('total_duration_seconds', 0)
        return True, f"✅ Generated {len(final_result.get('content', []))} slides in {total_time:.1f}s!"
        
    except Exception as e:
        st.session_state.is_loading = False
        st.session_state.error_message = str(e)
        return False, f"❌ Error: {str(e)}"


# ============================================================
# DISPLAY SLIDE FUNCTION - PROGRESS BAR OUTSIDE BOX
# ============================================================


def display_slide(slide_index):
    """Display current slide with progress bar OUTSIDE the white box."""
    if not st.session_state.slides_data:
        return
    
    slides = st.session_state.slides_data.get('content', [])
    if not slides or slide_index >= len(slides):
        return
    
    slide = slides[slide_index]
    total_slides = len(slides)
    progress_percent = ((slide_index + 1) / total_slides) * 100
    
    # Build the slide box HTML (WITHOUT progress bar)
    title = sanitize_for_html(slide.get("slide_title", ""))
    raw_content = slide.get("slide_content", "")
    
    # Determine content type
    if detect_code_content(raw_content):
        sanitized_content = f"<pre>{sanitize_for_html(raw_content)}</pre>"
    else:
        sanitized_content = f"<p>{sanitize_for_html(raw_content)}</p>"
    
    # Start building the slide HTML (NO progress bar inside)
    slide_html = f"""
    <div class="slide-box-wrapper">
        <h2>{title}</h2>
        {sanitized_content}
    """
    
    # Add image if available
    img_url = slide.get('image_description')
    if isinstance(img_url, str) and img_url.startswith('http'):
        slide_html += f'<img src="{img_url}" alt="Slide image" style="max-width: 90%; height: auto; display: block; margin: 24px auto; border-radius: 20px; box-shadow: 0 8px 24px rgba(37, 99, 235, 0.2);">'
    
    # Add learning resources (last slide only)
    if slide_index == total_slides - 1:
        urls = st.session_state.slides_data.get('urls', [])
        if urls:
            slide_html += '<div class="resources-section"><h4>📚 Learning Resources</h4>'
            for i, url_obj in enumerate(urls, 1):
                url_title = sanitize_for_html(url_obj.get('title', 'Documentation'))
                url = url_obj.get('url', '#')
                slide_html += f'<a href="{url}" target="_blank">{i}. {url_title}</a>'
            slide_html += '</div>'
    
    # Close the slide box (NO progress bar)
    slide_html += '</div>'
    
    # Render the slide box
    st.markdown(slide_html, unsafe_allow_html=True)
    
    # RENDER PROGRESS BAR OUTSIDE THE BOX
    progress_html = f"""
    <div class="progress-container">
        <div class="progress-bar">
            <div class="progress-fill" style="width: {progress_percent}%"></div>
        </div>
        <div class="slide-counter-badge">{slide_index + 1} / {total_slides}</div>
    </div>
    """
    st.markdown(progress_html, unsafe_allow_html=True)
    
    # Navigation buttons below
    st.markdown('<br>', unsafe_allow_html=True)
    col_left, col_center, col_right = st.columns([1, 8, 1])
    
    with col_left:
        if slide_index > 0:
            if st.button("⬅", key="prev_btn", help="Previous slide", use_container_width=True):
                st.session_state.current_slide -= 1
                st.rerun()
    
    with col_right:
        if slide_index < total_slides - 1:
            if st.button("➡", key="next_btn", help="Next slide", use_container_width=True):
                st.session_state.current_slide += 1
                st.rerun()


# ============================================================
# PAGE LAYOUT
# ============================================================


# Header
st.markdown(
    '<div class="header-container"><h1>🎓 LearnOnTheGo</h1></div>',
    unsafe_allow_html=True
)

# Search container
st.markdown('<div class="search-container">', unsafe_allow_html=True)

col1, col2 = st.columns([3, 1])

with col1:
    search_query = st.text_input(
        "Search",
        value=st.session_state.search_query,
        placeholder="e.g., Python, Machine Learning, Cloud Computing...",
        key="search_input",
        label_visibility="collapsed"
    )
    st.session_state.search_query = search_query

with col2:
    mode = st.radio(
        "Mode",
        options=["Technical", "Operational"],
        index=0 if st.session_state.mode == "technical" else 1,
        key="mode_radio",
        horizontal=True,
        label_visibility="collapsed"
    )
    st.session_state.mode = mode.lower()

st.markdown('</div>', unsafe_allow_html=True)

# Generate button
col1, col2, col3 = st.columns([1, 2, 1])
with col2:
    search_button = st.button("🔍 Generate Slides", key="search_btn", use_container_width=True)

# Error handling
if st.session_state.error_message:
    st.error(st.session_state.error_message)
    st.session_state.error_message = None

# Execute pipeline
if search_button and st.session_state.search_query:
    success, message = run_pipeline(st.session_state.search_query, st.session_state.mode)
    if success:
        st.success(message)
    else:
        st.error(message)

# Display slides
if st.session_state.slides_data:
    st.markdown("---")
    if st.session_state.current_slide >= len(st.session_state.slides_data.get('content', [])):
        st.session_state.current_slide = 0
    display_slide(st.session_state.current_slide)
else:
    st.info("👆 Enter a topic and click 'Generate Slides' to get started!")

# Footer
st.markdown(
    """<div class="footer-bar">
    <p><strong>LearnOnTheGo</strong> • Powered by AI • Built with Streamlit</p>
    <p>5-Stage Pipeline: Generate → Correct → Validate → Refine → Generate Images</p>
    <p>Gemini 2.5 Flash (text) • Gemini 2.5 Flash Image (images) • Perplexity Sonar Pro</p>
</div>""",
    unsafe_allow_html=True
)

print("✅ LearnOnTheGo - Progress bar moved outside box - Fixed!")