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

st.title("Transaction Summarizer")

# custom CSS to gray out Tinyllama
st.markdown("""
    <style>
        .disabled-option {
            color: gray;
            pointer-events: none;
            cursor: not-allowed;
        }
    </style>
""", unsafe_allow_html=True)

# model selection with grayed-out option
model_selection = st.radio(
    "Select model to use:",
    ["Gemma", "Tinyllama (disabled)"],
    index=0
)

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

# Initialize session state for jobs list
if "jobs" not in st.session_state:
    st.session_state.jobs = []


# Define the transaction processing function
def process_transactions(transactions):
    base_url = get_base_url()
    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)

    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()
        job_id = result['id']

        # Add the job to session state and set initial status
        st.session_state.jobs.append({
            "id": job_id,
            "status": "IN_QUEUE",
            "transactions": transactions,
            "result": None
        })

        status_url = f"{base_url}status/{job_id}"

        # Polling the job status
        while True:
            status_response = requests.get(status_url, headers=headers)
            status_data = status_response.json()
            status = status_data.get('status', '')

            # Update job status in session state
            for job in st.session_state.jobs:
                if job["id"] == job_id:
                    job["status"] = status

            if status == "COMPLETED":
                break
            elif status == "CANCELLED":
                return None

            time.sleep(2)  # Adjust interval as needed

        # Once status changes, retrieve and return the result
        result_response = requests.get(status_url, headers=headers)
        result_response.raise_for_status()
        result_data = result_response.json().get('output')

        # Update job result in session state
        for job in st.session_state.jobs:
            if job["id"] == job_id:
                job["result"] = result_data

        return result_data

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


# Creating tabs for different pages
tab1, tab2 = st.tabs(["Submit Transactions", "Upload CSV"])

# Tab 1: Submit Transactions
with tab1:
    st.header("Submit New Transactions")

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

    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()]

        with st.spinner("Processing..."):
            # Process the transactions and display the results
            result_data = process_transactions(new_transactions)

            if result_data:
                st.write("Transaction Summaries:")
                st.write(result_data)
            else:
                st.write("The job was cancelled or encountered an error.")

# Tab 2: Upload CSV
with tab2:
    st.header("Upload a CSV File of Transactions")

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

    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()

            with st.spinner("Processing..."):
                # Process the transactions and display the results
                result_data = process_transactions(transactions)

                if result_data:
                    df['summary'] = result_data
                    st.write("Summarized Data:")
                    st.write(df)

                    # Prepare the summarized data for download
                    csv_buffer = io.BytesIO()
                    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"
                    )

                else:
                    st.write("The job was cancelled or encountered an error.")

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

    # Reset button
    if st.button("Reset All Jobs"):
        st.session_state.jobs = []