import os
import json
from difflib import unified_diff

import gradio as gr

import rft_flightrecorder as fr

LOG_PATH = fr.DEFAULT_LOG_PATH


def text_diff(before_text: str, after_text: str):
    a = (before_text or "").splitlines()
    b = (after_text or "").splitlines()
    diff = list(unified_diff(a, b, fromfile="before", tofile="after", lineterm=""))
    return "\n".join(diff) if diff else "(no differences)"


def download_log():
    return LOG_PATH if os.path.exists(LOG_PATH) else None


def ui_start_session(model_id, run_mode, notes, sign_start, sk_hex):
    sid, msg = fr.start_session(LOG_PATH, model_id, run_mode, notes, sign_start, sk_hex)
    # fan-out the session id into all tabs
    return sid, sid, sid, sid, sid, msg


def ui_append_event(session_id, event_type, parent_hash, payload_text, sign_event, sk_hex, model_id, run_mode):
    ev, msg = fr.append_event(
        log_path=LOG_PATH,
        session_id=session_id,
        event_type=event_type,
        payload_text=payload_text,
        parent_event_hash=parent_hash,
        sign_event=sign_event,
        sk_hex=sk_hex,
        model_id=model_id,
        run_mode=run_mode,
    )
    return ev, msg


def ui_timeline(session_id):
    rows, msg = fr.session_timeline_rows(LOG_PATH, session_id)
    return rows, msg


def ui_verify(session_id, pk_hex, require_sigs):
    return fr.verify_session(LOG_PATH, session_id, pk_hex, require_sigs)


def ui_finalise(session_id, sign_anchor, sk_hex, model_id, run_mode):
    return fr.finalise_session(LOG_PATH, session_id, sign_anchor, sk_hex, model_id, run_mode)


def ui_export(session_id):
    return fr.export_session_bundle(LOG_PATH, session_id)


def ui_list_sessions():
    sessions, msg = fr.list_sessions(LOG_PATH)
    return gr.Dropdown(choices=sessions, value=(sessions[-1] if sessions else None)), msg


def ui_pick_session(sid):
    sid = (sid or "").strip()
    return sid, sid, sid, sid, sid


def ui_get_event(session_id, ev_hash):
    return fr.get_event_by_hash(LOG_PATH, session_id, ev_hash)


def ui_import_bundle(bundle_file, pk_hex, require_sigs, store_into_log):
    # gr.File returns a path string
    status, ok, report, stored_msg = fr.import_bundle_verify(
        bundle_path=bundle_file,
        pk_hex=pk_hex,
        require_signatures=require_sigs,
        store_into_log=store_into_log,
        log_path=LOG_PATH,
    )
    extra = (stored_msg or "")
    return status, ok, report + (("\n\n" + extra) if extra else "")


# ============================================================
# Quickstart (1-click) demo
# ============================================================

def _j(obj) -> str:
    return json.dumps(obj, ensure_ascii=False)


