from huggingface_hub import list_repo_files, hf_hub_download, snapshot_download
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
                    )


REPO_ID = "wicaksonolxn/dataoptim"
REPO_TYPE = "dataset"
LOCALROOT = Path("./repo_data")  # Changed to a specific subfolder
MARKER = LOCALROOT / ".hf_mirror_complete"


def ensure_cache_dir():
    cache_dir = Path.home() / ".cache" / "huggingface"
    try:
        cache_dir.mkdir(parents=True, exist_ok=True)
        test_file = cache_dir / "test_write"
        test_file.touch()
        test_file.unlink()
        return True
    except Exception as e:
        st.warning(f"Cannot create cache dir at {cache_dir}: {e}")
        # Try alternative cache location
        alt_cache = Path("./cache/huggingface")
        try:
            alt_cache.mkdir(parents=True, exist_ok=True)
            os.environ["HF_HOME"] = str(alt_cache.parent)
            return True
        except Exception as e2:
            st.error(f"Cannot create alternative cache: {e2}")
            return False


def mirror_repo_fixed(repo_id: str, repo_type: str, local_root: Path) -> bool:
    try:
        if not ensure_cache_dir():
            return False

        # Ensure local root exists
        local_root.mkdir(parents=True, exist_ok=True)

        st.write(f"⏬ Downloading repository '{repo_id}' to {local_root}")
        downloaded_path = snapshot_download(
            repo_id=repo_id,
            repo_type=repo_type,
            local_dir=str(local_root),
            local_dir_use_symlinks=False,
            resume_download=True,
            force_download=False,
        )

        # Create completion marker
        MARKER.touch()
        st.success(f"✅ Repository downloaded successfully to {downloaded_path}")
        return True

    except Exception as e:
        st.error(f"❌ Error downloading repository: {e}")
        return False


def mirror_repo_file_by_file(repo_id: str, repo_type: str, local_root: Path) -> bool:
    """
    Alternative file-by-file download method - more control but slower
    """
    try:
        # Ensure cache directory exists
        if not ensure_cache_dir():
            return False

        # Get list of files
        remote_files = list_repo_files(repo_id=repo_id, repo_type=repo_type)

        if not remote_files:
            st.error(f"No files found in repository '{repo_id}'")
            return False

        st.write(f"⏬ Copying **{len(remote_files)}** files from Hugging Face Hub...")
        bar = st.progress(0.0)

        # Ensure local root exists
        local_root.mkdir(parents=True, exist_ok=True)

        for i, remote_file in enumerate(remote_files, 1):
            try:
                # Download to local_root directly
                downloaded_file = hf_hub_download(
                    repo_id=repo_id,
                    filename=remote_file,
                    repo_type=repo_type,
                    local_dir=str(local_root),  # Download directly to target directory
                    local_dir_use_symlinks=False,
                    force_download=False,
                    resume_download=True,
                )

                # Update progress
                bar.progress(i / len(remote_files))

            except Exception as file_error:
                st.warning(f"Failed to download {remote_file}: {file_error}")
                continue

        # Create completion marker
        MARKER.touch()
        st.success("✅ Download completed – all files are now local.")
        return True

    except Exception as e:
        st.error(f"❌ Error in file-by-file download: {e}")
        return False


def verify_download(local_root: Path) -> dict:
    """Verify the downloaded repository structure"""
    if not local_root.exists():
        return {"status": "missing", "files": 0, "size": 0}

    files = list(local_root.rglob("*"))
    file_count = len([f for f in files if f.is_file()])
    total_size = sum(f.stat().st_size for f in files if f.is_file())

    return {
        "status": "exists",
        "files": file_count,
        "size": total_size,
        "size_mb": round(total_size / (1024 * 1024), 2)
    }


@st.cache_resource(show_spinner=False)
def ensure_repo_mirrored() -> bool:
    """
    Ensure repository is mirrored locally with better error handling
    """
    # Check if already mirrored
    if MARKER.exists():
        verification = verify_download(LOCALROOT)
        if verification["files"] > 0:
            st.success(
                f"✅ Repository already available locally ({verification['files']} files, {verification['size_mb']} MB)")
            return True
        else:
            # Marker exists but no files - remove marker and re-download
            MARKER.unlink()

    # Show current status
    verification = verify_download(LOCALROOT)
    if verification["status"] == "exists" and verification["files"] > 0:
        st.info(f"🔄 Found partial download ({verification['files']} files). Resuming...")

    # Try snapshot download first (recommended)
    st.info("🚀 Attempting full repository download...")
    if mirror_repo_fixed(REPO_ID, REPO_TYPE, LOCALROOT):
        return True

    # Fallback to file-by-file if snapshot fails
    st.warning("⚠️ Snapshot download failed. Trying file-by-file download...")
    if mirror_repo_file_by_file(REPO_ID, REPO_TYPE, LOCALROOT):
        return True

    # If both methods fail
    st.error("❌ All download methods failed. Please check your internet connection and repository access.")
    return False


def cleanup_partial_download():
    """Clean up partial downloads and reset"""
    if LOCALROOT.exists():
        import shutil
        shutil.rmtree(LOCALROOT)
    if MARKER.exists():
        MARKER.unlink()
    st.success("🗑️ Cleaned up partial downloads")

# Enhanced main function with better error handling


def main():
    """Main function with robust repository handling"""

    # Add cleanup option in sidebar for debugging
    with st.sidebar:
        st.markdown("### 🔧 Debug Options")
        if st.button("🗑️ Clean & Re-download"):
            cleanup_partial_download()
            st.rerun()

        # Show download status
        verification = verify_download(LOCALROOT)
        if verification["status"] == "exists":
            st.success(f"📁 Local files: {verification['files']}")
            st.success(f"💾 Size: {verification['size_mb']} MB")
    try:
        if ensure_repo_mirrored():
            enhanced_chat_interface()
        else:
            st.error("❌ Failed to download repository. Please try again or check your connection.")
            if st.button("🔄 Retry Download"):
                st.rerun()

    except Exception as e:
        st.error(f"❌ Critical error in main(): {e}")
        st.info("💡 Try using the 'Clean & Re-download' button in the sidebar")


def get_repo_info():
    try:
        files = list_repo_files(repo_id=REPO_ID, repo_type=REPO_TYPE)
        return {
            "total_files": len(files),
            "file_types": list(set(Path(f).suffix for f in files if Path(f).suffix)),
            "sample_files": files[:5]
        }
    except Exception as e:
        return {"error": str(e)}


if __name__ == "__main__":
    main()