app.py

from __future__ import annotations

import base64
import json
import mimetypes
from pathlib import Path
from typing import Any

import gradio as gr
from PIL import Image

from kgvis.plot import build_figure, empty_figure
from kgvis.schema import graph_from_dict, graph_to_dict, parse_graph


DEFAULT_TOGGLES = ["nodes", "edges", "node_labels", "edge_labels"]

CSS = """
.kgvis-shell { max-width: 1600px; margin: 0 auto; }
.kgvis-status textarea { font-family: ui-monospace, SFMono-Regular, Menlo, monospace; }
.kgvis-plot { min-height: 640px; }
"""


def _path_from_upload(upload: Any) -> Path | None:
    if upload is None:
        return None
    if isinstance(upload, (str, Path)):
        return Path(upload)
    if isinstance(upload, dict):
        for key in ("path", "name"):
            if upload.get(key):
                return Path(upload[key])
    name = getattr(upload, "name", None)
    if name:
        return Path(name)
    return None


def _upload_name(upload: Any, path: Path) -> str:
    if isinstance(upload, dict):
        return str(upload.get("orig_name") or upload.get("name") or path.name)
    return str(getattr(upload, "orig_name", None) or path.name)


def _image_state_from_upload(upload: Any) -> dict[str, Any]:
    path = _path_from_upload(upload)
    if path is None:
        raise ValueError("No image file was provided.")
    if not path.exists():
        raise FileNotFoundError(f"Uploaded image not found: {path}")

    with Image.open(path) as img:
        width, height = img.size
    if height == 0:
        raise ValueError("Image height is zero.")

    mime, _ = mimetypes.guess_type(path.name)
    if not mime:
        mime = "image/png"
    encoded = base64.b64encode(path.read_bytes()).decode("ascii")
    return {
        "name": _upload_name(upload, path),
        "data_uri": f"data:{mime};base64,{encoded}",
        "aspect": width / height,
    }


def _load_graph_text(text: str) -> tuple[dict[str, Any], dict[str, Any] | str]:
    if not text.strip():
        raise ValueError("Paste JSON or token graph text first.")
    try:
        raw = json.loads(text)
    except json.JSONDecodeError:
        raw = text
    graph = parse_graph(raw)
    return graph_to_dict(graph), raw


def _load_graph_upload(upload: Any) -> tuple[dict[str, Any], dict[str, Any] | str, str]:
    path = _path_from_upload(upload)
    if path is None:
        raise ValueError("No graph file was provided.")
    if not path.exists():
        raise FileNotFoundError(f"Uploaded graph file not found: {path}")
    graph_payload, raw_graph = _load_graph_text(path.read_text(encoding="utf-8"))
    return graph_payload, raw_graph, _upload_name(upload, path)


def _raw_display(raw: dict[str, Any] | str | None) -> str:
    if raw is None:
        return ""
    if isinstance(raw, str):
        return raw
    return json.dumps(raw, indent=2, ensure_ascii=False)


def _node_choices(graph_payload: dict[str, Any] | None) -> list[tuple[str, str]]:
    if not graph_payload:
        return []
    graph = graph_from_dict(graph_payload)
    return [
        (f"{node.label or node.id} ({node.id})", node.id)
        for node in graph.nodes
        if node.id
    ]


def _summary(graph_payload: dict[str, Any] | None, graph_name: str | None = None) -> str:
    if not graph_payload:
        return "No graph loaded."
    graph = graph_from_dict(graph_payload)
    name = f"`{graph_name}`: " if graph_name else ""
    schema = graph.meta.get("schema", "unknown")
    return f"{name}{len(graph.nodes)} nodes, {len(graph.edges)} edges (schema: {schema})"


def _dangling_markdown(graph_payload: dict[str, Any] | None) -> str:
    if not graph_payload:
        return "Dangling edges: 0"
    graph = graph_from_dict(graph_payload)
    node_ids = {node.id for node in graph.nodes}
    dangling = [
        edge
        for edge in graph.edges
        if edge.source not in node_ids or edge.target not in node_ids
    ]
    if not dangling:
        return "Dangling edges: 0"
    lines = [f"Dangling edges: {len(dangling)}"]
    lines.extend(
        f"- `{edge.predicate}`: `{edge.source}` -> `{edge.target}`"
        for edge in dangling
    )
    return "\n".join(lines)


