import shutil
from huggingface_hub import snapshot_download, hf_hub_download, list_repo_files, HfApi
from pathlib import Path
from huggingface_hub import list_repo_files, hf_hub_download
import streamlit as st
import os
from PIL import Image
import pandas as pd
import json
from datetime import datetime
from typing import List, Dict, Any, Optional
from RAG import EnhancedMultimodalRAGSystem
from config import *

# Page config
st.set_page_config(
    page_title="DTMI UGM Academic Assistant",
    page_icon="🎓",
    layout="wide",
    initial_sidebar_state="expanded"
)
st.markdown("""
<style>
/* Main Header */
.main-header {
    background: linear-gradient(135deg, #1e3c72 0%, #2a5298 50%, #4a90e2 100%);
    padding: 2rem;
    border-radius: 15px;
    color: white;
    text-align: center;
    margin-bottom: 2rem;
    box-shadow: 0 8px 32px rgba(0,0,0,0.1);
}

.main-header h1 {
    margin-bottom: 0.5rem;
    font-size: 2.5rem;
    font-weight: 700;
}

.main-header p {
    margin: 0.3rem 0;
    opacity: 0.9;
}

/* Chat Messages - Hitam Putih Simple */
.user-message {
    background: #2d2d2d;
    color: white;
    padding: 1.2rem;
    border-radius: 15px;
    margin: 1rem 0;
    border-left: 5px solid #0084ff;
    box-shadow: 0 4px 12px rgba(0,0,0,0.2);
    animation: slideInRight 0.3s ease-out;
}

.assistant-message {
    background: #f8f9fa;
    color: #2d2d2d;
    padding: 1.2rem;
    border-radius: 15px;
    margin: 1rem 0;
    border-left: 5px solid #28a745;
    box-shadow: 0 4px 12px rgba(0,0,0,0.1);
    animation: slideInLeft 0.3s ease-out;
}

@keyframes slideInRight {
    from { transform: translateX(20px); opacity: 0; }
    to { transform: translateX(0); opacity: 1; }
}

@keyframes slideInLeft {
    from { transform: translateX(-20px); opacity: 0; }
    to { transform: translateX(0); opacity: 1; }
}

/* Example Queries */
.example-query {
    background: #fff8e1;
    color: #333;
    padding: 1rem;
    border-radius: 10px;
    margin: 0.5rem 0;
    border-left: 4px solid #ff9800;
    cursor: pointer;
    transition: all 0.3s ease;
    box-shadow: 0 2px 8px rgba(255, 152, 0, 0.1);
}

.example-query:hover {
    background: #ffecb3;
    transform: translateY(-2px);
    box-shadow: 0 4px 12px rgba(255, 152, 0, 0.2);
}

/* Source Preview */
.source-preview {
    background: #f5f5f5;
    color: #333;
    padding: 1rem;
    border-radius: 10px;
    margin: 0.5rem 0;
    font-size: 0.9em;
    border-left: 3px solid #6c757d;
}

/* Buttons */
.stButton > button {
    border-radius: 10px !important;
    font-weight: 600 !important;
    transition: all 0.3s ease !important;
}

.stButton > button:hover {
    transform: translateY(-1px) !important;
    box-shadow: 0 4px 12px rgba(0,0,0,0.15) !important;
}
.chat-container {
    height: calc(100vh - 180px);
    overflow-y: auto;
    padding: 1rem;
    border: 1px solid #e0e0e0;
    border-radius: 10px;
    background-color: #fafafa;
    margin-bottom: 1rem;
}
.fixed-input {
    position: fixed;
    bottom: 2rem;
    width: 60%;
    max-width: 800px;
    left: 50%;
    transform: translateX(-50%);
    background-color: white;
    padding: 1rem;
    border-radius: 10px;
    box-shadow: 0 4px 16px rgba(0,0,0,0.1);
    z-index: 999;
}
.spacer {
    height: 120px; /* Tambahkan spacer agar konten tak tertutup input */
}
</style>
""", unsafe_allow_html=True)


@st.cache_resource
def initialize_rag_system():
    try:
        return EnhancedMultimodalRAGSystem()
    except Exception as e:
        st.error(f"❌ Error initializing RAG system: {e}")
        st.stop()


