import streamlit as st
from azure.cosmos import CosmosClient
import os
import traceback
import json

# Cosmos DB configuration
ENDPOINT = "https://acae-afd.documents.azure.com:443/"
SUBSCRIPTION_ID = "003fba60-5b3f-48f4-ab36-3ed11bc40816"
# You'll need to set these environment variables or use Azure Key Vault
DATABASE_NAME = os.environ.get("COSMOS_DATABASE_NAME")
CONTAINER_NAME = os.environ.get("COSMOS_CONTAINER_NAME")
Key = os.environ.get("Key")

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

def call_stored_procedure(container, long_text, partition_key_type):
    try:
        if partition_key_type == "first_word":
            partition_key = long_text.split()[0]
        elif partition_key_type == "first_two_words":
            partition_key = " ".join(long_text.split()[:2])
        elif partition_key_type == "first_line":
            partition_key = long_text.split('\n')[0].strip()
        elif partition_key_type == "json":
            # Assume the text is a JSON string and extract a specific field
            json_data = json.loads(long_text)
            partition_key = json_data.get("partitionKey", "default")
        else:
            partition_key = "default"

        result = container.scripts.execute_stored_procedure(
            sproc='processQTPrompts',  # Make sure this matches your stored procedure name
            params=[long_text],
            partition_key=partition_key
        )
        return True, result
    except Exception as e:
        error_msg = f"Error: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
        return False, error_msg

# Streamlit app
st.title("🌟 Cosmos DB Record Insertion and Stored Procedure Caller")

# 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 = st.text_input("Enter your Cosmos DB Primary Key", type="password", value=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("Please enter a valid key")
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)

    # Create tabs for different operations
    tab1, tab2 = st.tabs(["Insert Record", "Call Stored Procedure"])

    with tab1:
        st.subheader("📝 Enter Record Details")
        id = st.text_input("ID")
        name = st.text_input("Name")
        age = st.number_input("Age", min_value=0, max_value=150)
        city = st.text_input("City")

        if st.button("💾 Insert Record"):
            record = {
                "id": id,
                "name": name,
                "age": age,
                "city": city
            }
            
            success, response = insert_record(container, record)
            if success:
                st.success("✅ Record inserted successfully!")
                st.json(response)
            else:
                st.error(f"❌ Failed to insert record: {response}")

    with tab2:
        st.subheader("📝 Enter Long Text for Stored Procedure")
        long_text = st.text_area("Long Text", height=200)

        partition_key_type = st.selectbox(
            "Select Partition Key Type",
            ["first_word", "first_two_words", "first_line", "json", "default"]
        )

        if st.button("🚀 Call Stored Procedure"):
            if long_text:
                with st.spinner('Calling stored procedure...'):
                    success, result = call_stored_procedure(container, long_text, partition_key_type)
                    if success:
                        st.success("✅ Stored procedure executed successfully!")
                        st.json(result)
                    else:
                        st.error("❌ Failed to execute stored procedure")
                        st.text(result)  # This will show the detailed error message
            else:
                st.warning("Please enter some text before calling the stored procedure.")

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

    # Display connection info
    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}")