app.py

# app.py
import os
import json
import tempfile
from typing import Dict, Any, Optional, Tuple

import gradio as gr

from replayproof_sim import (
    SimConfig,
    SimState,
    reset_sim,
    step_sim,
    render_world_image,
    render_pov_image,
)
from replayproof_receipts import (
    Recorder,
    build_receipt_bundle_zip,
    verify_receipt_bundle_zip,
    replay_from_receipt_bundle_zip,
)

APP_TITLE = "ReplayProof — Agent POV + Verified Replay"
APP_SUB = "Play it. Prove it. Replay it."


def _new_session(
    seed: int,
    size: int,
    walls_pct: float,
    coins: int,
    hazards: int,
    pov_radius: int,
) -> Dict[str, Any]:
    cfg = SimConfig(
        size=int(size),
        walls_pct=float(walls_pct),
        coins=int(coins),
        hazards=int(hazards),
        pov_radius=int(pov_radius),
        max_steps=2000,
    )
    state = reset_sim(cfg, seed=int(seed))
    rec = Recorder.new_session(
        sim_spec="replayproof-sim-v0",
        cfg=cfg.to_dict(),
        seed=int(seed),
    )
    rec.record_event(state=state, action="RESET")
    return {"cfg": cfg, "state": state, "rec": rec}


def _render(state: SimState) -> Tuple[Any, Any, str]:
    world = render_world_image(state)
    pov = render_pov_image(state)
    status = (
        f"step={state.step}  score={state.score}  done={state.done}  "
        f"final_hash={(state.last_state_sha256 or '')[:12]}"
    )
    return world, pov, status


def _receipt_tail(rec: Recorder, n: int = 25) -> str:
    tail = rec.events[-n:]
    lines = []
    for e in tail:
        lines.append(
            f"{e['step']:04d}  a={e['action']:<6}  "
            f"state={e['state_sha256'][:12]}  event={e['event_sha256'][:12]}"
        )
    return "\n".join(lines) if lines else ""


def ui_new(seed, size, walls_pct, coins, hazards, pov_radius):
    sess = _new_session(seed, size, walls_pct, coins, hazards, pov_radius)
    world, pov, status = _render(sess["state"])
    receipts = _receipt_tail(sess["rec"])
    return sess, world, pov, status, receipts, None


def ui_step(sess: Optional[Dict[str, Any]]):
    if not sess:
        return sess, None, None, "No session. Click New Session.", ""

    cfg: SimConfig = sess["cfg"]
    state: SimState = sess["state"]
    rec: Recorder = sess["rec"]

    if state.done:
        world, pov, status = _render(state)
        return sess, world, pov, status + "  (done)", _receipt_tail(rec)

    new_state, action = step_sim(cfg, state)
    sess["state"] = new_state
    rec.record_event(state=new_state, action=action)

    world, pov, status = _render(new_state)
    return sess, world, pov, status, _receipt_tail(rec)


def ui_run(sess: Optional[Dict[str, Any]], steps: int):
    if not sess:
        return sess, None, None, "No session. Click New Session.", ""

    cfg: SimConfig = sess["cfg"]
    state: SimState = sess["state"]
    rec: Recorder = sess["rec"]

    n = int(steps)
    for _ in range(n):
        if state.done:
            break
        state, action = step_sim(cfg, state)
        rec.record_event(state=state, action=action)

    sess["state"] = state
    world, pov, status = _render(state)
    return sess, world, pov, status, _receipt_tail(rec)


def ui_export_bundle(sess: Optional[Dict[str, Any]], include_gif: bool, watermark: bool):
    if not sess:
        return None, "No session. Create a session first."

    cfg: SimConfig = sess["cfg"]
    state: SimState = sess["state"]
    rec: Recorder = sess["rec"]

    out_dir = tempfile.mkdtemp(prefix="replayproof_")
    zip_path = os.path.join(out_dir, f"replayproof_bundle_{rec.session_id}.zip")

    build_receipt_bundle_zip(
        zip_path=zip_path,
        recorder=rec,
        cfg=cfg,
        final_state=state,
        include_gif=bool(include_gif),
        watermark=bool(watermark),
    )

    msg = f"Bundle created: {os.path.basename(zip_path)}"
    return zip_path, msg


def ui_verify(zip_file):
    if zip_file is None:
        return "Upload a receipt bundle ZIP."

    zip_path = zip_file if isinstance(zip_file, str) else zip_file.name
    report = verify_receipt_bundle_zip(zip_path)
    return json.dumps(report, indent=2)


