# =============================================================================
# ───────────── IMPORTS ─────────────
# =============================================================================
import base64
import glob
import hashlib
import json
import os
import pandas as pd
import pytz
import random
import re
import shutil
import streamlit as st
import time
import traceback
import uuid
import zipfile
from PIL import Image
from azure.cosmos import CosmosClient, PartitionKey, exceptions
from datetime import datetime
from git import Repo
from github import Github
from gradio_client import Client
import tempfile
import io
import requests
import numpy as np
from urllib.parse import quote
import logging  # Added for logging

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# ... (Previous imports and external links remain unchanged)

# =============================================================================
# ───────────── APP CONFIGURATION ─────────────
# =============================================================================
# ... (App configuration remains unchanged)

# =============================================================================
# ───────────── HELPER FUNCTIONS ─────────────
# =============================================================================
# ... (Previous helper functions remain unchanged)

# =============================================================================
# ───────────── COSMOS DB FUNCTIONS ─────────────
# =============================================================================
def vector_search(container, query, search_type="basic", limit=10):
    """Perform vector search on Cosmos DB container with different search types."""
    try:
        if search_type == "basic":
            sql_query = f"SELECT TOP {limit} * FROM c WHERE CONTAINS(c.content, '{query}', true) ORDER BY c._ts DESC"
        elif search_type == "exact":
            sql_query = f"SELECT TOP {limit} * FROM c WHERE c.content = '{query}' ORDER BY c._ts DESC"
        elif search_type == "prefix":
            sql_query = f"SELECT TOP {limit} * FROM c WHERE STARTSWITH(c.content, '{query}', true) ORDER BY c._ts DESC"
        else:
            raise ValueError("Unsupported search type")
        
        logger.info(f"Executing vector search with query: {sql_query}")
        items = list(container.query_items(query=sql_query, enable_cross_partition_query=True))
        return items, sql_query
    except Exception as e:
        logger.error(f"Vector search error: {str(e)}")
        return [], f"Error: {str(e)}"

def delete_record(container, record):
    """Delete a record from Cosmos DB with detailed logging."""
    try:
        doc_id = record["id"]
        partition_key_value = record.get("pk", doc_id)
        logger.info(f"Attempting to delete record {doc_id} with partition key {partition_key_value}")
        container.delete_item(item=doc_id, partition_key=partition_key_value)
        logger.info(f"Successfully deleted record {doc_id}")
        return True, f"Record {doc_id} deleted. 🗑️"
    except exceptions.CosmosResourceNotFoundError as e:
        logger.warning(f"Record {doc_id} not found (already deleted): {str(e)}")
        return True, f"Record {doc_id} not found (already deleted). 🗑️"
    except exceptions.CosmosHttpResponseError as e:
        logger.error(f"HTTP error deleting {doc_id}: {str(e)}")
        return False, f"HTTP error deleting {doc_id}: {str(e)} 🚨"
    except Exception as e:
        logger.error(f"Unexpected error deleting {doc_id}: {str(e)}")
        return False, f"Unexpected error deleting {doc_id}: {str(e)} 😱"

def import_database(client, file_path):
    """Import a Cosmos DB database from a zip file."""
    try:
        with zipfile.ZipFile(file_path, 'r') as zip_ref:
            zip_ref.extractall("./temp_import")
        for db_file in glob.glob("./temp_import/*.json"):
            with open(db_file, 'r') as f:
                db_data = json.load(f)
                db_id = db_data.get("id")
                if db_id:
                    database = client.create_database_if_not_exists(id=db_id)
                    for container_data in db_data.get("containers", []):
                        container_id = container_data.get("id")
                        partition_key = container_data.get("partitionKey", "/pk")
                        container = database.create_container_if_not_exists(
                            id=container_id, partition_key=PartitionKey(path=partition_key)
                        )
                        for item in container_data.get("items", []):
                            container.upsert_item(item)
        shutil.rmtree("./temp_import")
        return True, "Database imported successfully! 📥"
    except Exception as e:
        logger.error(f"Import error: {str(e)}")
        return False, f"Import error: {str(e)} 😱"

# ... (Other Cosmos DB functions remain unchanged)

# =============================================================================
# ───────────── UI FUNCTIONS ─────────────
# =============================================================================
def vector_search_ui(container):
    st.sidebar.markdown("### 🌐 Vector Search")
    st.sidebar.markdown("*Search documents using vector-based techniques in Cosmos DB.*")
    with st.sidebar.form("vector_search_form"):
        query = st.text_input("Enter search query", key="vector_query")
        search_type = st.selectbox("Search Type", ["basic", "exact", "prefix"], help="Basic: Contains keyword, Exact: Exact match, Prefix: Starts with")
        search_submitted = st.form_submit_button("🔍 Run Search")
    
    if search_submitted and query:
        results, executed_query = vector_search(container, query, search_type)
        st.markdown("#### Vector Search Results")
        st.text(f"Executed Query: {executed_query}")
        if isinstance(results, list) and results:
            for doc in results:
                with st.expander(f"{doc.get('name', 'Unnamed')} - {doc.get('timestamp', 'No timestamp')}"):
                    st.json(doc)
        else:
            st.info("No results found or an error occurred.")

