"""JSON extractor — pull JSON out of messy LLM text.

Uses agentcast's tolerant extractor. Handles fenced blocks, inline JSON,
trailing prose, and unfenced multi-line objects.
"""

import json
import gradio as gr
from agentcast import extract_json


def extract(messy: str):
    if not messy.strip():
        return "_Paste some text to extract JSON from._", ""
    extracted = extract_json(messy)
    if extracted is None:
        return "❌ **No JSON found.**\n\nTry: fenced ` ```json ... ``` `, inline `{...}`, or top-level array `[...]`.", ""
    pretty = json.dumps(extracted, indent=2, ensure_ascii=False)
    summary = f"✅ **Extracted** ({type(extracted).__name__}, {len(pretty)} chars pretty-printed)"
    return summary, pretty


with gr.Blocks(title="JSON Extractor — for messy LLM output", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # JSON Extractor

        Paste messy LLM output, get clean JSON. Powered by [`agentcast`](https://pypi.org/project/agentcast-py/).

        Handles:
        - Fenced ` ```json ... ``` ` blocks
        - Fenced blocks with no language tag
        - Top-level arrays `[...]`
        - Inline JSON in prose
        - Multi-line unfenced objects
        - Refusals → returns `null`

        Test cases drawn from [`llm-output-extraction-cases`](https://huggingface.co/datasets/mukunda1729/llm-output-extraction-cases) (20 real-world patterns).
        """
    )

    with gr.Row():
        with gr.Column():
            txt = gr.Textbox(
                value='Sure! Here is the answer:\n\n```json\n{"name": "Widget Pro", "price": 29.99}\n```\n\nLet me know if you need anything else!',
                label="Messy LLM output",
                lines=12,
            )
            btn = gr.Button("Extract", variant="primary")
        with gr.Column():
            summary_out = gr.Markdown()
            json_out = gr.Code(language="json", label="Extracted JSON")
    btn.click(extract, inputs=txt, outputs=[summary_out, json_out])

    gr.Examples(
        examples=[
            ['Sure! Here is the answer:\n\n```json\n{"name": "Widget Pro", "price": 29.99}\n```\n\nLet me know!'],
            ['{"answer": 42}'],
            ['[{"k": 1}, {"k": 2}, {"k": 3}]'],
            ['Final:\n\n{\n  "event": "login",\n  "ts": "2026-04-26T12:00:00Z"\n}\n\nDone.'],
            ['I am sorry, I cannot answer that.'],
        ],
        inputs=txt,
    )

    gr.Markdown(
        """
        ---
        Part of [The Agent Reliability Stack](https://mukundakatta.github.io/agent-stack/) · MIT licensed
        """
    )


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)