def _edge_rows(graph_payload: dict[str, Any] | None) -> list[list[str]]:
    if not graph_payload:
        return []
    graph = graph_from_dict(graph_payload)
    node_ids = {node.id for node in graph.nodes}
    return [
        [
            edge.source,
            edge.predicate,
            edge.target,
            "yes" if edge.source not in node_ids or edge.target not in node_ids else "no",
        ]
        for edge in graph.edges
    ]


def _node_details(
    graph_payload: dict[str, Any] | None,
    selected_node_id: str | None,
) -> dict[str, Any]:
    if not graph_payload:
        return {"message": "No graph loaded."}
    if not selected_node_id:
        return {"message": "Select a node to see details."}

    graph = graph_from_dict(graph_payload)
    node = next((candidate for candidate in graph.nodes if candidate.id == selected_node_id), None)
    if node is None:
        return {"message": "Selected node not found."}

    node_ids = {candidate.id for candidate in graph.nodes}
    outgoing = [
        {
            "predicate": edge.predicate,
            "source": edge.source,
            "target": edge.target,
        }
        for edge in graph.edges
        if edge.source == selected_node_id and edge.target in node_ids
    ]
    incoming = [
        {
            "predicate": edge.predicate,
            "source": edge.source,
            "target": edge.target,
        }
        for edge in graph.edges
        if edge.target == selected_node_id and edge.source in node_ids
    ]
    dangling = [
        {
            "predicate": edge.predicate,
            "source": edge.source,
            "target": edge.target,
        }
        for edge in graph.edges
        if (
            edge.source == selected_node_id
            and edge.target not in node_ids
            or edge.target == selected_node_id
            and edge.source not in node_ids
        )
    ]

    return {
        "id": node.id,
        "label": node.label,
        "attributes": node.attributes,
        "outgoing_edges": outgoing,
        "incoming_edges": incoming,
        "dangling_edges": dangling,
    }


def _render_figure(
    graph_payload: dict[str, Any] | None,
    image_payload: dict[str, Any] | None,
    toggles: list[str] | None,
    layout_mode: str | None,
    show_image: bool,
    background_opacity: float | None,
    show_bboxes: bool,
    selected_node_id: str | None,
):
    if not graph_payload:
        return empty_figure("Upload a graph file to begin.")

    graph = graph_from_dict(graph_payload)
    active_toggles = toggles or []
    background = None
    background_aspect = None
    if show_image and image_payload:
        background = image_payload.get("data_uri")
        background_aspect = image_payload.get("aspect")

    return build_figure(
        graph,
        show_nodes="nodes" in active_toggles,
        show_edges="edges" in active_toggles,
        show_node_labels="node_labels" in active_toggles and "nodes" in active_toggles,
        show_edge_labels="edge_labels" in active_toggles and "edges" in active_toggles,
        selected_node_id=selected_node_id,
        layout_mode=layout_mode or "force",
        background_image=background,
        background_aspect=background_aspect,
        background_opacity=background_opacity if background_opacity is not None else 0.35,
        show_bboxes=show_bboxes,
        highlight_node_id=selected_node_id,
    )


def _render_visual_outputs(
    graph_payload: dict[str, Any] | None,
    image_payload: dict[str, Any] | None,
    toggles: list[str] | None,
    layout_mode: str | None,
    show_image: bool,
    background_opacity: float | None,
    show_bboxes: bool,
    selected_node_id: str | None,
):
    return (
        _render_figure(
            graph_payload,
            image_payload,
            toggles,
            layout_mode,
            show_image,
            background_opacity,
            show_bboxes,
            selected_node_id,
        ),
        _node_details(graph_payload, selected_node_id),
    )