def ui_quickstart_run(sign_all: bool, current_sk: str, current_pk: str):
    status_lines = []

    # Decide keys
    if sign_all:
        sk_hex, pk_hex = fr.gen_keys()
        status_lines.append("[OK] Generated fresh Ed25519 keypair for this demo run.")
    else:
        sk_hex, pk_hex = (current_sk or ""), (current_pk or "")
        status_lines.append("[OK] Running unsigned demo (hash-chain only).")

    model_id = "rft-flightrecorder-demo"
    run_mode = "deterministic"

    # Start
    sid, start_msg = fr.start_session(
        LOG_PATH,
        model_id=model_id,
        run_mode=run_mode,
        notes="Quickstart demo: prompt → tool_call/result → output → memory_write → finalise → export",
        sign_start=bool(sign_all),
        sk_hex=sk_hex,
    )

    if not sid:
        status_lines.append(f"[FAIL] start_session: {start_msg}")
        quick_msg = "\n".join(status_lines)
        return (
            (sk_hex if sign_all else current_sk),
            (pk_hex if sign_all else current_pk),
            "", "", "", "", "",  # session ids fanout
            quick_msg,            # start_status
            [], "No timeline.",
            "FAIL", False, "Quickstart failed at start_session.",
            None, "Not finalised.",
            None, "No export.",
            quick_msg,            # quick_status
        )

    status_lines.append(f"[OK] Started session: {sid}")

    # Append a realistic sequence of events
    demo_events = [
        ("prompt", {"text": "Summarise why tamper-evident memory logs matter for agent systems."}),
        ("tool_call", {"tool": "search", "input": {"q": "tamper-evident audit log hash chain"}, "id": "call_01"}),
        ("tool_result", {"id": "call_01", "ok": True, "items": [{"title": "Hash chaining overview", "source": "demo"}]}),
        ("output", {"text": "Tamper-evident logs make history verifiable: edits break the chain and verification fails."}),
        ("memory_write", {"key": "policy.audit_mode", "before": "off", "after": "on", "reason": "enable strict audit trail"}),
        ("note", {"checkpoint": "demo_complete", "expected_next": "finalise + export"}),
    ]

    for etype, payload in demo_events:
        ev, msg = fr.append_event(
            log_path=LOG_PATH,
            session_id=sid,
            event_type=etype,
            payload_text=_j(payload),
            parent_event_hash="",
            sign_event=bool(sign_all),
            sk_hex=sk_hex,
            model_id=model_id,
            run_mode=run_mode,
        )
        if not ev:
            status_lines.append(f"[FAIL] append_event({etype}): {msg}")
            break
        status_lines.append(f"[OK] appended {etype} (seq={ev.get('seq')})")

    # Load timeline now (even if partial)
    tl_rows, tl_msg = fr.session_timeline_rows(LOG_PATH, sid)

    # Verify (pre-finalise)
    verify_status, verify_ok, verify_report = fr.verify_session(
        LOG_PATH,
        sid,
        pk_hex=(pk_hex if sign_all else ""),
        require_signatures=bool(sign_all),
    )
    status_lines.append(f"[{'OK' if verify_ok else 'FAIL'}] verify_session (pre-finalise): {verify_status}")

    anchor = None
    fin_msg = "Not finalised."
    export_path = None
    export_msg = "No export."

    if verify_ok:
        anchor, fin_msg = fr.finalise_session(
            LOG_PATH,
            sid,
            sign_anchor=bool(sign_all),
            sk_hex=sk_hex,
            model_id=model_id,
            run_mode=run_mode,
        )
        status_lines.append(f"[OK] finalise_session: {fin_msg}")

        export_path, export_msg = fr.export_session_bundle(LOG_PATH, sid)
        status_lines.append(f"[OK] export_session_bundle: {export_msg}")
    else:
        status_lines.append("[SKIP] Finalise/export skipped because verification failed.")

    quick_msg = f"Quickstart complete. session_id={sid}\n\n" + "\n".join(status_lines)

    return (
        (sk_hex if sign_all else current_sk),
        (pk_hex if sign_all else current_pk),
        sid, sid, sid, sid, sid,        # session ids fanout
        quick_msg,                      # start_status
        tl_rows, tl_msg,
        verify_status, verify_ok, verify_report,
        anchor, fin_msg,
        export_path, export_msg,
        quick_msg,                      # quick_status
    )


