import streamlit as st
import os
import pandas as pd
from azure.identity import AzureCliCredential, DefaultAzureCredential
from azure.mgmt.resource import ResourceManagementClient
from azure.cosmos import CosmosClient, exceptions
from github import Github
from git import Repo
import shutil
from datetime import datetime
import base64
import subprocess
import sys

# Check if playwright is installed
try:
    from playwright.sync_api import sync_playwright
    PLAYWRIGHT_INSTALLED = True
except ImportError:
    PLAYWRIGHT_INSTALLED = False

# Azure 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")

# Playwright function for Azure login
def azure_login_with_playwright(username, password):
    if not PLAYWRIGHT_INSTALLED:
        st.error("Playwright is not installed. Please install it using the provided button.")
        return False

    try:
        with sync_playwright() as p:
            browser = p.chromium.launch(headless=False)  # Set to True for headless mode
            page = browser.new_page()
            page.goto("https://portal.azure.com")
            
            # Wait for and fill in the email field
            page.wait_for_selector('input[type="email"]')
            page.fill('input[type="email"]', username)
            page.click('input[type="submit"]')
            
            # Wait for and fill in the password field
            page.wait_for_selector('input[type="password"]')
            page.fill('input[type="password"]', password)
            page.click('input[type="submit"]')
            
            # Wait for login to complete (adjust as needed)
            page.wait_for_selector('button[id="nav_menu"]', timeout=60000)
            
            browser.close()
        
        return True
    except Exception as e:
        st.error(f"An error occurred during browser simulation: {str(e)}")
        return False

# Function to install Playwright and its dependencies
def install_playwright_full():
    try:
        # Install Playwright package
        subprocess.run([sys.executable, "-m", "pip", "install", "playwright"], 
                       capture_output=True, text=True, check=True)
        
        # Install Playwright browsers and dependencies
        subprocess.run([sys.executable, "-m", "playwright", "install"], 
                       capture_output=True, text=True, check=True)
        
        st.success("Playwright and its dependencies installed successfully!")
        return True
    except subprocess.CalledProcessError as e:
        st.error(f"Failed to install Playwright or its dependencies. Error: {e.stderr}")
        return False

# 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 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>'

# Azure Resource Management functions
def get_azure_resources(credential, subscription_id):
    resource_client = ResourceManagementClient(credential, subscription_id)
    resources = resource_client.resources.list()
    return [{"name": r.name, "type": r.type, "id": r.id} for r in resources]

def build_resource_tree(resources):
    tree = {}
    for resource in resources:
        resource_type = resource['type']
        if resource_type not in tree:
            tree[resource_type] = []
        tree[resource_type].append(resource)
    return tree

# Cosmos DB functions
def insert_record(container, record):
    try:
        response = container.create_item(body=record)
        return True, response
    except exceptions.CosmosHttpResponseError as e:
        return False, f"HTTP error occurred: {str(e)}. Status code: {e.status_code}"
    except Exception as e:
        return False, f"An unexpected error occurred: {str(e)}"

def fetch_all_records(container):
    try:
        query = "SELECT * FROM c"
        items = list(container.query_items(query=query, enable_cross_partition_query=True))
        return pd.DataFrame(items)
    except exceptions.CosmosHttpResponseError as e:
        st.error(f"HTTP error occurred while fetching records: {str(e)}. Status code: {e.status_code}")
        return pd.DataFrame()
    except Exception as e:
        st.error(f"An unexpected error occurred while fetching records: {str(e)}")
        return pd.DataFrame()

# Streamlit app
st.set_page_config(layout="wide")
st.title("🌟 Azure Resource Manager and GitHub Integration")

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