def handle_graph_upload(
    graph_file: Any,
    image_payload: dict[str, Any] | None,
    toggles: list[str] | None,
    layout_mode: str | None,
    show_image: bool,
    background_opacity: float | None,
    show_bboxes: bool,
):
    if graph_file is None:
        return (
            None,
            None,
            None,
            gr.update(choices=[], value=None),
            "No graph loaded.",
            "Dangling edges: 0",
            [],
            empty_figure("Upload or paste a graph to begin."),
            {"message": "No graph loaded."},
            "",
            "Graph cleared.",
        )

    try:
        graph_payload, raw_graph, graph_name = _load_graph_upload(graph_file)
    except Exception as exc:
        return (
            None,
            None,
            None,
            gr.update(choices=[], value=None),
            f"Failed to parse graph: {exc}",
            "Dangling edges: 0",
            [],
            empty_figure("Failed to parse graph file."),
            {"message": "No graph loaded."},
            "",
            f"Upload failed: {exc}",
        )

    return (
        graph_payload,
        raw_graph,
        graph_name,
        gr.update(choices=_node_choices(graph_payload), value=None),
        _summary(graph_payload, graph_name),
        _dangling_markdown(graph_payload),
        _edge_rows(graph_payload),
        _render_figure(
            graph_payload,
            image_payload,
            toggles,
            layout_mode,
            show_image,
            background_opacity,
            show_bboxes,
            None,
        ),
        _node_details(graph_payload, None),
        _raw_display(raw_graph),
        f"Loaded {graph_name}.",
    )


def handle_graph_text(
    graph_text: str | None,
    image_payload: dict[str, Any] | None,
    toggles: list[str] | None,
    layout_mode: str | None,
    show_image: bool,
    background_opacity: float | None,
    show_bboxes: bool,
):
    if not graph_text or not graph_text.strip():
        return (
            None,
            None,
            None,
            gr.update(choices=[], value=None),
            "No graph loaded.",
            "Dangling edges: 0",
            [],
            empty_figure("Paste a graph to begin."),
            {"message": "No graph loaded."},
            "",
            "Graph cleared.",
        )

    try:
        graph_payload, raw_graph = _load_graph_text(graph_text)
    except Exception as exc:
        return (
            None,
            None,
            None,
            gr.update(choices=[], value=None),
            f"Failed to parse graph: {exc}",
            "Dangling edges: 0",
            [],
            empty_figure("Failed to parse graph text."),
            {"message": "No graph loaded."},
            "",
            f"Parse failed: {exc}",
        )

    return (
        graph_payload,
        raw_graph,
        "pasted graph",
        gr.update(choices=_node_choices(graph_payload), value=None),
        _summary(graph_payload, "pasted graph"),
        _dangling_markdown(graph_payload),
        _edge_rows(graph_payload),
        _render_figure(
            graph_payload,
            image_payload,
            toggles,
            layout_mode,
            show_image,
            background_opacity,
            show_bboxes,
            None,
        ),
        _node_details(graph_payload, None),
        _raw_display(raw_graph),
        "Parsed pasted graph.",
    )


def handle_image_upload(
    image_file: Any,
    graph_payload: dict[str, Any] | None,
    toggles: list[str] | None,
    layout_mode: str | None,
    show_image: bool,
    background_opacity: float | None,
    show_bboxes: bool,
    selected_node_id: str | None,
):
    if image_file is None:
        return (
            None,
            _render_figure(
                graph_payload,
                None,
                toggles,
                layout_mode,
                show_image,
                background_opacity,
                show_bboxes,
                selected_node_id,
            ),
            "Image cleared.",
        )

    try:
        image_payload = _image_state_from_upload(image_file)
    except Exception as exc:
        return (
            None,
            _render_figure(
                graph_payload,
                None,
                toggles,
                layout_mode,
                show_image,
                background_opacity,
                show_bboxes,
                selected_node_id,
            ),
            f"Image upload failed: {exc}",
        )

    return (
        image_payload,
        _render_figure(
            graph_payload,
            image_payload,
            toggles,
            layout_mode,
            show_image,
            background_opacity,
            show_bboxes,
            selected_node_id,
        ),
        f"Loaded image {image_payload['name']}.",
    )


