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

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.pdf_report import write_pdf_report
from drp.simulate import make_demo_bundle_zip, fork_patch_bundle


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


def _tmp_dir() -> str:
    return tempfile.mkdtemp(prefix="drp_")


def _tmp_path(name: str) -> str:
    d = _tmp_dir()
    return os.path.join(d, name)


def _as_alignment_rows(alignment: List[Dict[str, Any]]) -> List[List[Any]]:
    """
    Convert alignment dict rows to Dataframe rows.
    """
    rows = []
    for r in alignment:
        rows.append([r.get("i"), r.get("status"), r.get("kind_a"), r.get("step_a"), r.get("kind_b"), r.get("step_b")])
    return rows


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, str, Dict[str, Any], List[List[Any]], str, str, 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)

    # Summary header (quick context above JSON)
    s = diff.get("summary", {})
    parts = []
    if s.get("first_divergence_index") is None:
        parts.append("Timelines identical")
    else:
        parts.append(f"Identical until step/index {s.get('first_divergence_index')}")
    parts.append(f"{s.get('diff_event_count', 0)} events differ")
    if s.get("missing_event_count", 0):
        parts.append(f"{s.get('missing_event_count')} missing")
    if s.get("final_reward_delta") is not None:
        parts.append(f"Final reward delta: {s.get('final_reward_delta'):.6g}")
    summary_line = " · ".join(parts)

    # Export artifacts
    out_dir = _tmp_dir()
    out_md = os.path.join(out_dir, "report.md")
    out_json = os.path.join(out_dir, "diff.json")
    out_pdf = os.path.join(out_dir, "report.pdf")
    out_zip = os.path.join(out_dir, "drp-diff-report.zip")

    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)

    # PDF: render from markdown lines (plain, reliable)
    md_lines = md.splitlines()
    write_pdf_report(out_pdf, "DRP Differential Report", md_lines)

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

    alignment_rows = _as_alignment_rows(diff.get("alignment", []))

    status = "✅ Diff complete."
    return status, summary_line, diff, alignment_rows, out_zip, out_pdf, out_md


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 forensic reports.")

    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)
        summary = gr.Textbox(label="Summary (fast context)", interactive=False)

        alignment_table = gr.Dataframe(
            headers=["i", "status", "kind_a", "step_a", "kind_b", "step_b"],
            label="Timeline alignment (scan for first divergence)",
            interactive=False,
            wrap=True,
        )

        diff_json = gr.JSON(label="Diff JSON (summary + per-event diffs)")

        with gr.Row():
            report_zip = gr.File(label="Download drp-diff-report.zip (MD + PDF + JSON)")
            report_pdf = gr.File(label="Download report.pdf")
            report_md = gr.File(label="Download report.md")

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

    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(
        "Exporter tip: emit DRP events at LLM calls, tool calls/results, memory writes, planner steps, guardrails, and state snapshots."
    )

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