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


# 🎉 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!

# 🐙 GitHub configuration
def download_github_repo(url, local_path):
    # 🚚 Let's download that GitHub repo!
    if os.path.exists(local_path):
        shutil.rmtree(local_path)
    Repo.clone_from(url, local_path)

def create_zip_file(source_dir, output_filename):
    # 📦 Zipping up files like a pro!
    shutil.make_archive(output_filename, 'zip', source_dir)

def create_repo(g, repo_name):
    # 🛠️ Creating a new GitHub repo. Magic!
    user = g.get_user()
    return user.create_repo(repo_name)

def push_to_github(local_path, repo, github_token):
    # 🚀 Pushing code to GitHub. Hold on tight!
    repo_url = f"https://{github_token}@github.com/{repo.full_name}.git"
    local_repo = Repo(local_path)
    
    if 'origin' in [remote.name for remote in local_repo.remotes]:
        origin = local_repo.remote('origin')
        origin.set_url(repo_url)
    else:
        origin = local_repo.create_remote('origin', repo_url)
    
    if not local_repo.heads:
        local_repo.git.checkout('-b', 'main')
        current_branch = 'main'
    else:
        current_branch = local_repo.active_branch.name
    
    local_repo.git.add(A=True)
    
    if local_repo.is_dirty():
        local_repo.git.commit('-m', 'Initial commit')
    
    origin.push(refspec=f'{current_branch}:{current_branch}')

def get_base64_download_link(file_path, file_name):
    # 🧙‍♂️ Generating a magical download link!
    with open(file_path, "rb") as file:
        contents = file.read()
    base64_encoded = base64.b64encode(contents).decode()
    return f'<a href="data:application/zip;base64,{base64_encoded}" download="{file_name}">⬇️ Download {file_name}</a>'


# 🧭 New functions for dynamic sidebar navigation
def get_databases(client):
    # 📚 Fetching list of databases. So many options!
    return [db['id'] for db in client.list_databases()]

def get_containers(database):
    # 📂 Getting containers. Containers within containers!
    return [container['id'] for container in database.list_containers()]

def get_documents(container, limit=1000):
    # 📝 Retrieving documents. Shhh, don't tell anyone!
    query = "SELECT * FROM c"
    items = list(container.query_items(query=query, enable_cross_partition_query=True, max_item_count=limit))
    return items


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

def update_record(container, updated_record):
    try:
        container.upsert_item(body=updated_record)
        return True, f"Record with id {updated_record['id']} successfully updated. 🛠️"
    except exceptions.CosmosHttpResponseError as e:
        return False, f"HTTP error occurred: {str(e)} 🚨"
    except Exception as e:
        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)
        return True, f"Successfully deleted record with name: {name} and id: {id} 🗑️"
    except exceptions.CosmosResourceNotFoundError:
        return False, f"Record with id {id} not found. It may have been already deleted. 🕵️‍♂️"
    except exceptions.CosmosHttpResponseError as e:
        return False, f"HTTP error occurred: {str(e)} 🚨"
    except Exception as e:
        return False, f"An unexpected error occurred: {traceback.format_exc()} 😱"

# 🎲 Function to generate a unique UUID
def generate_unique_id():
    # 🧙‍♂️ Generating a unique UUID!
    return str(uuid.uuid4())

# 📦 Function to archive current container
def archive_current_container(database_name, container_name, client):
    try:
        base_dir = "./cosmos_archive_current_container"
        if os.path.exists(base_dir):
            shutil.rmtree(base_dir)
        os.makedirs(base_dir)
        
        db_client = client.get_database_client(database_name)
        container_client = db_client.get_container_client(container_name)
        items = list(container_client.read_all_items())
        
        container_dir = os.path.join(base_dir, container_name)
        os.makedirs(container_dir)
        
        for item in items:
            item_id = item.get('id', f"unknown_{datetime.now().strftime('%Y%m%d%H%M%S')}")
            with open(os.path.join(container_dir, f"{item_id}.json"), 'w') as f:
                json.dump(item, f, indent=2)
        
        archive_name = f"{container_name}_archive_{datetime.now().strftime('%Y%m%d%H%M%S')}"
        shutil.make_archive(archive_name, 'zip', base_dir)
        
        return get_base64_download_link(f"{archive_name}.zip", f"{archive_name}.zip")
    except Exception as e:
        return f"An error occurred while archiving data: {str(e)} 😢"


