from statebase import StateBase
import time

sb = StateBase(api_key="your-key")

def run_long_task(session_id, task_definition):
    # 1. Rehydrate state — resume where we left off
    state = sb.sessions.get(session_id).state
    cursor = state.get("cursor", 0)
    completed = state.get("completed_items", [])

    # 2. Process remaining items
    for item in task_definition["items"][cursor:]:
        if item.id in completed:
            continue

        # 3. Checkpoint BEFORE the side effect
        sb.sessions.update_state(
            session_id=session_id,
            state={"cursor": cursor, "current_item": item.id},
            reasoning=f"Starting item {item.id}"
        )

        # 4. Do the work (may take minutes)
        result = process_item(item)

        # 5. Checkpoint AFTER success
        sb.sessions.update_state(
            session_id=session_id,
            state={
                "cursor": cursor + 1,
                "completed_items": completed + [item.id],
                f"result_{item.id}": result
            },
            reasoning=f"Completed item {item.id}"
        )
        cursor += 1

    return {"status": "complete", "results": ...}

def run_with_heartbeat(session_id, task, heartbeat_every=60):
    state = sb.sessions.get(session_id).state
    last_checkpoint = state.get("last_checkpoint", 0)

    for step in task.steps:
        step.run()

        if time.time() - last_checkpoint > heartbeat_every:
            sb.sessions.update_state(
                session_id=session_id,
                state={"progress": step.index, "last_checkpoint": time.time()},
                reasoning=f"Heartbeat checkpoint at step {step.index}"
            )
            last_checkpoint = time.time()

def send_invoice(session_id, invoice_id):
    state = sb.sessions.get(session_id).state
    sent = state.get("sent_invoices", [])

    if invoice_id in sent:
        return "already sent"  # Replay-safe

    email_service.send(invoice_id)
    sb.sessions.update_state(
        session_id=session_id,
        state={"sent_invoices": sent + [invoice_id]},
        reasoning=f"Invoice {invoice_id} sent"
    )

sb.sessions.update_state(
    session_id=session_id,
    state={
        "status": "awaiting_approval",
        "pending_action": action,
        "requested_at": time.time()
    },
    reasoning="Waiting for human approval"
)