import json
import os
import tempfile
from typing import Any, Dict, Optional, Tuple

import gradio as gr

from drp.bundle import verify_bundle
from drp.diff import diff_bundles
from drp.report import render_report_markdown
from drp.simulate import make_demo_bundle_zip, fork_patch_bundle


APP_TITLE = "TimelineDiff — Differential Reproducibility Protocol (DRP)"


def _tmp_path(name: str) -> str:
    d = tempfile.mkdtemp(prefix="drp_")
    return os.path.join(d, name)


def ui_verify(bundle_file) -> Tuple[str, Dict[str, Any]]:
    if not bundle_file:
        return "Upload a bundle zip first.", {"ok": False}

    ok, summary = verify_bundle(bundle_file.name)
    msg = "✅ Bundle verifies (hash chain OK)." if ok else "❌ Bundle failed verification."
    return msg, summary


def ui_diff(bundle_a, bundle_b) -> Tuple[str, Dict[str, Any], str]:
    if not bundle_a or not bundle_b:
        return "Upload two bundles.", {"error": "missing input"}, ""

    diff = diff_bundles(bundle_a.name, bundle_b.name)
    md = render_report_markdown(diff)

    # export report artifacts
    out_zip = _tmp_path("drp-diff-report.zip")
    out_md = _tmp_path("report.md")
    out_json = _tmp_path("diff.json")

    with open(out_md, "w", encoding="utf-8") as f:
        f.write(md)
    with open(out_json, "w", encoding="utf-8") as f:
        json.dump(diff, f, ensure_ascii=False, indent=2)

    import zipfile
    with zipfile.ZipFile(out_zip, "w", compression=zipfile.ZIP_DEFLATED) as z:
        z.write(out_md, arcname="report.md")
        z.write(out_json, arcname="diff.json")

    headline = "✅ Diff complete."
    return headline, diff, out_zip


def ui_generate(seed_a: int, seed_b: int, chaos: float) -> Tuple[str, str, str]:
    chaos = float(max(0.0, min(1.0, chaos)))

    path_a = _tmp_path("demo-A.zip")
    path_b = _tmp_path("demo-B.zip")
    make_demo_bundle_zip(path_a, seed=int(seed_a), chaos=chaos, label="A")
    make_demo_bundle_zip(path_b, seed=int(seed_b), chaos=chaos, label="B")

    return "✅ Generated demo bundles.", path_a, path_b


def ui_fork(source_bundle, fork_index: int, patch_kind: str, patch_step: str, patch_payload: str) -> Tuple[str, str]:
    if not source_bundle:
        return "Upload a source bundle first.", ""

    patch_payload_json: Optional[Dict[str, Any]] = None
    if patch_payload.strip():
        try:
            patch_payload_json = json.loads(patch_payload)
        except Exception as e:
            return f"Invalid JSON patch payload: {e}", ""

    out = _tmp_path("forked.zip")
    fork_patch_bundle(
        out,
        source_zip=source_bundle.name,
        fork_at_index=int(fork_index),
        patch_kind=patch_kind.strip() or None,
        patch_step=patch_step.strip() or None,
        patch_payload_json=patch_payload_json,
    )
    return "✅ Fork bundle created (patched + re-hash-chained).", out


with gr.Blocks(title=APP_TITLE) as demo:
    gr.Markdown(f"# {APP_TITLE}\nDiff two agent timelines. Find first divergence. Export a forensic report.")

    with gr.Tab("Diff two bundles"):
        with gr.Row():
            bundle_a = gr.File(label="Bundle A (.zip)", file_types=[".zip"])
            bundle_b = gr.File(label="Bundle B (.zip)", file_types=[".zip"])
        run = gr.Button("Compute differential report", variant="primary")

        status = gr.Textbox(label="Status", interactive=False)
        diff_json = gr.JSON(label="Diff JSON (summary + per-event diffs)")
        report_zip = gr.File(label="Download drp-diff-report.zip")

        run.click(fn=ui_diff, inputs=[bundle_a, bundle_b], outputs=[status, diff_json, report_zip])

    with gr.Tab("Verify bundle integrity"):
        bundle_v = gr.File(label="Bundle (.zip)", file_types=[".zip"])
        verify_btn = gr.Button("Verify hash chain", variant="secondary")
        v_msg = gr.Textbox(label="Result", interactive=False)
        v_json = gr.JSON(label="Verification summary")
        verify_btn.click(fn=ui_verify, inputs=[bundle_v], outputs=[v_msg, v_json])

    with gr.Tab("Generate demo bundles"):
        gr.Markdown("Use this to test the diff UI without integrating exporters yet.")
        with gr.Row():
            seed_a = gr.Number(value=1, label="Seed A", precision=0)
            seed_b = gr.Number(value=2, label="Seed B", precision=0)
            chaos = gr.Slider(0, 1, value=0.35, step=0.01, label="Chaos (divergence likelihood)")
        gen_btn = gr.Button("Generate", variant="primary")
        gen_msg = gr.Textbox(label="Status", interactive=False)
        demo_a = gr.File(label="Demo Bundle A")
        demo_b = gr.File(label="Demo Bundle B")
        gen_btn.click(fn=ui_generate, inputs=[seed_a, seed_b, chaos], outputs=[gen_msg, demo_a, demo_b])

    with gr.Tab("Fork / patch a bundle"):
        gr.Markdown("Counterfactual workflow: patch an event at index N, then re-hash-chain into a new bundle.")
        src = gr.File(label="Source bundle (.zip)", file_types=[".zip"])
        fork_index = gr.Number(value=0, label="Fork at event index", precision=0)

        with gr.Row():
            patch_kind = gr.Textbox(label="Patch kind (optional)", placeholder="e.g. tool_result")
            patch_step = gr.Textbox(label="Patch step (optional)", placeholder="e.g. t12.tool_result")
        patch_payload = gr.Textbox(
            label="Patch payload JSON (optional)",
            lines=8,
            placeholder='{"ok": true, "value": 42}',
        )

        fork_btn = gr.Button("Create fork bundle", variant="primary")
        fork_msg = gr.Textbox(label="Status", interactive=False)
        fork_file = gr.File(label="Forked bundle (.zip)")
        fork_btn.click(fn=ui_fork, inputs=[src, fork_index, patch_kind, patch_step, patch_payload], outputs=[fork_msg, fork_file])

    gr.Markdown(
        "Tip: In a real agent exporter, emit DRP events at: LLM calls, tool calls/results, memory writes, planner steps, guardrails."
    )

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