import gradio as gr
from transformers import pipeline

MODEL_ID = "HuggingFaceTB/SmolLM2-1.7B-Instruct"

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

SYSTEM_PROMPT = (
    "You are a writing assistant helping to incorporate a quotation into a "
    "paragraph. Your job has two parts:\n\n"
    "1. Insert the quotation at the most natural point in the paragraph. "
    "Wrap the quotation in double quote marks. Do not change the wording of "
    "the quotation. Do not change the wording of the surrounding paragraph.\n\n"
    "2. Immediately after the quotation, add ONE sentence that explains how "
    "the quotation supports, complicates, or extends the paragraph's main "
    "idea. The added sentence should match the voice of the paragraph.\n\n"
    "Output only the revised paragraph, with the quotation and the one new "
    "sentence in place. No preamble, no notes."
)


def incorporate_with_use(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=400,
        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 — Richer") as demo:
    gr.Markdown(
        """
        # Quote Incorporation — Space 2 (Richer)

        Insert a quotation into a paragraph **and** add one sentence after it
        that uses the quote. Uses a 1.7B-parameter model on Hugging Face
        free CPU — slower than Space 1 (~30–90s per request), and noticeably
        smarter, but still hits real limits.

        *Watch for the constraint: latency, generic explanations, and quote
        drift. Those are the seeds of Space 3.*
        """
    )

    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 + Use Quote", variant="primary"
            )
            gr.Markdown(
                "*This may take 30–90 seconds on free CPU. The model is "
                "doing real work — paragraph parsing, quote placement, "
                "and connective reasoning all on a 2-vCPU container.*"
            )

        with gr.Column():
            result_out = gr.Textbox(
                label="Revised Paragraph (with quote and one explanatory sentence)",
                lines=12,
            )

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

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


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