src/streamlit_app.py

import ast
import streamlit as st
import requests
import json
from uuid import uuid4
from io import BytesIO
from datetime import datetime
from reportlab.lib.pagesizes import letter
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.units import inch
from reportlab.lib import colors

config : dict = {'api-url' : 'https://chat.voxio.in/agents'}
API_KEY = 'vk_JZgdwJHVmqsWeNvO9b2QvcThfUxl16moNr3wxoDyf2QER97rpi5w3XJCL0H9NRKE'

st.set_page_config(layout = 'wide')

if 'messages' not in st.session_state :
    st.session_state.messages = []
if 'session_id' not in st.session_state :
    st.session_state.session_id = None
if 'waiting_for_var' not in st.session_state :
    st.session_state.waiting_for_var = None


def extract_from_messages(messages) :
    """Scan already-stored messages to pull out actions and llm_response.
    Works because messages persist reliably even when separate state vars don't."""
    actions      = []
    llm_response = None
    summary      = None

    for msg in messages :
        if msg.get('role') != 'assistant' :
            continue
        try :
            data = ast.literal_eval(msg['content'])   # str(dict) → dict
            if not isinstance(data, dict) :
                continue
            for a in data.get('actions') or [] :
                if a and a not in actions :
                    actions.append(a)
            if data.get('llm_response') :
                llm_response = data['llm_response']
            if data.get('summary') :
                summary = data['summary']
        except (ValueError, SyntaxError) :
            pass   # plain-string content — skip

    return actions, llm_response, summary


def generate_pdf(summary: str, feedback: str) -> BytesIO:
    buf = BytesIO()
    doc = SimpleDocTemplate(buf, pagesize=letter,
                            topMargin=0.75*inch, bottomMargin=0.75*inch,
                            leftMargin=inch, rightMargin=inch)
    styles = getSampleStyleSheet()
    h1 = ParagraphStyle('H1', parent=styles['Heading1'], fontSize=20,
                         textColor=colors.HexColor('#1a56db'), spaceAfter=6)
    body = ParagraphStyle('Body', parent=styles['Normal'], fontSize=11, leading=17)
    story = []
    story.append(Paragraph("Summary", h1))
    story.append(Spacer(1, 6))
    for line in (summary or '').split('\n') :
        if line.strip() :
            story.append(Paragraph(line.strip(), body))
    story.append(Spacer(1, 20))
    story.append(Paragraph("Feedback", h1))
    story.append(Spacer(1, 6))
    for line in (feedback or '').split('\n') :
        if line.strip() :
            story.append(Paragraph(line.strip(), body))
    doc.build(story)
    buf.seek(0)
    return buf


def process_stream(response) :
    for line in response.iter_lines() :
        if line :
            try :
                data : dict = json.loads(line)

                if 'error' in data :
                    st.error(f'Server Error: {data["error"]}')
                    st.session_state.messages.append({'role' : 'error' , 'content' : data['error']})

                elif 'detail' in data :
                    st.error(f'API Error: {data["detail"]}')
                    st.session_state.messages.append({'role' : 'error' , 'content' : str(data['detail'])})

                elif data.get('waiting_for_input') :
                    req_var = data.get('required_variable')
                    st.session_state.waiting_for_var = req_var

                else :
                    content = data.get('out' , data)
                    st.session_state.messages.append({'role' : 'assistant' , 'content' : str(content)})

            except json.JSONDecodeError :
                st.error(f'Failed to decode JSON line: {line}')


# ── Derive actions / llm_response from persisted messages ─────────────────────
actions, llm_response, summary = extract_from_messages(st.session_state.messages)


with st.sidebar :

    st.title("Settings")

    if st.session_state.session_id :

        st.info(f'Active Session: `{st.session_state.session_id}`')

        if actions :
            st.divider()
            for a in actions :
                st.warning(f'⚡ Action: `{a}`')

        if st.button('End Session' , type = 'primary') :
            st.session_state.session_id     = None
            st.session_state.messages       = []
            st.session_state.waiting_for_var = None
            st.rerun()

    else :
        st.write("No active session.")

    if llm_response :
        st.divider()
        st.subheader("📄 Report Ready")
        pdf_bytes = generate_pdf(summary or '', llm_response)
        st.download_button(
            label="📥 Download PDF",
            data=pdf_bytes,
            file_name=f"report_{datetime.now().strftime('%Y%m%d_%H%M')}.pdf",
            mime="application/pdf",
        )
        with st.expander("👁️ Preview") :
            if summary :
                st.markdown("## Summary")
                st.markdown(summary)
            st.markdown("## Feedback")
            st.markdown(llm_response)


if not st.session_state.session_id :

    if st.button("🚀 Start Session" , type = "primary") :

        payload = {
            "session_id" : str(uuid4()) ,
            "api_key" : API_KEY
        }

        st.session_state.session_id = payload['session_id']

        with st.spinner('Initializing Workflow...') :

            try :

                with requests.post(f'{config["api-url"]}/start_session', json=payload, stream=True) as r:

                    if r.status_code == 200:
                        process_stream(r)
                        st.rerun()

                    else :
                        st.error(f"Error {r.status_code}: {r.text}")
                        st.session_state.session_id = None

            except Exception as e :

                st.error(f"Connection Error: {e}")
                st.session_state.session_id = None

chat_container = st.container()

with chat_container :

    for msg in st.session_state.messages :

        with st.chat_message(msg['role']) :

            if msg['role'] == 'assistant' :
                st.write(msg['content'])

            else :
                st.write(msg['content'])

if st.session_state.session_id and st.session_state.waiting_for_var :

    with st.chat_message("assistant") :

        st.write(f"I need a value for: **{st.session_state.waiting_for_var}**")

    if user_text := st.chat_input("Enter response...") :

        st.session_state.messages.append({'role' : 'user' , 'content' : user_text})

        payload = {
            'session_id' : st.session_state.session_id ,
            'user_input' : {st.session_state.waiting_for_var : user_text}
        }

        st.session_state.waiting_for_var = None

        with st.spinner('Processing...') :

            try :

                with requests.post(
                    f'{config["api-url"]}/process_input' ,
                    json = payload ,
                    stream = True
                ) as r :

                    if r.status_code == 200 :

                        process_stream(r)
                        st.rerun()

                    else :
                        st.error("Server Error processing input.")

            except Exception as e :
                st.error(f"Connection Error: {e}")