import streamlit as st
from azure.cosmos import CosmosClient, PartitionKey
import os
import pandas as pd
from streamlit.components.v1 import html

# 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")
CONTAINER_NAME = os.environ.get("COSMOS_CONTAINER_NAME")
Key = os.environ.get("Key")

def insert_record(record):
    try:
        response = container.create_item(body=record)
        return True, response
    except Exception as e:
        return False, str(e)

def call_stored_procedure(record):
    try:
        response = container.scripts.execute_stored_procedure(
            sproc="processPrompt",
            params=[record],
            partition_key=record['id']
        )
        return True, response
    except Exception as e:
        error_message = f"Error type: {type(e).__name__}\nError message: {str(e)}"
        if hasattr(e, 'sub_status'):
            error_message += f"\nSub-status: {e.sub_status}"
        if hasattr(e, 'response'):
            error_message += f"\nResponse: {e.response}"
        return False, error_message

def fetch_all_records():
    query = "SELECT * FROM c"
    items = list(container.query_items(query=query, enable_cross_partition_query=True))
    return pd.DataFrame(items)

def delete_records(ids):
    try:
        for id in ids:
            container.delete_item(item=id, partition_key=id)
        return True, f"Successfully deleted {len(ids)} records"
    except Exception as e:
        return False, f"Error deleting records: {str(e)}"

def create_html_component(records):
    html_template = """
    <style>
        .record-list {
            font-family: Arial, sans-serif;
        }
        .record-item {
            display: flex;
            align-items: center;
            margin-bottom: 10px;
            padding: 10px;
            border: 1px solid #ddd;
            border-radius: 5px;
        }
        .record-image {
            width: 50px;
            height: 50px;
            margin-right: 10px;
        }
        .record-details {
            flex-grow: 1;
        }
        .record-title {
            font-weight: bold;
        }
        .record-description {
            font-style: italic;
            color: #666;
        }
        .record-actions {
            display: flex;
            align-items: center;
        }
        .record-actions button {
            margin-left: 5px;
        }
    </style>
    <div class="record-list">
    """
    
    for record in records:
        html_template += f"""
        <div class="record-item">
            <img src="https://via.placeholder.com/50" class="record-image" alt="Record thumbnail">
            <div class="record-details">
                <div class="record-title">{record['name']}</div>
                <div class="record-description">{record['document']}</div>
            </div>
            <div class="record-actions">
                <button onclick="Streamlit.setComponentValue({{action: 'extend', id: '{record['id']}'}})">Extend</button>
                <label>
                    Public
                    <input type="checkbox" onclick="Streamlit.setComponentValue({{action: 'toggle_public', id: '{record['id']}'}})">
                </label>
                <button onclick="Streamlit.setComponentValue({{action: 'like', id: '{record['id']}'}})">👍</button>
                <button onclick="Streamlit.setComponentValue({{action: 'dislike', id: '{record['id']}'}})">👎</button>
            </div>
        </div>
        """
    
    html_template += "</div>"
    return html_template

# Streamlit app
st.title("🌟 Cosmos DB Record Management")

# Login section
if 'logged_in' not in st.session_state:
    st.session_state.logged_in = False

if not st.session_state.logged_in:
    st.subheader("🔐 Login")
    input_key = Key
    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
    client = CosmosClient(ENDPOINT, credential=st.session_state.primary_key)
    database = client.get_database_client(DATABASE_NAME)
    container = database.get_container_client(CONTAINER_NAME)

    # Fetch all records
    records_df = fetch_all_records()

    # Sidebar for data editing
    st.sidebar.title("Edit Records")
    edited_df = st.sidebar.data_editor(records_df, num_rows="dynamic", key="sidebar_editor")

    # Main Content Area
    st.subheader("📊 All Records")
    
    # Create and render the HTML component
    html_component = create_html_component(edited_df.to_dict('records'))
    component_value = html(html_component, height=400)
    
    # Handle component interactions
    if component_value:
        action = component_value.get('action')
        record_id = component_value.get('id')
        if action and record_id:
            if action == 'extend':
                st.write(f"Extending record {record_id}")
            elif action == 'toggle_public':
                st.write(f"Toggling public status for record {record_id}")
            elif action == 'like':
                st.write(f"Liked record {record_id}")
            elif action == 'dislike':
                st.write(f"Disliked record {record_id}")

    # Add delete and download buttons
    col1, col2 = st.columns(2)
    with col1:
        if st.button("🗑️ Delete Selected"):
            st.warning("Deletion functionality needs to be implemented for the new component")

    with col2:
        if st.download_button("📥 Download Data", edited_df.to_csv(index=False), "cosmos_db_data.csv", "text/csv"):
            st.success("Data downloaded successfully!")

    # Input fields for new record
    st.subheader("📝 Enter New Record Details")
    new_id = st.text_input("ID")
    new_name = st.text_input("Name")
    new_document = st.text_area("Document")
    new_evaluation_text = st.text_area("Evaluation Text")
    new_evaluation_score = st.number_input("Evaluation Score", min_value=0, max_value=100, step=1)

    col1, col2 = st.columns(2)
    
    # Insert Record button
    with col1:
        if st.button("💾 Insert Record"):
            record = {
                "id": new_id,
                "name": new_name,
                "document": new_document,
                "evaluationText": new_evaluation_text,
                "evaluationScore": new_evaluation_score
            }
            
            success, response = insert_record(record)
            if success:
                st.success("✅ Record inserted successfully!")
                st.json(response)
            else:
                st.error(f"❌ Failed to insert record: {response}")
            st.rerun()

    # Call Procedure button
    with col2:
        if st.button("🔧 Call Procedure"):
            record = {
                "id": new_id,
                "name": new_name,
                "document": new_document,
                "evaluationText": new_evaluation_text,
                "evaluationScore": new_evaluation_score
            }
            
            success, response = call_stored_procedure(record)
            if success:
                st.success("✅ Stored procedure executed successfully!")
                st.json(response)
            else:
                st.error(f"❌ Failed to execute stored procedure: {response}")

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

    # Display connection info
    st.sidebar.markdown("---")
    st.sidebar.subheader("🔗 Connection Information")
    st.sidebar.text(f"Endpoint: {ENDPOINT}")
    st.sidebar.text(f"Subscription ID: {SUBSCRIPTION_ID}")
    st.sidebar.text(f"Database: {DATABASE_NAME}")
    st.sidebar.text(f"Container: {CONTAINER_NAME}")

# Preview section
st.markdown("---")
st.markdown("### Select a record to preview.")