import streamlit as st
import requests
import json
import time

st.title("Transaction Summarizer")

# Input for selecting model
model_selection = st.radio("Select model to use:", ["Gemma", "Tinyllama"])

# Input for submitting new transactions
new_transactions_input = st.text_area("Enter your transactions (comma-separated)", key="input_area")
submit_button = st.button("Submit New Transactions", key="submit_button")

if submit_button:
    # Split transactions and strip whitespace
    new_transactions = [i.strip() for i in new_transactions_input.split(',') if i.strip()]
else:
    new_transactions = []

# Determine URL based on model selection
if model_selection == "Gemma":
    base_url = "https://api.runpod.ai/v2/cdvoxheesm31dc/"
elif model_selection == "Tinyllama":
    base_url = "https://api.runpod.ai/v2/0wnm75vx5o77s1/"

# Processing flag
processing = False

# When submit button is clicked and there are transactions to process
if submit_button and new_transactions:
    url = base_url + "run"
    
    # Retrieve API key from Streamlit secrets
    api_key = st.secrets["api_key"]
    
    headers = {
        'Content-Type': 'application/json',
        'Authorization': api_key
    }
    data = {
        'input': {
            'transaction': new_transactions
        }
    }
    
    json_data = json.dumps(data)

    # Show a spinner while waiting for the response
    with st.spinner("Processing..."):
        try:
            processing = True

            # Send POST request to start processing
            response = requests.post(url, headers=headers, data=json_data)
            response.raise_for_status()  # Raise an error for bad status codes

            # Parse response to get job ID
            result = response.json()
            job_id = result['id']
            st.write(f"New Job ID: {job_id}")

            # Keep checking status until it's no longer 'IN_QUEUE'
            status_url = f"{base_url}status/{job_id}"
            status = "IN_QUEUE"
            while status != "COMPLETED":
                status_response = requests.get(status_url, headers=headers)
                status_data = status_response.json()
                status = status_data.get('status', '')
                time.sleep(2)  # Adjust interval as needed

            # Once status changes, display final status
            st.write("Final status:", status_data)
            
        except requests.exceptions.RequestException as e:
            st.error(f"An error occurred: {e}")
        finally:
            processing = False
            st.spinner(False)  # Turn off the spinner after response is received

# Disable input elements while processing
if processing:
    st.text_area("Enter your transactions (comma-separated)", value=new_transactions_input, key="disabled_area", disabled=True)
    st.radio("Select model to use:", ["Gemma", "Tinyllama"], index=0, key="disabled_radio", disabled=True)
    st.button("Submit New Transactions", disabled=True, key="disabled_button")
else:
    st.text_area("Enter your transactions (comma-separated)", value=new_transactions_input, key="enabled_area")
    st.radio("Select model to use:", ["Gemma", "Tinyllama"], index=0, key="enabled_radio")
    st.button("Submit New Transactions", key="enabled_button")

if st.button("Reset", key="reset_button"):
    st.text_area("Enter your transactions (comma-separated)", value="", key="reset_area")