# 🎈 Let's modify the main app to be more fun!
def main():
    st.title("🐙Git🌌Cosmos💫 - Azure Cosmos DB and Github Agent")

    # 🚦 Initialize session state
    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

    # 🔐 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  # Can't proceed without a key

    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)
            
            # 🗄️ Sidebar for database, container, and document selection
            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.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.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)
                    
                    limit_to_1000 = st.sidebar.checkbox("🔢 Limit to top 1000 documents", value=True)
                    documents = get_documents(container, limit=1000 if limit_to_1000 else None)
                    
                    if documents:
                        document_ids = [doc.get('id', 'Unknown') for doc in documents]
                        selected_document = st.sidebar.selectbox("📄 Select Document", document_ids, index=document_ids.index(st.session_state.selected_document_id) if st.session_state.selected_document_id in document_ids else 0)
                        
                        if selected_document != st.session_state.selected_document_id:
                            st.session_state.selected_document_id = selected_document
                            st.rerun()
                        
                        if st.session_state.selected_document_id:
                            st.subheader(f"📄 Document Details: {st.session_state.selected_document_id}")
                            selected_doc = next((doc for doc in documents if doc.get('id') == st.session_state.selected_document_id), None)
                            if selected_doc:
                                # 🎨 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)
    
                                if selected_view == 'Show as Markdown':
                                    # 🖌️ Show as Markdown
                                    items = get_documents(container)
                                    for item in items:
                                        st.markdown(f"### **🆔 ID: {item.get('id', 'Unknown')}**")
                                        content = item.get('content', '')
                                        if isinstance(content, dict) or isinstance(content, list):
                                            content = json.dumps(content, indent=2)
                                        st.markdown(content)
                                elif selected_view == 'Show as Code Editor':
                                    # 💻 Show as Code Editor
                                    items = get_documents(container)
                                    for item in items:
                                        st.markdown(f"### **🆔 ID: {item.get('id', 'Unknown')}**")
                                        st.code(json.dumps(item, indent=2), language='python')
                                elif selected_view == 'Show as Edit and Save':
                                    # ✏️ Show as Edit and Save
                                    st.markdown("#### Edit the document fields below:")
                                    editable_id = st.text_input("ID", value=selected_doc.get('id', ''), key='edit_id')
                                    # Remove 'id' from the document for editing other fields
                                    editable_doc = selected_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)
                                    if st.button("💾 Save Changes"):
                                        try:
                                            updated_doc = json.loads(doc_str)
                                            updated_doc['id'] = editable_id  # Include the possibly edited ID
                                            success, message = update_record(container, updated_doc)
                                            if success:
                                                st.success(message)
                                                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 == 'Clone Document':
                                    # 🧬 Clone Document
                                    if st.button("📄 Clone Document"):
                                        cloned_doc = selected_doc.copy()
                                        # Generate a unique ID
                                        cloned_doc['id'] = generate_unique_id()
                                        success, message = insert_record(container, cloned_doc)
                                        if success:
                                            st.success(f"Document cloned with new id: {cloned_doc['id']} 🎉")
                                            st.session_state.selected_document_id = cloned_doc['id']
                                            # Switch to 'Show as Edit and Save' mode
                                            st.rerun()
                                        else:
                                            st.error(message)
                                elif selected_view == 'New Record':
                                    # 🆕 New Record
                                    st.markdown("#### Create a new document:")
                                    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  # Use the provided 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']
                                                # Switch to 'Show as Edit and Save' mode
                                                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:
                    df = pd.DataFrame(documents)
                    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.rerun()

if __name__ == "__main__":
    main()