def display_example_queries():
    """Display clickable example queries"""
    st.markdown("### 💡 Contoh Pertanyaan")

    for category, queries in EXAMPLE_QUERIES.items():
        with st.expander(f"{category}", expanded=True):
            for query in queries:
                if st.button(f"💬 {query}", key=f"example_{hash(query)}", use_container_width=True):
                    st.session_state.user_input = query
                    st.rerun()


def display_tables_in_chat(table_data: List[Dict]):
    """Display tables directly in chat"""
    if not table_data:
        return

    st.markdown("### 📊 Tabel Data")

    for i, table_info in enumerate(table_data, 1):
        with st.expander(f"📊 {table_info['title']} (Hal. {table_info['page']}, {table_info['year']})", expanded=True):

            # Table metadata
            col1, col2, col3 = st.columns(3)
            with col1:
                st.metric("📄 Halaman", table_info['page'])
            with col2:
                st.metric("📅 Tahun", table_info['year'])
            with col3:
                st.metric("📊 Score", f"{table_info['score']:.3f}")
            # Display table data
            try:
                if table_info.get("data_type") == "dataframe" and isinstance(table_info["data"], pd.DataFrame):
                    st.dataframe(table_info["data"], use_container_width=True)
                    # Download CSV
                    csv_data = table_info["data"].to_csv(index=False)
                    st.download_button(
                        label="💾 Download CSV",
                        data=csv_data,
                        file_name=f"table_{i}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
                        mime="text/csv"
                    )

                elif table_info.get("data_type") == "json":
                    st.json(table_info["data"])

                    # Download JSON
                    json_str = json.dumps(table_info["data"], indent=2, ensure_ascii=False)
                    st.download_button(
                        label="💾 Download JSON",
                        data=json_str,
                        file_name=f"data_{i}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
                        mime="application/json"
                    )

                # Show description
                if table_info.get('description'):
                    st.markdown("**📝 Deskripsi:**")
                    st.text(table_info['description'])

            except Exception as e:
                st.error(f"❌ Error displaying table: {e}")


def display_single_image_compact(img_info: Dict, index: int):
    """Display single image in compact format - CLEAN VERSION"""
    try:
        image_path = img_info["path"]

        # Check if file exists
        if not os.path.exists(image_path):
            st.error(f"❌ Gambar {index} tidak ditemukan")
            return

        # Load and display image
        image = Image.open(image_path)

        # Display image with nice styling
        st.image(image,
                 caption=f"📖 {img_info.get('title', 'Gambar')} - Hal. {img_info.get('page', 'N/A')} ({img_info.get('year', 'N/A')})",
                 use_container_width=True)

        # Compact metadata
        col1, col2 = st.columns(2)
        with col1:
            st.metric("📊 Relevance Score", f"None")
            # {img_info.get('score', 0):.2f}")
        with col2:
            st.metric("📐 Ukuran", f"{image.width}×{image.height}px")

        # Expandable details
        with st.expander(f"📝 Detail Gambar {index}", expanded=False):
            if img_info.get('description'):
                st.markdown("**📄 Deskripsi:**")
                st.text(img_info['description'])
            if img_info.get('caption'):
                st.markdown("**💬 Caption:**")
                st.text(img_info['caption'])

    except Exception as e:
        st.error(f"❌ Error loading image {index}: {str(e)}")