def ui_replay(zip_file, export_gif_flag: bool, watermark: bool):
    if zip_file is None:
        return None, None, "Upload a receipt bundle ZIP."

    zip_path = zip_file if isinstance(zip_file, str) else zip_file.name
    replay = replay_from_receipt_bundle_zip(
        zip_path,
        export_gif_flag=bool(export_gif_flag),
        watermark=bool(watermark),
    )

    world = replay.get("world_img")
    pov = replay.get("pov_img")
    msg = replay.get("message", "")
    gif_path = replay.get("gif_path")

    return world, pov, (gif_path if gif_path else None), msg


with gr.Blocks(title=APP_TITLE) as demo:
    gr.Markdown(
        f"# {APP_TITLE}\n"
        f"**{APP_SUB}**\n\n"
        "Interactive agent POV sandbox with signed, hash-chained receipts. "
        "Export a clip + receipt bundle, then verify and replay anywhere."
    )

    with gr.Row():
        seed = gr.Number(value=12345, label="Seed", precision=0)
        size = gr.Slider(8, 20, value=12, step=1, label="Grid Size")
        walls_pct = gr.Slider(0.0, 0.35, value=0.18, step=0.01, label="Walls %")
        coins = gr.Slider(0, 12, value=5, step=1, label="Coins")
        hazards = gr.Slider(0, 10, value=4, step=1, label="Hazards")
        pov_radius = gr.Slider(2, 6, value=4, step=1, label="POV Radius")

    sess_state = gr.State(None)

    with gr.Row():
        btn_new = gr.Button("New Session", variant="primary")
        btn_step = gr.Button("Step")
        btn_run_50 = gr.Button("Run 50")
        btn_run_250 = gr.Button("Run 250")

    with gr.Row():
        world_img = gr.Image(label="World View", type="pil")
        pov_img = gr.Image(label="Agent POV", type="pil")

    status = gr.Textbox(label="Status", value="", interactive=False)

    with gr.Tabs():
        with gr.Tab("Receipts"):
            receipts_box = gr.Textbox(label="Latest receipts (tail)", lines=14, interactive=False)
            with gr.Row():
                include_gif = gr.Checkbox(value=True, label="Include GIF in bundle")
                watermark = gr.Checkbox(value=True, label="Watermark final hash on frames")
            btn_export = gr.Button("Download Receipt Bundle (ZIP)", variant="primary")
            bundle_file = gr.File(label="Bundle ZIP")
            export_msg = gr.Textbox(label="Export message", interactive=False)

        with gr.Tab("Verify"):
            verify_upload = gr.File(label="Upload receipt bundle ZIP")
            btn_verify = gr.Button("Verify Bundle", variant="primary")
            verify_report = gr.Textbox(label="Verification report (JSON)", lines=18)

        with gr.Tab("Replay"):
            replay_upload = gr.File(label="Upload receipt bundle ZIP")
            with gr.Row():
                replay_export_gif_flag = gr.Checkbox(value=True, label="Export replay GIF")
                replay_watermark = gr.Checkbox(value=True, label="Watermark final hash on frames")
            btn_replay = gr.Button("Replay From Bundle", variant="primary")
            with gr.Row():
                replay_world = gr.Image(label="Replayed World", type="pil")
                replay_pov = gr.Image(label="Replayed POV", type="pil")
            replay_gif = gr.File(label="Replayed GIF (if exported)")
            replay_msg = gr.Textbox(label="Replay result", interactive=False)

    btn_new.click(
        ui_new,
        inputs=[seed, size, walls_pct, coins, hazards, pov_radius],
        outputs=[sess_state, world_img, pov_img, status, receipts_box, bundle_file],
    )

    btn_step.click(
        ui_step,
        inputs=[sess_state],
        outputs=[sess_state, world_img, pov_img, status, receipts_box],
    )

    btn_run_50.click(
        lambda s: ui_run(s, 50),
        inputs=[sess_state],
        outputs=[sess_state, world_img, pov_img, status, receipts_box],
    )

    btn_run_250.click(
        lambda s: ui_run(s, 250),
        inputs=[sess_state],
        outputs=[sess_state, world_img, pov_img, status, receipts_box],
    )

    btn_export.click(
        ui_export_bundle,
        inputs=[sess_state, include_gif, watermark],
        outputs=[bundle_file, export_msg],
    )

    btn_verify.click(
        ui_verify,
        inputs=[verify_upload],
        outputs=[verify_report],
    )

    btn_replay.click(
        ui_replay,
        inputs=[replay_upload, replay_export_gif_flag, replay_watermark],
        outputs=[replay_world, replay_pov, replay_gif, replay_msg],
    )

if __name__ == "__main__":
    demo.launch()