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

# Cosmos DB configuration
ENDPOINT = "https://acae-afd.documents.azure.com:443/"
SUBSCRIPTION_ID = "003fba60-5b3f-48f4-ab36-3ed11bc40816"
DATABASE_NAME = os.environ.get("COSMOS_DATABASE_NAME")
Key = os.environ.get("Key")

# GitHub configuration
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_zip_file(source_dir, output_filename):
    shutil.make_archive(output_filename, 'zip', source_dir)

def create_repo(g, repo_name):
    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)
    
    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):
    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>'

# Function to fetch all databases and containers
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

# Function to fetch items from a container
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

# Function to render sidebar
def render_sidebar(client, structure):
    st.sidebar.title("📁 Cosmos DB Manager")
    
    # 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)
    
    # CRUD 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 section
    st.sidebar.subheader("🐙 GitHub Operations")
    github_token = os.environ.get("GITHUB")
    source_repo = st.sidebar.text_input("🔗 Source Repo URL", value="https://github.com/AaronCWacker/AIExamples-8-24-Streamlit")
    new_repo_name = st.sidebar.text_input("📝 New Repo Name", value=f"AIExample-Clone-{datetime.now().strftime('%Y%m%d_%H%M%S')}")
    
    if st.sidebar.button("📥 Clone Repository"):
        clone_github_repo(github_token, source_repo, new_repo_name)
    
    if st.sidebar.button("📤 Push to New Repository"):
        push_to_new_repo(github_token, 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)
    
    # 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.rerun()
    
    return selected_db, selected_container, selected_item

# GitHub operations
def clone_github_repo(github_token, 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"
            create_zip_file(local_path, zip_filename[:-4])
            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(github_token, source_repo, new_repo_name):
    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.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.")

# Function to archive all data
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:
        return f"An error occurred while archiving data: {str(e)}"

# New function to convert text to JSON
def text_to_json(text):
    lines = text.split('\n')
    json_data = {}
    for line in lines:
        if ':' in line:
            key, value = line.split(':', 1)
            json_data[key.strip()] = value.strip()
        else:
            # If there's no colon, use the whole line as a key with an empty string as value
            json_data[line.strip()] = ""
    return json_data

# Cosmos DB operations
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 (Text format)", 
                             value="🌟 Cosmos DB and GitHub Integration\nCreate New Item\nItem ID\nItem Data (Text format)\n💾 Save New Item")
    
    if st.button("💾 Save New Item"):
        try:
            container = client.get_database_client(db_name).get_container_client(container_name)
            json_data = text_to_json(item_data)
            new_item = {"id": item_id, "content": json_data}
            container.create_item(body=new_item)
            st.success("New item created successfully!")
            st.json(new_item)  # Display the created item in JSON format
            st.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 = {}
    for key, value in item.items():
        if key not in ['_rid', '_self', '_etag', '_attachments', '_ts']:
            if key == 'content':
                # Convert the JSON content back to text for editing
                text_content = "\n".join([f"{k}: {v}" for k, v in value.items()])
                edited_text = st.text_area(key, text_content)
                edited_item[key] = text_to_json(edited_text)
            else:
                edited_item[key] = st.text_input(key, value)
    
    if st.button("💾 Save Changes"):
        try:
            container.upsert_item(body=edited_item)
            st.success("Item updated successfully!")
            st.json(edited_item)  # Display the updated item in JSON format
            st.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}")
    st.write("Are you sure you want to delete this item?")
    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.rerun()
        except Exception as e:
            st.error(f"Error deleting item: {str(e)}")

# Main Streamlit app
def main():
    st.set_page_config(layout="wide")
    st.title("🌟 Cosmos DB and GitHub Integration")

    # 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")
        input_key = Key  # Get the key from the environment variable
        if st.button("🚀 Login"):
            if input_key:
                st.session_state.primary_key = input_key
                st.session_state.logged_in = True
                st.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 = render_sidebar(client, structure)
            
            # 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.")

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

if __name__ == "__main__":
    main()