def display_single_image_full(img_info: Dict):
    """Display single image in full format - CLEAN VERSION"""
    try:
        image_path = img_info["path"]

        if not os.path.exists(image_path):
            st.error("❌ Gambar tidak ditemukan")
            return

        # Load image
        image = Image.open(image_path)

        # Display with title
        st.markdown(f"### 🖼️ {img_info.get('title', 'Gambar')}")

        # Create columns for image and metadata
        col1, col2 = st.columns([3, 1])

        with col1:
            st.image(image, use_column_width=True)

        with col2:
            st.markdown("**📋 Informasi Gambar**")
            st.metric("📄 Halaman", img_info.get('page', 'N/A'))
            st.metric("📅 Tahun", img_info.get('year', 'N/A'))
            # st.metric("📊 Score", f"{img_info.get('score', 0):.3f}")
            st.metric("📐 Dimensi", f"{image.width} × {image.height}")

            # Download button
            with open(image_path, "rb") as file:
                st.download_button(
                    label="💾 Download Gambar",
                    data=file.read(),
                    file_name=os.path.basename(image_path),
                    mime="image/png",
                    use_container_width=True
                )

        # Description below image
        if img_info.get('description'):
            st.markdown("**📝 Deskripsi Gambar:**")
            st.info(img_info['description'])

        if img_info.get('caption'):
            st.markdown("**💬 Caption:**")
            st.info(img_info['caption'])

    except Exception as e:
        st.error(f"❌ Error loading image: {str(e)}")


def display_images_in_chat(image_data: List[Dict], show_details: bool = True):
    """Display images directly in chat - CLEAN VERSION"""
    if not image_data:
        return

    st.markdown("### 🖼️ Gambar Terkait")
    if len(image_data) == 1:
        st.markdown(f"*Ditemukan 1 gambar relevan*")
    else:
        st.markdown(f"*Ditemukan {len(image_data)} gambar relevan*")
    if len(image_data) > 1:
        cols = st.columns(min(len(image_data), 2))  # Max 2 columns
        for i, img_info in enumerate(image_data):
            with cols[i % 2]:
                display_single_image_compact(img_info, i+1)
    else:
        display_single_image_full(image_data[0])


