import streamlit as st
from azure.cosmos import CosmosClient, exceptions
import os
import pandas as pd
import traceback
import shutil
from github import Github
from git import Repo
from datetime import datetime
import base64
import json
import uuid  # 🎲 For generating unique IDs
from urllib.parse import quote  # 🔗 For encoding URLs
from gradio_client import Client  # 🌐 For connecting to Gradio apps

# 🎉 Welcome to our fun-filled Cosmos DB and GitHub Integration app!
st.set_page_config(layout="wide")

# 🌌 Cosmos DB configuration
ENDPOINT = "https://acae-afd.documents.azure.com:443/"
DATABASE_NAME = os.environ.get("COSMOS_DATABASE_NAME")
CONTAINER_NAME = os.environ.get("COSMOS_CONTAINER_NAME")
Key = os.environ.get("Key")  # 🔑 Don't forget your key!

# 🏠 Your local app URL (Change this to your app's URL)
LOCAL_APP_URL = "https://huggingface.co/spaces/awacke1/AzureCosmosDBUI"

# 📝 Initialize logging system
if 'logs' not in st.session_state:
    st.session_state.logs = []

def log_event(message):
    """Adds a log entry to the log session state."""
    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    st.session_state.logs.append(f"[{timestamp}] {message}")

# 🏠 Add logs to sidebar
st.sidebar.title("📝 Log Viewer")
st.sidebar.write("\n".join(st.session_state.logs))

# 🤖 OpenAI configuration
# openai.api_key = os.environ.get("OPENAI_API_KEY")
# MODEL = "gpt-3.5-turbo"  # Replace with your desired model

# 🌟 Cosmos DB functions
def insert_record(container, record):
    try:
        container.create_item(body=record)
        log_event("Record inserted successfully!")
        return True, "Record inserted successfully! 🎉"
    except exceptions.CosmosHttpResponseError as e:
        log_event(f"HTTP error occurred: {str(e)} 🚨")
        return False, f"HTTP error occurred: {str(e)} 🚨"
    except Exception as e:
        log_event(f"An unexpected error occurred: {str(e)} 😱")
        return False, f"An unexpected error occurred: {str(e)} 😱"

def update_record(container, updated_record):
    try:
        container.upsert_item(body=updated_record)
        log_event(f"Record with id {updated_record['id']} successfully updated.")
        return True, f"Record with id {updated_record['id']} successfully updated. 🛠️"
    except exceptions.CosmosHttpResponseError as e:
        log_event(f"HTTP error occurred: {str(e)} 🚨")
        return False, f"HTTP error occurred: {str(e)} 🚨"
    except Exception as e:
        log_event(f"An unexpected error occurred: {traceback.format_exc()} 😱")
        return False, f"An unexpected error occurred: {traceback.format_exc()} 😱"

def delete_record(container, name, id):
    try:
        container.delete_item(item=id, partition_key=id)
        log_event(f"Successfully deleted record with name: {name} and id: {id}")
        return True, f"Successfully deleted record with name: {name} and id: {id} 🗑️"
    except exceptions.CosmosResourceNotFoundError:
        log_event(f"Record with id {id} not found. It may have been already deleted.")
        return False, f"Record with id {id} not found. It may have been already deleted. 🕵️‍♂️"
    except exceptions.CosmosHttpResponseError as e:
        log_event(f"HTTP error occurred: {str(e)} 🚨")
        return False, f"HTTP error occurred: {str(e)} 🚨"
    except Exception as e:
        log_event(f"An unexpected error occurred: {traceback.format_exc()} 😱")
        return False, f"An unexpected error occurred: {traceback.format_exc()} 😱"

# 🎲 Function to generate a unique UUID
def generate_unique_id():
    return str(uuid.uuid4())

def get_databases(client):
    return [db['id'] for db in client.list_databases()]

def get_containers(database):
    return [container['id'] for container in database.list_containers()]

def get_documents(container, limit=None):
    query = "SELECT * FROM c ORDER BY c._ts DESC"
    items = list(container.query_items(query=query, enable_cross_partition_query=True, max_item_count=limit))
    return items

def save_to_cosmos_db(container, query, response1, response2):
    try:
        if container:
            record = {
                "id": generate_unique_id(),
                "query": query,
                "response1": response1,
                "response2": response2
            }
            try:
                container.create_item(body=record)
                st.success(f"Record saved successfully with ID: {record['id']}")
                log_event(f"Record saved to Cosmos DB with ID: {record['id']}")
                # Refresh the documents display
                st.session_state.documents = get_documents(container)
            except exceptions.CosmosHttpResponseError as e:
                st.error(f"Error saving record to Cosmos DB: {e}")
                log_event(f"Error saving record to Cosmos DB: {e}")
        else:
            st.error("Cosmos DB container is not initialized.")
            log_event("Error: Cosmos DB container is not initialized.")
    except Exception as e:
        st.error(f"An unexpected error occurred: {str(e)}")
        log_event(f"An unexpected error occurred while saving to Cosmos DB: {str(e)}")