def import_export_ui(client, database_name, container_name):
    st.sidebar.markdown("### 📦 Import/Export DB")
    st.sidebar.markdown("*Easily import or export your Cosmos DB databases.*")
    col1, col2 = st.sidebar.columns(2)
    with col1:
        if st.button("📤 Export"):
            download_link = archive_current_container(database_name, container_name, client)
            st.sidebar.markdown(download_link, unsafe_allow_html=True) if download_link.startswith('<a') else st.sidebar.error(download_link)
    with col2:
        uploaded_file = st.file_uploader("📥 Import", type=["zip"], key="import_file")
        if uploaded_file and st.button("Import Now"):
            success, message = import_database(client, uploaded_file)
            if success:
                st.sidebar.success(message)
                st.rerun()
            else:
                st.sidebar.error(message)

# ... (Other UI functions remain unchanged)

# =============================================================================
# ───────────── MAIN FUNCTION ─────────────
# =============================================================================
def main():
    st.markdown("### 🐙 GitCosmos - Cosmos & Git Hub")
    st.markdown(f"[🔗 Portal]({CosmosDBUrl})")
    
    if "chat_history" not in st.session_state:
        st.session_state.chat_history = []
    if "current_container" not in st.session_state:
        st.session_state.current_container = None
    if not Key:
        st.error("Missing Cosmos Key 🔑❌")
        return
    st.session_state.primary_key = Key
    st.session_state.logged_in = True

    # Sidebar: Hierarchical Navigation
    st.sidebar.title("🐙 Navigator")
    
    # Vector Search (Moved to top)
    if st.session_state.current_container:
        vector_search_ui(st.session_state.current_container)

    # Import/Export (Second feature)
    if st.session_state.current_container:
        import_export_ui(st.session_state.client, st.session_state.selected_database, st.session_state.selected_container)

    # Databases Section
    st.sidebar.subheader("🗃️ Databases")
    if "client" not in st.session_state:
        st.session_state.client = CosmosClient(ENDPOINT, credential=Key)
    databases = get_databases(st.session_state.client)
    selected_db = st.sidebar.selectbox("Select Database", databases, key="db_select")
    if selected_db != st.session_state.get("selected_database"):
        st.session_state.selected_database = selected_db
        st.session_state.selected_container = None
        st.session_state.current_container = None
        if 'active_search' in st.session_state:
            del st.session_state.active_search
        st.rerun()

    # Containers Section
    if st.session_state.selected_database:
        database = st.session_state.client.get_database_client(st.session_state.selected_database)
        st.sidebar.subheader("📁 Containers")
        if st.sidebar.button("🆕 New Container"):
            with st.sidebar.form("new_container_form"):
                container_id = st.text_input("Container ID", "new-container")
                partition_key = st.text_input("Partition Key", "/pk")
                if st.form_submit_button("Create"):
                    container = create_new_container(database, container_id, partition_key)
                    if container:
                        st.success(f"Container '{container_id}' created!")
                        st.rerun()
        containers = get_containers(database)
        selected_container = st.sidebar.selectbox("Select Container", containers, key="container_select")
        if selected_container != st.session_state.get("selected_container"):
            st.session_state.selected_container = selected_container
            st.session_state.current_container = database.get_container_client(selected_container)
            if 'active_search' in st.session_state:
                del st.session_state.active_search
            st.rerun()

        # Actions Section
        st.sidebar.subheader("⚙️ Actions")

        # Items Section
        st.sidebar.subheader("📑 Items")
        if st.session_state.current_container:
            if st.sidebar.button("➕ New Item"):
                new_item_default(st.session_state.current_container)
            st.sidebar.text_input("New Field Key", key="new_field_key")
            st.sidebar.text_input("New Field Value", key="new_field_value")
            if st.sidebar.button("➕ Add Field") and "doc_editor" in st.session_state:
                add_field_to_doc()
            if st.sidebar.button("🤖 New AI Record"):
                new_ai_record(st.session_state.current_container)
            if st.sidebar.button("🔗 New Links Record"):
                new_links_record(st.session_state.current_container)
            search_documents_ui(st.session_state.current_container)

    # Central Area: Editable Documents with Search Filter
    if st.session_state.current_container:
        search_keyword = st.session_state.get('active_search', None)
        edit_all_documents(st.session_state.current_container, search_keyword)
    else:
        st.info("Select a database and container to view and edit documents.")

    # Additional Features
    update_file_management_section()
    add_video_generation_ui(st.session_state.current_container if st.session_state.current_container else None)

    # Logout
    if st.session_state.logged_in and st.sidebar.button("🚪 Logout"):
        st.session_state.logged_in = False
        st.session_state.client = None
        st.session_state.selected_database = None
        st.session_state.selected_container = None
        st.session_state.current_container = None
        if 'active_search' in st.session_state:
            del st.session_state.active_search
        st.rerun()

if __name__ == "__main__":
    main()