import streamlit as st
import uuid
from agent_workflow import workflow_app

st.set_page_config(page_title="Agentic AI Document Review", layout="wide")
st.title("Agentic AI Document Review System")

if "thread_id" not in st.session_state:
    st.session_state.thread_id = str(uuid.uuid4())
    st.session_state.config = {"configurable": {"thread_id": st.session_state.thread_id}}

raw_input = st.text_area("Enter initial notes/requirements:", height=150)

if st.button("Start Workflow") and raw_input:
    initial_state = {
        "raw_input": raw_input,
        "current_draft": "",
        "review_feedback": "",
        "status": "started",
        "revision_count": 0
    }
    st.session_state.workflow_started = True
    workflow_app.invoke(initial_state, st.session_state.config)
    st.rerun()

if getattr(st.session_state, "workflow_started", False):
    state = workflow_app.get_state(st.session_state.config)
    if state.values:
        st.write(f"**Status:** {state.values.get('status')}")
        st.write(f"**Revisions:** {state.values.get('revision_count')}")
        st.write(f"**Latest Feedback:** {state.values.get('review_feedback')}")
        
        st.subheader("Current Draft")
        st.markdown(state.values.get("current_draft", ""))

        next_node = state.next
        if "human_approval" in next_node:
            st.warning("Waiting for Human Approval")
            if st.button("Approve & Finalize"):
                workflow_app.invoke(None, st.session_state.config)
                st.success("Workflow Approved!")
                st.rerun()
        elif not next_node:
            st.success("Workflow Completed!")