# 🎤 Function to process text input
def process_text(text_input):
    if text_input:
        if 'messages' not in st.session_state:
            st.session_state.messages = []
        
        st.session_state.messages.append({"role": "user", "content": text_input})

        with st.chat_message("user"):
            st.markdown(text_input)
            
        with st.chat_message("assistant"):
            search_glossary(text_input)

# 🎈 Main function
def main():
    st.title("🐙Git🌌Cosmos💫 - Azure Cosmos DB and Github Agent")

    if 'logged_in' not in st.session_state:
        st.session_state.logged_in = False
    if 'selected_records' not in st.session_state:
        st.session_state.selected_records = []
    if 'client' not in st.session_state:
        st.session_state.client = None
    if 'selected_database' not in st.session_state:
        st.session_state.selected_database = None
    if 'selected_container' not in st.session_state:
        st.session_state.selected_container = None
    if 'selected_document_id' not in st.session_state:
        st.session_state.selected_document_id = None
    if 'current_index' not in st.session_state:
        st.session_state.current_index = 0
    if 'cloned_doc' not in st.session_state:
        st.session_state.cloned_doc = None

    # ⚙️ Check query parameters for any action
    try:
        query_params = st.query_params
        query = query_params.get('q') or query_params.get('query') or ''
        if query:
            process_text(query)
            st.stop()
    except Exception as e:
        st.markdown(' ')
    
    # 🔐 Automatic Login
    if Key:
        st.session_state.primary_key = Key
        st.session_state.logged_in = True
    else:
        st.error("Cosmos DB Key is not set in environment variables. 🔑❌")
        return

    if st.session_state.logged_in:
        # 🌌 Initialize Cosmos DB client
        try:
            if st.session_state.client is None:
                st.session_state.client = CosmosClient(ENDPOINT, credential=st.session_state.primary_key)
            
            st.sidebar.title("🐙Git🌌Cosmos💫🗄️Navigator")
            databases = get_databases(st.session_state.client)
            selected_db = st.sidebar.selectbox("🗃️ Select Database", databases)
            
            if selected_db != st.session_state.selected_database:
                st.session_state.selected_database = selected_db
                st.session_state.selected_container = None
                st.session_state.selected_document_id = None
                st.session_state.current_index = 0
                st.rerun()
            
            if st.session_state.selected_database:
                database = st.session_state.client.get_database_client(st.session_state.selected_database)
                containers = get_containers(database)
                selected_container = st.sidebar.selectbox("📁 Select Container", containers)
                
                if selected_container != st.session_state.selected_container:
                    st.session_state.selected_container = selected_container
                    st.session_state.selected_document_id = None
                    st.session_state.current_index = 0
                    st.rerun()
                
                if st.session_state.selected_container:
                    container = database.get_container_client(st.session_state.selected_container)
                    
                    # 📦 Add Export button
                    if st.button("📦 Export Container Data"):
                        download_link = archive_current_container(st.session_state.selected_database, st.session_state.selected_container, st.session_state.client)
                        if download_link.startswith('<a'):
                            st.markdown(download_link, unsafe_allow_html=True)
                        else:
                            st.error(download_link)
                    
                    # Fetch documents
                    documents = get_documents(container)
                    total_docs = len(documents)
                    
                    if total_docs > 5:
                        documents_to_display = documents[:5]
                        st.info("Showing top 5 most recent documents.")
                    else:
                        documents_to_display = documents
                        st.info(f"Showing all {len(documents_to_display)} documents.")
                    
                    if documents_to_display:
                        # 🎨 Add Viewer/Editor selection
                        view_options = ['Show as Markdown', 'Show as Code Editor', 'Show as Edit and Save', 'Clone Document', 'New Record']
                        selected_view = st.selectbox("Select Viewer/Editor", view_options, index=2)
                        
                        if selected_view == 'Show as Markdown':
                            doc = documents[st.session_state.current_index]
                            st.markdown(f"#### Document ID: {doc.get('id', '')}")
                            content = json.dumps(doc, indent=2)
                            st.markdown(f"```json\n{content}\n```")

                            col_prev, col_next = st.columns([1, 1])
                            with col_prev:
                                if st.button("⬅️ Previous", key='prev_markdown'):
                                    if st.session_state.current_index > 0:
                                        st.session_state.current_index -= 1
                                        st.rerun()
                            with col_next:
                                if st.button("➡️ Next", key='next_markdown'):
                                    if st.session_state.current_index < total_docs - 1:
                                        st.session_state.current_index += 1
                                        st.rerun()

                        elif selected_view == 'Show as Code Editor':
                            doc = documents[st.session_state.current_index]
                            st.markdown(f"#### Document ID: {doc.get('id', '')}")
                            doc_str = st.text_area("Edit Document", value=json.dumps(doc, indent=2), height=300, key=f'code_editor_{st.session_state.current_index}')
                            
                            col_prev, col_next = st.columns([1, 1])
                            with col_prev:
                                if st.button("⬅️ Previous", key='prev_code'):
                                    if st.session_state.current_index > 0:
                                        st.session_state.current_index -= 1
                                        st.rerun()
                            with col_next:
                                if st.button("➡️ Next", key='next_code'):
                                    if st.session_state.current_index < total_docs - 1:
                                        st.session_state.current_index += 1
                                        st.rerun()
                            
                            if st.button("💾 Save Changes", key=f'save_button_{st.session_state.current_index}'):
                                try:
                                    updated_doc = json.loads(doc_str)
                                    success, message = update_record(container, updated_doc)
                                    if success:
                                        st.success(f"Document {updated_doc['id']} saved successfully.")
                                        st.session_state.selected_document_id = updated_doc['id']
                                        st.rerun()
                                    else:
                                        st.error(message)
                                except json.JSONDecodeError as e:
                                    st.error(f"Invalid JSON: {str(e)} 🚫")
                                    
                        elif selected_view == 'Show as Edit and Save':
                            st.markdown("#### Edit the document fields below:")
                            num_cols = len(documents_to_display)
                            cols = st.columns(num_cols)

                            for idx, (col, doc) in enumerate(zip(cols, documents_to_display)):
                                with col:
                                    st.markdown(f"##### Document ID: {doc.get('id', '')}")
                                    editable_id = st.text_input("ID", value=doc.get('id', ''), key=f'edit_id_{idx}')
                                    editable_doc = doc.copy()
                                    editable_doc.pop('id', None)
                                    doc_str = st.text_area("Document Content (in JSON format)", value=json.dumps(editable_doc, indent=2), height=300, key=f'doc_str_{idx}')
                                    
                                    col_save, col_ai = st.columns(2)
                                    with col_save:
                                        if st.button("💾 Save Changes", key=f'save_button_{idx}'):
                                            try:
                                                updated_doc = json.loads(doc_str)
                                                updated_doc['id'] = editable_id
                                                success, message = update_record(container, updated_doc)
                                                if success:
                                                    st.success(f"Document {updated_doc['id']} saved successfully.")
                                                    st.session_state.selected_document_id = updated_doc['id']
                                                    st.rerun()
                                                else:
                                                    st.error(message)
                                            except json.JSONDecodeError as e:
                                                st.error(f"Invalid JSON: {str(e)} 🚫")
                                    with col_ai:
                                        if st.button("🤖 Run With AI", key=f'run_with_ai_button_{idx}'):
                                            search_glossary(json.dumps(editable_doc, indent=2))

                        elif selected_view == 'Clone Document':
                            st.markdown("#### Clone a document:")
                            for idx, doc in enumerate(documents_to_display):
                                st.markdown(f"##### Document ID: {doc.get('id', '')}")
                                if st.button("📄 Clone Document", key=f'clone_button_{idx}'):
                                    cloned_doc = doc.copy()
                                    cloned_doc['id'] = generate_unique_id()
                                    st.session_state.cloned_doc = cloned_doc
                                    st.session_state.cloned_doc_str = json.dumps(cloned_doc, indent=2)
                                    st.session_state.clone_mode = True
                                    st.rerun()

                            if st.session_state.get('clone_mode', False):
                                st.markdown("#### Edit Cloned Document:")
                                cloned_doc_str = st.text_area("Cloned Document Content (in JSON format)", value=st.session_state.cloned_doc_str, height=300)
                                if st.button("💾 Save Cloned Document"):
                                    try:
                                        new_doc = json.loads(cloned_doc_str)
                                        success, message = insert_record(container, new_doc)
                                        if success:
                                            st.success(f"Cloned document saved with id: {new_doc['id']} 🎉")
                                            st.session_state.selected_document_id = new_doc['id']
                                            st.session_state.clone_mode = False
                                            st.session_state.cloned_doc = None
                                            st.session_state.cloned_doc_str = ''
                                            st.rerun()
                                        else:
                                            st.error(message)
                                    except json.JSONDecodeError as e:
                                        st.error(f"Invalid JSON: {str(e)} 🚫")

                        elif selected_view == 'New Record':
                            st.markdown("#### Create a new document:")
                            if st.button("🤖 Insert Auto-Generated Record"):
                                success, message = insert_auto_generated_record(container)
                                if success:
                                    st.success(message)
                                    st.rerun()
                                else:
                                    st.error(message)
                            else:
                                new_id = st.text_input("ID", value=generate_unique_id(), key='new_id')
                                new_doc_str = st.text_area("Document Content (in JSON format)", value='{}', height=300)
                                if st.button("➕ Create New Document"):
                                    try:
                                        new_doc = json.loads(new_doc_str)
                                        new_doc['id'] = new_id
                                        success, message = insert_record(container, new_doc)
                                        if success:
                                            st.success(f"New document created with id: {new_doc['id']} 🎉")
                                            st.session_state.selected_document_id = new_doc['id']
                                            st.rerun()
                                        else:
                                            st.error(message)
                                    except json.JSONDecodeError as e:
                                        st.error(f"Invalid JSON: {str(e)} 🚫")

                    else:
                        st.sidebar.info("No documents found in this container. 📭")
            
            # 🎉 Main content area
            st.subheader(f"📊 Container: {st.session_state.selected_container}")
            if st.session_state.selected_container:
                if documents_to_display:
                    df = pd.DataFrame(documents_to_display)
                    st.dataframe(df)
                else:
                    st.info("No documents to display. 🧐")

            # 🐙 GitHub section
            st.subheader("🐙 GitHub Operations")
            github_token = os.environ.get("GITHUB")  # Read GitHub token from environment variable
            source_repo = st.text_input("Source GitHub Repository URL", value="https://github.com/AaronCWacker/AIExamples-8-24-Streamlit")
            new_repo_name = st.text_input("New Repository Name (for cloning)", value=f"AIExample-Clone-{datetime.now().strftime('%Y%m%d_%H%M%S')}")

            col1, col2 = st.columns(2)
            with col1:
                if st.button("📥 Clone Repository"):
                    if github_token and source_repo:
                        try:
                            local_path = f"./temp_repo_{datetime.now().strftime('%Y%m%d%H%M%S')}"
                            download_github_repo(source_repo, local_path)
                            zip_filename = f"{new_repo_name}.zip"
                            create_zip_file(local_path, zip_filename[:-4])
                            st.markdown(get_base64_download_link(zip_filename, zip_filename), unsafe_allow_html=True)
                            st.success("Repository cloned successfully! 🎉")
                        except Exception as e:
                            st.error(f"An error occurred: {str(e)} 😢")
                        finally:
                            if os.path.exists(local_path):
                                shutil.rmtree(local_path)
                            if os.path.exists(zip_filename):
                                os.remove(zip_filename)
                    else:
                        st.error("Please ensure GitHub token is set in environment variables and source repository URL is provided. 🔑❓")
        
            with col2:
                if st.button("📤 Push to New Repository"):
                    if github_token and source_repo:
                        try:
                            g = Github(github_token)
                            new_repo = create_repo(g, new_repo_name)
                            local_path = f"./temp_repo_{datetime.now().strftime('%Y%m%d%H%M%S')}"
                            download_github_repo(source_repo, local_path)
                            push_to_github(local_path, new_repo, github_token)
                            st.success(f"Repository pushed successfully to {new_repo.html_url} 🚀")
                        except Exception as e:
                            st.error(f"An error occurred: {str(e)} 😢")
                        finally:
                            if os.path.exists(local_path):
                                shutil.rmtree(local_path)
                    else:
                        st.error("Please ensure GitHub token is set in environment variables and source repository URL is provided. 🔑❓")

        except exceptions.CosmosHttpResponseError as e:
            st.error(f"Failed to connect to Cosmos DB. HTTP error: {str(e)} 🚨")
        except Exception as e:
            st.error(f"An unexpected error occurred: {str(e)} 😱")

    # 🚪 Logout button
    if st.session_state.logged_in and st.sidebar.button("🚪 Logout"):
        st.session_state.logged_in = False
        st.session_state.selected_records.clear()
        st.session_state.client = None
        st.session_state.selected_database = None
        st.session_state.selected_container = None
        st.session_state.selected_document_id = None
        st.session_state.current_index = 0
        st.rerun()

if __name__ == "__main__":
    main()