with gr.Blocks(title="RFT Agent Flight Recorder — Black Box Trace + Third-Party Verification") as demo:
    gr.Markdown(
        "# RFT Agent Flight Recorder — Black Box Trace + Third-Party Verification\n"
        "This Space records a **tamper-evident, hash-chained event timeline** for AI/agent runs.\n\n"
        "**Core deliverable:** export a ZIP bundle that **anyone can verify**.\n\n"
        "**Key safety note:** Public demo. Do not paste production private keys here."
    )

    # ------------------------------------------------------------
    # Quickstart FIRST (most important UX improvement)
    # ------------------------------------------------------------
    with gr.Tab("Quickstart (1-click)"):
        gr.Markdown(
            "## Quickstart (1-click)\n"
            "If you just want to see the full workflow **without guessing what to click**, use this.\n\n"
            "**This will:** start a session → append events → verify → finalise → export a ZIP proof bundle → "
            "and auto-fill the other tabs with the `session_id`, timeline, reports, and export file."
        )
        quick_sign = gr.Checkbox(label="Sign everything (Ed25519) + generate fresh keys for this run", value=False)
        quick_run = gr.Button("Run Quickstart Demo Now")
        quick_status = gr.Textbox(label="Quickstart status", lines=14)

    with gr.Tab("Keys"):
        sk = gr.Textbox(label="Ed25519 Private Key (hex)", lines=2)
        pk = gr.Textbox(label="Ed25519 Public Key (hex)", lines=2)
        gen = gr.Button("Generate Keypair")
        gen.click(fn=fr.gen_keys, outputs=[sk, pk])

    with gr.Tab("Sessions"):
        with gr.Row():
            list_btn = gr.Button("Refresh session list")
            sessions_dd = gr.Dropdown(label="Existing sessions", choices=[], value=None)
        sessions_msg = gr.Textbox(label="Status", lines=1)

        sid_start = gr.Textbox(label="session_id (Start/Record)", lines=1)
        sid_tl = gr.Textbox(label="session_id (Timeline)", lines=1)
        sid_verify = gr.Textbox(label="session_id (Verify)", lines=1)
        sid_final = gr.Textbox(label="session_id (Finalise/Export)", lines=1)
        sid_record = gr.Textbox(label="session_id (Record Event)", lines=1)

        list_btn.click(fn=ui_list_sessions, outputs=[sessions_dd, sessions_msg])
        sessions_dd.change(fn=ui_pick_session, inputs=[sessions_dd], outputs=[sid_start, sid_record, sid_tl, sid_verify, sid_final])

    with gr.Tab("Start Session"):
        model_id = gr.Textbox(label="Model ID", value="audit-demo")
        run_mode = gr.Radio(["deterministic", "creative"], label="Run mode", value="deterministic")
        notes = gr.Textbox(label="Notes (optional)", lines=3)
        sign_start = gr.Checkbox(label="Sign session_start event", value=False)
        start_btn = gr.Button("Start New Session")
        start_status = gr.Textbox(label="Status", lines=2)

        start_btn.click(
            fn=ui_start_session,
            inputs=[model_id, run_mode, notes, sign_start, sk],
            outputs=[sid_start, sid_record, sid_tl, sid_verify, sid_final, start_status],
        )

    with gr.Tab("Record Event"):
        event_type = gr.Dropdown(
            choices=[
                "prompt",
                "output",
                "tool_call",
                "tool_result",
                "memory_read",
                "memory_write",
                "retrieval",
                "policy_block",
                "error",
                "note",
            ],
            value="note",
            label="event_type",
        )
        parent_hash = gr.Textbox(label="parent_event_hash_sha256 (optional). If empty, defaults to previous event.", lines=1)
        payload_text = gr.Textbox(
            label="payload (JSON or plain text)",
            lines=10,
            placeholder='Example JSON:\n{\n  "tool":"search",\n  "input":{"q":"..."},\n  "output":{"items":[...]}\n}\n',
        )
        sign_event = gr.Checkbox(label="Sign this event (Ed25519)", value=False)
        append_btn = gr.Button("Append Event")
        event_out = gr.JSON(label="event.json")
        append_status = gr.Textbox(label="Status", lines=2)

        append_btn.click(
            fn=ui_append_event,
            inputs=[sid_record, event_type, parent_hash, payload_text, sign_event, sk, model_id, run_mode],
            outputs=[event_out, append_status],
        )

        flightlog_file = gr.File(label="flightlog.jsonl (download)")
        gr.Button("Download flightlog.jsonl").click(fn=download_log, outputs=[flightlog_file])

    with gr.Tab("Timeline"):
        refresh = gr.Button("Load timeline")
        tl_status = gr.Textbox(label="Status", lines=1)
        tl = gr.Dataframe(
            headers=[
                "seq",
                "ts_utc",
                "event_type",
                "model_id",
                "run_mode",
                "parent_hash",
                "prev_hash",
                "event_hash",
                "signed",
            ],
            datatype=["number", "str", "str", "str", "str", "str", "str", "str", "str"],
            row_count=10,
            col_count=(9, "fixed"),
            label="Event timeline",
            wrap=True,
        )

        refresh.click(fn=ui_timeline, inputs=[sid_tl], outputs=[tl, tl_status])

        with gr.Accordion("View event by hash", open=False):
            ev_hash_in = gr.Textbox(label="event_hash_sha256", lines=1)
            ev_get = gr.Button("Get event")
            ev_status = gr.Textbox(label="Status", lines=1)
            ev_json = gr.JSON(label="event.json")
            ev_get.click(fn=ui_get_event, inputs=[sid_tl, ev_hash_in], outputs=[ev_json, ev_status])

    with gr.Tab("Verify Session"):
        require_sigs = gr.Checkbox(label="Require signatures on every event", value=False)
        verify_btn = gr.Button("Verify")
        verify_msg = gr.Textbox(label="Result", lines=1)
        verify_ok = gr.Checkbox(label="Valid", value=False)
        verify_report = gr.Textbox(label="Report", lines=14)

        verify_btn.click(
            fn=ui_verify,
            inputs=[sid_verify, pk, require_sigs],
            outputs=[verify_msg, verify_ok, verify_report],
        )

    with gr.Tab("Finalise + Export"):
        sign_anchor = gr.Checkbox(label="Sign session anchor + session_end event", value=False)
        fin_btn = gr.Button("Finalise session")
        anchor_out = gr.JSON(label="session_anchor.json")
        fin_status = gr.Textbox(label="Status", lines=2)

        fin_btn.click(
            fn=ui_finalise,
            inputs=[sid_final, sign_anchor, sk, model_id, run_mode],
            outputs=[anchor_out, fin_status],
        )

        export_btn = gr.Button("Export session bundle (ZIP)")
        export_status = gr.Textbox(label="Export status", lines=1)
        export_file = gr.File(label="Bundle download")

        export_btn.click(
            fn=ui_export,
            inputs=[sid_final],
            outputs=[export_file, export_status],
        )

    with gr.Tab("Import Bundle"):
        bundle = gr.File(label="Upload rft_flight_bundle_*.zip")
        store_into_log = gr.Checkbox(label="Store imported events into local flightlog.jsonl (only if PASS)", value=False)
        import_require_sigs = gr.Checkbox(label="Require signatures on every imported event", value=False)
        import_btn = gr.Button("Verify bundle")
        import_status = gr.Textbox(label="Result", lines=1)
        import_ok = gr.Checkbox(label="Valid", value=False)
        import_report = gr.Textbox(label="Report", lines=14)

        import_btn.click(
            fn=ui_import_bundle,
            inputs=[bundle, pk, import_require_sigs, store_into_log],
            outputs=[import_status, import_ok, import_report],
        )

    with gr.Tab("Diagnostics"):
        diag_btn = gr.Button("Run diagnostics")
        diag_out = gr.JSON(label="diagnostics.json")
        diag_btn.click(fn=lambda: fr.diagnostics(LOG_PATH), outputs=[diag_out])

    with gr.Tab("Diff Helper"):
        before = gr.Textbox(label="Before (text/JSON)", lines=8)
        after = gr.Textbox(label="After (text/JSON)", lines=8)
        diff_btn = gr.Button("Generate unified diff")
        diff_out = gr.Textbox(label="Diff", lines=14)
        diff_btn.click(fn=text_diff, inputs=[before, after], outputs=[diff_out])

    # Wire Quickstart click NOW that all components exist
    quick_run.click(
        fn=ui_quickstart_run,
        inputs=[quick_sign, sk, pk],
        outputs=[
            sk,
            pk,
            sid_start, sid_record, sid_tl, sid_verify, sid_final,
            start_status,
            tl, tl_status,
            verify_msg, verify_ok, verify_report,
            anchor_out, fin_status,
            export_file, export_status,
            quick_status,
        ],
    )

demo.launch()