import streamlit as st
from azure.cosmos import CosmosClient, exceptions
import os
import pandas as pd
import json
from datetime import datetime
import shutil
from github import Github
from git import Repo
import base64
from dotenv import load_dotenv

load_dotenv()

# Cosmos DB configuration
ENDPOINT = os.getenv("ENDPOINT")
PRIMARY_KEY = os.getenv("PRIMARY_KEY")

# GitHub configuration
GITHUB_TOKEN = os.getenv("GITHUB_TOKEN")

# Helper Functions

def download_github_repo(url, local_path):
    if os.path.exists(local_path):
        shutil.rmtree(local_path)
    Repo.clone_from(url, local_path)

def create_repo(github_token, repo_name):
    g = Github(github_token)
    user = g.get_user()
    return user.create_repo(repo_name)

def push_to_github(local_path, repo, github_token):
    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)
    
    local_repo.git.add(A=True)
    if local_repo.is_dirty():
        local_repo.git.commit('-m', 'Initial commit')
    
    origin.push(refspec=f"{local_repo.active_branch.name}:{local_repo.active_branch.name}")

def get_base64_download_link(file_path, file_name):
    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>'

def fetch_db_structure(client):
    structure = {}
    for database in client.list_databases():
        db_name = database['id']
        structure[db_name] = []
        db_client = client.get_database_client(db_name)
        for container in db_client.list_containers():
            structure[db_name].append(container['id'])
    return structure

def fetch_items(client, db_name, container_name):
    container = client.get_database_client(db_name).get_container_client(container_name)
    items = list(container.read_all_items())
    return items

def create_item(client, db_name, container_name):
    st.subheader("Create New Item")
    item_id = st.text_input("Item ID")
    item_data = st.text_area("Item Data (in JSON format)")
    
    if st.button("Save New Item"):
        try:
            container = client.get_database_client(db_name).get_container_client(container_name)
            json_data = json.loads(item_data)
            json_data['id'] = item_id
            container.create_item(body=json_data)
            st.success("New item created successfully!")
            st.json(json_data)
            st.experimental_rerun()
        except Exception as e:
            st.error(f"Error creating item: {str(e)}")

def edit_item(client, db_name, container_name, item_id):
    container = client.get_database_client(db_name).get_container_client(container_name)
    item = container.read_item(item=item_id, partition_key=item_id)
    
    st.subheader(f"Editing Item: {item_id}")
    edited_item = st.text_area("Edit Item Data (in JSON format)", value=json.dumps(item, indent=2))
    
    if st.button("Save Changes"):
        try:
            json_data = json.loads(edited_item)
            container.upsert_item(body=json_data)
            st.success("Item updated successfully!")
            st.json(json_data)
            st.experimental_rerun()
        except Exception as e:
            st.error(f"Error updating item: {str(e)}")

def delete_item(client, db_name, container_name, item_id):
    st.subheader(f"Delete Item: {item_id}")
    if st.button("Confirm Delete"):
        try:
            container = client.get_database_client(db_name).get_container_client(container_name)
            container.delete_item(item=item_id, partition_key=item_id)
            st.success(f"Item {item_id} deleted successfully!")
            st.experimental_rerun()
        except Exception as e:
            st.error(f"Error deleting item: {str(e)}")

def archive_all_data(client):
    try:
        base_dir = "./cosmos_archive"
        if os.path.exists(base_dir):
            shutil.rmtree(base_dir)
        os.makedirs(base_dir)

        for database in client.list_databases():
            db_name = database['id']
            db_dir = os.path.join(base_dir, db_name)
            os.makedirs(db_dir)

            db_client = client.get_database_client(db_name)
            for container in db_client.list_containers():
                container_name = container['id']
                container_dir = os.path.join(db_dir, container_name)
                os.makedirs(container_dir)

                container_client = db_client.get_container_client(container_name)
                items = list(container_client.read_all_items())

                with open(os.path.join(container_dir, f"{container_name}.json"), 'w') as f:
                    json.dump(items, f, indent=2)

        archive_name = f"cosmos_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:
        st.error(f"An error occurred while archiving data: {str(e)}")

