import gradio as gr
from transformers import pipeline

MODEL_ID = "HuggingFaceTB/SmolLM2-360M-Instruct"

generator = pipeline(
    "text-generation",
    model=MODEL_ID,
    device="cpu",
)

SYSTEM_PROMPT = (
    "You are a writing assistant. You are given a paragraph and a quotation. "
    "Insert the quotation into the paragraph at a sensible point. Wrap the "
    "quotation in double quote marks. Do not change the wording of the "
    "paragraph. Do not change the wording of the quotation. Output only the "
    "revised paragraph, with the quotation inserted."
)


def incorporate_quote(paragraph, quote):
    paragraph = (paragraph or "").strip()
    quote = (quote or "").strip()
    if not paragraph or not quote:
        return "Please provide both a paragraph and a quotation."

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {
            "role": "user",
            "content": f"Paragraph:\n{paragraph}\n\nQuotation:\n{quote}",
        },
    ]

    output = generator(
        messages,
        max_new_tokens=300,
        do_sample=False,
    )

    return output[0]["generated_text"][-1]["content"].strip()


EXAMPLES = [
    [
        "Joan Didion's evening walks are not nature writing in the conventional "
        "sense. They are an attempt to mark time, to register the gradient of "
        "light from one shade of blue to the next.",
        "In certain latitudes there comes a span of time approaching and "
        "following the summer solstice, some weeks in all, when the twilights "
        "turn long and blue.",
    ],
    [
        "Writing teachers often tell their students to use specific imagery "
        "rather than abstractions. The advice is harder to follow than it "
        "sounds, especially for students who have been rewarded all through "
        "school for sounding general and serious.",
        "Show, don't tell.",
    ],
]


with gr.Blocks(title="Quote Incorporation — Simple") as demo:
    gr.Markdown(
        """
        # Quote Incorporation — Space 1 (Simple)

        Drop a quotation into a paragraph with quote marks. Uses a small
        instruction-tuned model (SmolLM2-360M) on Hugging Face free CPU.
        Fast, basic, no cleverness.
        """
    )

    with gr.Row():
        with gr.Column():
            paragraph_in = gr.Textbox(
                label="Paragraph",
                placeholder="Paste a paragraph of writing here…",
                lines=8,
            )
            quote_in = gr.Textbox(
                label="Quotation",
                placeholder="Paste the quotation (without quote marks)…",
                lines=3,
            )
            submit_btn = gr.Button("Incorporate Quote", variant="primary")

        with gr.Column():
            result_out = gr.Textbox(
                label="Revised Paragraph",
                lines=12,
            )

    submit_btn.click(
        fn=incorporate_quote,
        inputs=[paragraph_in, quote_in],
        outputs=[result_out],
    )

    gr.Examples(
        examples=EXAMPLES,
        inputs=[paragraph_in, quote_in],
    )


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