def create_app() -> gr.Blocks:
    with gr.Blocks(title="KGVis", fill_width=True) as app:
        graph_state = gr.State(None)
        raw_state = gr.State(None)
        graph_name_state = gr.State(None)
        image_state = gr.State(None)

        with gr.Column(elem_classes=["kgvis-shell"]):
            gr.Markdown("# KGVis")

            with gr.Row():
                with gr.Column(scale=1, min_width=300):
                    graph_file = gr.File(
                        label="Upload Graph",
                        file_count="single",
                        file_types=[".json", ".txt"],
                        type="filepath",
                    )
                    graph_text = gr.Textbox(
                        label="Paste Graph",
                        placeholder="Paste JSON or <|graph|> token text here.",
                        interactive=True,
                        lines=10,
                        max_lines=24,
                    )
                    parse_graph_button = gr.Button(
                        "Parse Graph",
                        variant="primary",
                    )
                    image_file = gr.File(
                        label="Upload Image",
                        file_count="single",
                        file_types=[".png", ".jpg", ".jpeg"],
                        type="filepath",
                    )
                    status = gr.Textbox(
                        label="Status",
                        value="Upload a graph file, or paste a graph and click Parse Graph.",
                        interactive=False,
                        lines=2,
                        elem_classes=["kgvis-status"],
                    )
                    toggles = gr.CheckboxGroup(
                        label="View Options",
                        choices=[
                            ("Show nodes", "nodes"),
                            ("Show edges", "edges"),
                            ("Show node labels", "node_labels"),
                            ("Show edge labels", "edge_labels"),
                        ],
                        value=DEFAULT_TOGGLES,
                    )
                    layout_mode = gr.Radio(
                        label="Layout",
                        choices=[
                            ("Force layout", "force"),
                            ("Image layout (bbox)", "bbox"),
                        ],
                        value="force",
                    )
                    show_bboxes = gr.Checkbox(label="Show bounding boxes", value=False)
                    show_image = gr.Checkbox(label="Show image", value=True)
                    background_opacity = gr.Slider(
                        label="Background opacity",
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05,
                        value=0.35,
                    )
                    selected_node = gr.Dropdown(
                        label="Selected Node",
                        choices=[],
                        value=None,
                        interactive=True,
                    )

                with gr.Column(scale=3, min_width=640):
                    summary = gr.Markdown("No graph loaded.")
                    plot = gr.Plot(
                        label="Graph",
                        value=empty_figure("Upload or paste a graph to begin."),
                        elem_classes=["kgvis-plot"],
                    )
                    dangling = gr.Markdown("Dangling edges: 0")

                    with gr.Row():
                        with gr.Column(scale=1):
                            node_details = gr.JSON(
                                label="Node Details",
                                value={"message": "No graph loaded."},
                            )
                        with gr.Column(scale=1):
                            edge_table = gr.Dataframe(
                                label="Edges",
                                headers=["Source", "Predicate", "Target", "Dangling"],
                                datatype=["str", "str", "str", "str"],
                                interactive=False,
                            )

                    raw_json = gr.Textbox(
                        label="Raw Graph",
                        value="",
                        interactive=False,
                        lines=12,
                        max_lines=24,
                    )

        graph_file.change(
            handle_graph_upload,
            inputs=[
                graph_file,
                image_state,
                toggles,
                layout_mode,
                show_image,
                background_opacity,
                show_bboxes,
            ],
            outputs=[
                graph_state,
                raw_state,
                graph_name_state,
                selected_node,
                summary,
                dangling,
                edge_table,
                plot,
                node_details,
                raw_json,
                status,
            ],
        )

        parse_graph_button.click(
            handle_graph_text,
            inputs=[
                graph_text,
                image_state,
                toggles,
                layout_mode,
                show_image,
                background_opacity,
                show_bboxes,
            ],
            outputs=[
                graph_state,
                raw_state,
                graph_name_state,
                selected_node,
                summary,
                dangling,
                edge_table,
                plot,
                node_details,
                raw_json,
                status,
            ],
        )

        image_file.change(
            handle_image_upload,
            inputs=[
                image_file,
                graph_state,
                toggles,
                layout_mode,
                show_image,
                background_opacity,
                show_bboxes,
                selected_node,
            ],
            outputs=[image_state, plot, status],
        )

        for control in (
            toggles,
            layout_mode,
            show_image,
            background_opacity,
            show_bboxes,
            selected_node,
        ):
            control.change(
                _render_visual_outputs,
                inputs=[
                    graph_state,
                    image_state,
                    toggles,
                    layout_mode,
                    show_image,
                    background_opacity,
                    show_bboxes,
                    selected_node,
                ],
                outputs=[plot, node_details],
            )

    return app


if __name__ == "__main__":
    create_app().launch(css=CSS)