def download_all_code():
    try:
        base_dir = "."
        exclude_dirs = ['.git', '__pycache__', 'cosmos_archive']
        exclude_files = ['.env', '.gitignore', 'cosmos_archive.zip']

        def zipdir(path, ziph):
            for root, dirs, files in os.walk(path):
                dirs[:] = [d for d in dirs if d not in exclude_dirs]
                for file in files:
                    if file not in exclude_files and not file.endswith('.zip'):
                        ziph.write(os.path.join(root, file),
                                   os.path.relpath(os.path.join(root, file),
                                                   os.path.join(base_dir, '..')))

        zip_filename = f"CosmosDBAICode_{datetime.now().strftime('%Y%m%d_%H%M%S')}.zip"
        shutil.make_archive(zip_filename[:-4], 'zip', base_dir)
        return get_base64_download_link(zip_filename, zip_filename)
    except Exception as e:
        st.error(f"An error occurred while creating the code archive: {str(e)}")

def render_sidebar(client, structure):
    st.sidebar.title("📁 Cosmos DB Manager")

    # Navigation
    app_mode = st.sidebar.radio("Navigate", ["App", "Tutorial"])

    if app_mode == "App":
        # Database and Container selection
        selected_db = st.sidebar.selectbox("Select Database", list(structure.keys()))
        selected_container = st.sidebar.selectbox("Select Container", structure[selected_db])
        
        # Fetch items for the selected container
        items = fetch_items(client, selected_db, selected_container)
        
        # Item selection
        item_ids = [item['id'] for item in items]
        selected_item = st.sidebar.selectbox("Select Item", ["Create New Item"] + item_ids)
        
        # Action buttons
        if selected_item != "Create New Item":
            if st.sidebar.button("Edit Item"):
                st.session_state.action = "edit"
                st.session_state.selected_item = selected_item
            if st.sidebar.button("Delete Item"):
                st.session_state.action = "delete"
                st.session_state.selected_item = selected_item
        else:
            if st.sidebar.button("Create New Item"):
                st.session_state.action = "create"
        
        # GitHub Operations
        st.sidebar.subheader("🐙 GitHub Operations")
        source_repo = st.sidebar.text_input("Source Repo URL")
        new_repo_name = st.sidebar.text_input("New Repo Name", value=f"Clone-{datetime.now().strftime('%Y%m%d_%H%M%S')}")
        
        if st.sidebar.button("Clone Repository"):
            clone_github_repo(source_repo, new_repo_name)
        
        if st.sidebar.button("Push to New Repository"):
            push_to_new_repo(source_repo, new_repo_name)
        
        # Archive data button
        if st.sidebar.button("Archive All Data"):
            download_link = archive_all_data(client)
            st.sidebar.markdown(download_link, unsafe_allow_html=True)
        
        # Download All Code button
        if st.sidebar.button("Download All Code"):
            download_link = download_all_code()
            st.sidebar.markdown(download_link, unsafe_allow_html=True)
        
        # Logout button
        if st.sidebar.button("Logout"):
            st.session_state.logged_in = False
            st.session_state.action = None
            st.session_state.selected_item = None
            st.experimental_rerun()
        
        return selected_db, selected_container, selected_item, app_mode
    else:
        # For tutorial mode
        return None, None, None, app_mode

def clone_github_repo(source_repo, new_repo_name):
    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"
            shutil.make_archive(zip_filename[:-4], 'zip', local_path)
            st.sidebar.markdown(get_base64_download_link(zip_filename, zip_filename), unsafe_allow_html=True)
            st.sidebar.success("Repository cloned successfully!")
        except Exception as e:
            st.sidebar.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.sidebar.error("Please ensure GitHub token is set in environment variables and source repository URL is provided.")