# Login section
if not st.session_state.logged_in:
    st.subheader("🔐 Login")
    st.write("Please enter your Azure Subscription ID. You can find this in the Azure Portal under Subscriptions.")
    subscription_id = st.text_input("Azure Subscription ID", value="003fba60-5b3f-48f4-ab36-3ed11bc40816")
    
    login_method = st.radio("Choose login method", ["Automatic", "Manual (Browser Simulation)"])
    
    if login_method == "Manual (Browser Simulation)":
        if not PLAYWRIGHT_INSTALLED:
            st.warning("Playwright is not installed. Please install it to use browser simulation.")
            if st.button("Install Playwright and Dependencies"):
                if install_playwright_full():
                    st.success("Playwright and dependencies installed successfully. Please restart the app.")
                    st.stop()
        else:
            username = st.text_input("Azure Username/Email")
            password = st.text_input("Azure Password", type="password")
    
    if st.button("🚀 Login"):
        if login_method == "Automatic":
            try:
                # First, try AzureCliCredential
                credential = AzureCliCredential()
                credential.get_token("https://management.azure.com/.default")
            except Exception:
                # If AzureCliCredential fails, fall back to DefaultAzureCredential
                try:
                    credential = DefaultAzureCredential()
                    credential.get_token("https://management.azure.com/.default")
                except Exception as e:
                    st.error(f"Failed to authenticate automatically. Please try the Manual (Browser Simulation) method. Error: {str(e)}")
                    st.stop()
        else:
            if not PLAYWRIGHT_INSTALLED:
                st.error("Playwright is not installed. Please install it to use browser simulation.")
                st.stop()
            if not username or not password:
                st.error("Please enter both username and password for manual login.")
                st.stop()
            try:
                if azure_login_with_playwright(username, password):
                    credential = DefaultAzureCredential()
                    credential.get_token("https://management.azure.com/.default")
                else:
                    st.error("Failed to authenticate using browser simulation.")
                    st.stop()
            except Exception as e:
                st.error(f"Failed to authenticate using browser simulation. Error: {str(e)}")
                st.stop()
        
        st.session_state.credential = credential
        st.session_state.subscription_id = subscription_id
        st.session_state.logged_in = True
        st.success("Successfully logged in!")
        st.rerun()
else:
    # Main app content
    col1, col2 = st.columns([1, 3])

    with col1:
        st.subheader("📊 Azure Resources")
        resources = get_azure_resources(st.session_state.credential, st.session_state.subscription_id)
        resource_tree = build_resource_tree(resources)

        for resource_type, resource_list in resource_tree.items():
            with st.expander(resource_type):
                for resource in resource_list:
                    st.write(f"- {resource['name']}")

    with col2:
        st.subheader("🔧 Resource Management")

        # Display resources in a table
        df_resources = pd.DataFrame(resources)
        st.dataframe(df_resources)

        # Save resource to Cosmos DB
        st.subheader("💾 Save Resource to Cosmos DB")
        selected_resource = st.selectbox("Select a resource to save", df_resources['name'])
        if selected_resource:
            resource = df_resources[df_resources['name'] == selected_resource].iloc[0]
            if st.button("Save to Cosmos DB"):
                try:
                    cosmos_client = CosmosClient(ENDPOINT, credential=st.session_state.credential)
                    database = cosmos_client.get_database_client(DATABASE_NAME)
                    container = database.get_container_client(CONTAINER_NAME)
                    
                    record = {
                        "id": resource['id'],
                        "name": resource['name'],
                        "type": resource['type']
                    }
                    
                    success, response = insert_record(container, record)
                    if success:
                        st.success("Resource saved to Cosmos DB successfully!")
                    else:
                        st.error(f"Failed to save resource: {response}")
                except Exception as e:
                    st.error(f"An error occurred: {str(e)}")

        # GitHub Integration
        st.subheader("🐙 GitHub Integration")
        github_token = os.environ.get("GITHUB")
        if github_token:
            source_repo = st.text_input("Source GitHub Repository URL")
            new_repo_name = st.text_input("New Repository Name (for cloning)", value=f"Azure-Resource-Clone-{datetime.now().strftime('%Y%m%d_%H%M%S')}")
            
            if st.button("📥 Clone Repository"):
                if 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 provide a source repository URL.")
        else:
            st.warning("GitHub token not found in environment variables. GitHub integration is disabled.")

        # Logout button
        if st.button("🚪 Logout"):
            st.session_state.logged_in = False
            st.session_state.credential = None
            st.session_state.subscription_id = None
            st.rerun()

    # Display connection info
    st.sidebar.subheader("🔗 Connection Information")
    st.sidebar.text(f"Subscription ID: {st.session_state.subscription_id}")
    st.sidebar.text(f"Cosmos DB Endpoint: {ENDPOINT}")
    st.sidebar.text(f"Database: {DATABASE_NAME}")
    st.sidebar.text(f"Container: {CONTAINER_NAME}")