import streamlit as st
import pandas as pd
import requests
import json
import time
import io

st.title("Transaction Summarizer")

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

# 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", type="primary")

# File uploader for CSV files
uploaded_file = st.file_uploader("Upload a CSV file of transactions", type=["csv"])

# Session state to keep track of job ID and status
if "job_id" not in st.session_state:
    st.session_state.job_id = None
if "job_status" not in st.session_state:
    st.session_state.job_status = None

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


def process_transactions(transactions):
    url = base_url + "runsync"

    # Retrieve API key from Streamlit secrets
    api_key = st.secrets["api_key"]

    headers = {
        'Content-Type': 'application/json',
        'Authorization': api_key
    }
    data = {
        'input': {
            'transaction': transactions
        }
    }

    json_data = json.dumps(data)

    # Show a spinner while waiting for the response
    with st.spinner("Processing..."):
        try:
            # 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()
            st.session_state.job_id = result['id']
            st.write(f"New Job ID: {st.session_state.job_id}")

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

            # Once status changes, retrieve and display the result
            result_url = f"{base_url}{st.session_state.job_id}"
            result_response = requests.get(result_url, headers=headers)
            result_response.raise_for_status()
            result_data = result_response.json()
            return result_data

        except requests.exceptions.RequestException as e:
            st.error(f"An error occurred: {e}")
            return None


if submit_button and new_transactions_input:
    # Split transactions and strip whitespace
    new_transactions = [i.strip() for i in new_transactions_input.split(',') if i.strip()]
    result_data = process_transactions(new_transactions)
    if result_data:
        st.write("Transaction Summaries:")
        st.write(result_data)

if uploaded_file is not None:
    # Read the uploaded CSV file
    try:
        df = pd.read_csv(uploaded_file)
        st.write("Uploaded Data:")
        st.write(df)

        # Process the transactions in the CSV file
        transactions = df['transaction'].tolist()
        result_data = process_transactions(transactions)

        if result_data:
            # Assuming result_data is a list of summaries matching the input transactions
            df['summary'] = result_data
            st.write("Summarized Data:")
            st.write(df)

            # Prepare the summarized data for download
            csv_buffer = io.StringIO()
            df.to_csv(csv_buffer, index=False)
            csv_buffer.seek(0)

            # Download link for the summarized CSV
            st.download_button(
                label="Download Summarized CSV",
                data=csv_buffer,
                file_name="summarized_transactions.csv",
                mime="text/csv"
            )

    except Exception as e:
        st.error(f"An error occurred while processing the CSV file: {e}")

# Cancel button
if st.session_state.job_id and st.session_state.job_status == "IN_QUEUE":
    cancel_button = st.button("Cancel Request")
    if cancel_button:
        cancel_url = f"{base_url}cancel/{st.session_state.job_id}"
        try:
            cancel_response = requests.post(cancel_url, headers=headers)
            cancel_response.raise_for_status()
            cancel_result = cancel_response.json()
            st.session_state.job_status = "CANCELLED"
            st.write(f"Job {st.session_state.job_id} has been cancelled.")
        except requests.exceptions.RequestException as e:
            st.error(f"An error occurred while cancelling: {e}")

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