import gradio as gr
import asyncio
from langgraph.graph import StateGraph, END


# --- Dummy state and graph for demo ---
class SheamiState(dict):
    messages: list[str]
    units_processed: int
    units_total: int
    process_desc: str


async def fn_step(state: SheamiState):
    # simulate work
    await asyncio.sleep(0.5)
    state["units_processed"] += 1
    state["process_desc"] = (
        f"Processing step {state['units_processed']} of {state['units_total']}"
    )
    state["messages"].append(state["process_desc"])
    return state


def has_more(state: SheamiState) -> str:
    return "continue" if state["units_processed"] < state["units_total"] else "done"


builder = StateGraph(SheamiState)
builder.add_node("step", fn_step)
builder.add_conditional_edges("step", has_more, {"continue": "step", "done": END})
builder.set_entry_point("step")
app = builder.compile()


# --- Gradio integration ---
async def run_flow(user_input: str, progress=gr.Progress()):
    state = {
        "messages": [],
        "units_processed": 0,
        "units_total": 5,
        "process_desc": "Starting...",
    }
    async for s in app.astream(state, stream_mode="values"):
        # update progress bar
        print(s)
        if s["units_total"] > 0:
            progress((s["units_processed"], s["units_total"]))
        yield s["messages"][0] if len(s["messages"]) > 0 else "nothing to show"


with gr.Blocks() as demo:
    chatbot = gr.Textbox()
    msg = gr.Textbox(label="Message")
    run_btn = gr.Button("Run")

    run_btn.click(
        run_flow,
        inputs=[msg],
        outputs=[chatbot],
    )

demo.launch()