def enhanced_chat_interface():
    if 'messages' not in st.session_state:
        st.session_state.messages = []
    if 'user_input' not in st.session_state:
        st.session_state.user_input = ""
    rag_system = initialize_rag_system()
    st.markdown("""
    <div class="main-header">
        <h1>🎓 DTMI UGM Academic Assistant</h1>
        <p>Asisten Cerdas Multimodal untuk Informasi Akademik DTMI UGM</p>
        <p>💬 Tanyakan apapun tentang kurikulum, silabus, gambar, dan tabel data</p>
    </div>
    """, unsafe_allow_html=True)

    # Sidebar with controls
    with st.sidebar:
        st.markdown("### ⚙️ Pengaturan")

        # Content type preferences
        st.markdown("### 🎯 Preferensi Konten")
        content_preferences = []
        for content_type, description in CONTENT_TYPE_DESCRIPTIONS.items():
            if st.checkbox(description, key=f"pref_{content_type}"):
                content_preferences.append(content_type)

        # Retrieval settings
        st.markdown("### 🔍 Pengaturan Pencarian")
        max_results = st.slider("Jumlah Konteks Maksimal", 5, 20, 10)

        # Display settings
        st.markdown("### 🎭 Tampilan")
        show_images_inline = st.checkbox("🖼️ Tampilkan Gambar", value=True)
        show_tables_inline = st.checkbox("📊 Tampilkan Tabel", value=True)
        compact_mode = st.checkbox("📱 Mode Kompak", value=False)

        # Chat statistics
        if st.session_state.messages:
            st.markdown("### 📊 Statistik")
            total_messages = len(st.session_state.messages)
            st.metric("💬 Total Pesan", total_messages)
            st.metric("🗣️ Percakapan", total_messages // 2)

        # Clear chat
        if st.button("🗑️ Hapus Chat", type="secondary", use_container_width=True):
            st.session_state.messages = []
            st.rerun()

    # Main chat area
    col1, col2 = st.columns([3, 1] if not compact_mode else [1, 0])

    with col1:
        # Display chat history

        for message in st.session_state.messages:
            if message["role"] == "user":
                st.markdown(f"""
                <div class="user-message">
                    <strong>👤 Anda:</strong><br>
                    {message["content"]}
                </div>
                """, unsafe_allow_html=True)
            else:
                st.markdown(f"""
                <div class="assistant-message">
                    <strong>🤖 Assistant:</strong><br>
                    {message["content"]}
                </div>
                """, unsafe_allow_html=True)

                # 🎯 DISPLAY MULTIMODAL CONTENT
                if "result_data" in message:
                    result_data = message["result_data"]

                    # Show quick stats if has multimodal content
                    if result_data.get("has_images") or result_data.get("has_tables"):
                        st.markdown("---")  # Separator

                        col_stats1, col_stats2, col_stats3 = st.columns(3)
                        with col_stats1:
                            st.metric("🖼️ Gambar", len(result_data.get("image_data", [])))
                        with col_stats2:
                            st.metric("📊 Tabel", len(result_data.get("table_data", [])))
                        with col_stats3:
                            st.metric("📚 Sumber", result_data.get("total_sources", 0))

                    # 🖼️ DISPLAY IMAGES
                    if show_images_inline and result_data.get("has_images"):
                        display_images_in_chat(result_data.get("image_data", []))

                    # 📊 DISPLAY TABLES
                    if show_tables_inline and result_data.get("has_tables"):
                        display_tables_in_chat(result_data.get("table_data", []))

                # Collapsible sources
                if "sources" in message and message["sources"]:
                    with st.expander("📚 Lihat Sumber Informasi", expanded=False):
                        for i, source in enumerate(message["sources"][:3], 1):
                            content_type = source['metadata']['content_type']
                            year = source['metadata'].get('year', 'N/A')
                            page = source['metadata'].get('page', 'N/A')

                            st.markdown(f"""
                            **📋 Sumber {i}:** {CONTENT_TYPE_DESCRIPTIONS.get(content_type, content_type)}  
                            **📅 Tahun:** {year} | **📄 Halaman:** {page}  
                            **📝 Preview:** {source['content'][:150]}...
                            """)
                            st.markdown("---")

        # Chat input
        user_input = st.chat_input(
            "💬 Tanyakan tentang kurikulum, gambar, tabel, atau informasi lainnya...", key="chat_input")

        # Handle example query selection
        if st.session_state.user_input:
            user_input = st.session_state.user_input
            st.session_state.user_input = ""

        # 🚀 PROCESS USER INPUT
        if user_input:
            # Add user message
            st.session_state.messages.append({"role": "user", "content": user_input})

            # Show loading
            with st.spinner("🔍 Mencari informasi relevan..."):
                try:
                    result_data = rag_system.query(
                        user_input,
                        k=max_results,
                        content_filter=content_preferences if content_preferences else None
                    )

                    # Save assistant message with complete data
                    assistant_message = {
                        "role": "assistant",
                        "content": result_data["answer"],
                        "sources": result_data["sources"],
                        "result_data": result_data
                    }

                    st.session_state.messages.append(assistant_message)

                except Exception as e:
                    st.error(f"❌ Terjadi kesalahan: {e}")
                    st.session_state.messages.append({
                        "role": "assistant",
                        "content": "Maaf, terjadi kesalahan dalam memproses pertanyaan Anda. Silakan coba lagi."
                    })

            st.rerun()

    # Sidebar dengan example queries (only if not compact)
    if not compact_mode:
        with col2:
            display_example_queries()

            # Quick actions
            st.markdown("### ⚡ Aksi Cepat")

            quick_actions = [
                ("🖼️ Cari Gambar", "Tampilkan gambar formulir atau diagram"),
                ("📊 Lihat Tabel", "Tabel kurikulum semester 1"),
                ("🎓 Info Program", "Informasi program studi teknik mesin"),
                ("📚 Silabus", "Silabus mata kuliah wajib")
            ]

            for label, query in quick_actions:
                if st.button(label, use_container_width=True):
                    st.session_state.user_input = query
                    st.rerun()
            if st.session_state.messages:
                st.markdown("### 📤 Export")
                if st.button("💾 Download Chat", use_container_width=True):
                    chat_export = ""
                    for msg in st.session_state.messages:
                        role = "User" if msg["role"] == "user" else "Assistant"
                        chat_export += f"**{role}:** {msg['content']}\n\n"

                    st.download_button(
                        label="📄 Download Markdown",
                        data=chat_export,
                        file_name=f"chat_dtmi_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md",
                        mime="text/markdown",
                        use_container_width=True
                    )


def main():
    enhanced_chat_interface()


if __name__ == "__main__":
    main()