def push_to_new_repo(source_repo, new_repo_name):
    if GITHUB_TOKEN and source_repo:
        try:
            new_repo = create_repo(GITHUB_TOKEN, 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.sidebar.success(f"Repository pushed successfully to {new_repo.html_url}")
        except Exception as e:
            st.sidebar.error(f"An error occurred: {str(e)}")
        finally:
            if os.path.exists(local_path):
                shutil.rmtree(local_path)
    else:
        st.sidebar.error("Please ensure GitHub token is set in environment variables and source repository URL is provided.")

def main():
    st.set_page_config(layout="wide")
    st.title("🌟 Cosmos DB and GitHub Integration App with Tutorial")

    # Initialize session state
    if 'logged_in' not in st.session_state:
        st.session_state.logged_in = False
    if 'action' not in st.session_state:
        st.session_state.action = None
    if 'selected_item' not in st.session_state:
        st.session_state.selected_item = None

    # Login section
    if not st.session_state.logged_in:
        st.subheader("Login")
        if st.button("Login"):
            if PRIMARY_KEY:
                st.session_state.primary_key = PRIMARY_KEY
                st.session_state.logged_in = True
                st.experimental_rerun()
            else:
                st.error("Invalid key. Please check your environment variables.")
    else:
        # Initialize Cosmos DB client
        try:
            client = CosmosClient(ENDPOINT, credential=st.session_state.primary_key)
            # Fetch database structure
            structure = fetch_db_structure(client)
            # Render sidebar and get selections
            selected_db, selected_container, selected_item, app_mode = render_sidebar(client, structure)
            
            if app_mode == "App":
                # Main area
                if st.session_state.action == "create":
                    create_item(client, selected_db, selected_container)
                elif st.session_state.action == "edit":
                    edit_item(client, selected_db, selected_container, st.session_state.selected_item)
                elif st.session_state.action == "delete":
                    delete_item(client, selected_db, selected_container, st.session_state.selected_item)
                else:
                    st.write("Select an action from the sidebar to get started.")
            elif app_mode == "Tutorial":
                display_tutorial()
        except exceptions.CosmosHttpResponseError as e:
            st.error(f"Failed to connect to Cosmos DB: {str(e)}")
        except Exception as e:
            st.error(f"An unexpected error occurred: {str(e)}")

def display_tutorial():
    st.header("📖 In-App Tutorial")

    st.markdown("""
    ## Introduction

    In this tutorial, we'll build a web application using **Streamlit** that allows users to perform CRUD (Create, Read, Update, Delete) operations on Azure Cosmos DB. Additionally, we'll integrate GitHub functionalities to clone repositories, push changes, and archive data.

    ### What You'll Learn

    - How to connect to Azure Cosmos DB using Python.
    - Performing CRUD operations on Cosmos DB containers.
    - Integrating GitHub functionalities within the application.
    - Building an interactive user interface with Streamlit.

    ### Use Cases

    - **Database Administration**: Simplify database management tasks for administrators.
    - **Educational Tools**: Teach students about NoSQL databases and version control.
    - **Data Archiving**: Easily archive and download database contents.

    ---

    ## Prerequisites

    Before we start, ensure you have the following:

    - An **Azure Cosmos DB** account.
    - A **GitHub** account and a personal access token with the necessary permissions.
    - Python 3.7 or higher installed on your system.
    - Basic knowledge of Python and Streamlit.

    ---

    ## Setting Up the Environment

    ### 1. Clone the Repository

    Use the following command to clone the repository:

    """)
    st.code("""
    git clone https://github.com/yourusername/cosmosdb-streamlit-app.git
    cd cosmosdb-streamlit-app
    """, language='bash')

    st.markdown("""
    ### 2. Create a Virtual Environment

    Create and activate a virtual environment to manage dependencies.

    """)
    st.code("""
    python -m venv venv
    # On Windows
    venv\\Scripts\\activate
    # On Unix or MacOS
    source venv/bin/activate
    """, language='bash')

    st.markdown("""
    ### 3. Install Dependencies

    Install the required Python packages using `requirements.txt`.

    """)
    st.code("""
    pip install -r requirements.txt
    """, language='bash')

    st.markdown("""
    Contents of `requirements.txt`:

    """)
    st.code("""
    streamlit
    azure-cosmos
    PyGithub
    GitPython
    python-dotenv
    """, language='text')

    st.markdown("""
    ### 4. Set Up Environment Variables

    Create a `.env` file to store your sensitive information.

    """)
    st.code("""
    touch .env
    """, language='bash')

    st.markdown("""
    Add the following variables to your `.env` file:

    """)
    st.code("""
    ENDPOINT=https://your-cosmos-db-account.documents.azure.com:443/
    PRIMARY_KEY=your_cosmos_db_primary_key
    GITHUB_TOKEN=your_github_personal_access_token
    """, language='text')

    st.warning("**Note:** Never commit your `.env` file to version control.")

    st.markdown("---")

    st.markdown("""
    ## Application Architecture

    ### Overview

    Our application consists of several key components:

    - **Cosmos DB Operations**: Functions to connect and interact with Azure Cosmos DB.
    - **GitHub Integration**: Functions to clone repositories, create new ones, and push changes.
    - **Streamlit Interface**: The front-end of our app that users will interact with.

    ### Data Flow

    1. **User Interaction**: Users perform actions via the Streamlit sidebar and main interface.
    2. **Database Operations**: The app connects to Cosmos DB to perform CRUD operations based on user input.
    3. **GitHub Actions**: Users can clone repositories or push local changes to GitHub.
    4. **Archiving Data**: The app can archive database contents into a downloadable ZIP file.

    ---

    ## Implementing the Application

    Below are code snippets for key parts of the application.

    ### Importing Libraries

    """)
    st.code("""
    import streamlit as st
    from azure.cosmos import CosmosClient, exceptions
    import os
    import json
    from datetime import datetime
    import shutil
    from github import Github
    from git import Repo
    import base64
    from dotenv import load_dotenv

    load_dotenv()
    """, language='python')

    st.markdown("""
    ### Configuring Cosmos DB and GitHub

    """)
    st.code("""
    # Cosmos DB configuration
    ENDPOINT = os.getenv("ENDPOINT")
    PRIMARY_KEY = os.getenv("PRIMARY_KEY")

    # GitHub configuration
    GITHUB_TOKEN = os.getenv("GITHUB_TOKEN")
    """, language='python')

    st.markdown("""
    ### Helper Functions

    **GitHub Operations**
    """)
    st.code("""
    def download_github_repo(url, local_path):
        # Function implementation

    def create_repo(github_token, repo_name):
        # Function implementation

    def push_to_github(local_path, repo, github_token):
        # Function implementation
    """, language='python')

    st.markdown("""
    **Base64 Download Link**
    """)
    st.code("""
    def get_base64_download_link(file_path, file_name):
        # Function implementation
    """, language='python')

    st.markdown("""
    ### Cosmos DB Operations

    **Fetching Database Structure**
    """)
    st.code("""
    def fetch_db_structure(client):
        # Function implementation
    """, language='python')

    st.markdown("""
    **CRUD Operations**

    *Create Item*
    """)
    st.code("""
    def create_item(client, db_name, container_name):
        # Function implementation
    """, language='python')

    st.markdown("""
    *Edit Item*
    """)
    st.code("""
    def edit_item(client, db_name, container_name, item_id):
        # Function implementation
    """, language='python')

    st.markdown("""
    *Delete Item*
    """)
    st.code("""
    def delete_item(client, db_name, container_name, item_id):
        # Function implementation
    """, language='python')

    st.markdown("""
    ### Archiving Data

    """)
    st.code("""
    def archive_all_data(client):
        # Function implementation
    """, language='python')

    st.markdown("""
    ### Streamlit Interface

    **Rendering the Sidebar**
    """)
    st.code("""
    def render_sidebar(client, structure):
        # Function implementation
    """, language='python')

    st.markdown("""
    **Main Function**
    """)
    st.code("""
    def main():
        # Function implementation

    if __name__ == "__main__":
        main()
    """, language='python')

    st.markdown("---")

    st.markdown("""
    ## Running the Application

    ### Start the Streamlit App

    Run the following command in your terminal:

    """)
    st.code("""
    streamlit run app.py
    """, language='bash')

    st.markdown("""
    ### Login

    - Navigate to `http://localhost:8501` in your web browser.
    - Click on the **Login** button. If your `PRIMARY_KEY` is correctly set in the `.env` file, you should be logged in.

    ### Using the App

    - **Database Selection**: Use the sidebar to select a database and container.
    - **Item Operations**:
      - **Create New Item**: Input the item ID and JSON data to create a new item.
      - **Edit Item**: Modify the JSON data of an existing item.
      - **Delete Item**: Remove an item from the container.
    - **GitHub Operations**:
      - **Clone Repository**: Enter the source repository URL and clone it locally.
      - **Push to New Repository**: Create a new repository and push local changes.
    - **Archive Data**: Click **Archive All Data** to download a ZIP file of your databases.
    - **Download All Code**: Click **Download All Code** to download the app's source code.

    ---

    ## Summary

    In this tutorial, we've built a comprehensive application that integrates Azure Cosmos DB and GitHub functionalities using Streamlit. This app allows users to manage their Cosmos DB databases through a graphical interface, perform GitHub operations, and archive data easily.

    ### Next Steps

    - **Enhance Security**: Implement authentication mechanisms for multiple users.
    - **Extend Functionality**: Add support for other database operations like indexing.
    - **User Interface Improvements**: Make the app more user-friendly with better error handling and feedback.

    """)

if __name__ == "__main__":
    main()