from transformers import BartForConditionalGeneration, BartTokenizer
import gradio as gr
import torch

# ── 1. Load Model ──────────────────────────────────────────────────
print("Loading model...")

MODEL_NAME = "sshleifer/distilbart-cnn-12-6"
tokenizer  = BartTokenizer.from_pretrained(MODEL_NAME)
model      = BartForConditionalGeneration.from_pretrained(MODEL_NAME)
device     = "cpu"
model      = model.to(device)

print("Model ready ✅")

# ── 2. Summarization Function ──────────────────────────────────────
def summarize_text(user_input, max_len, min_len, bullet_mode, chat_html):
    user_input = user_input.strip()

    if not user_input:
        return chat_html, ""

    if len(user_input.split()) < 30:
        bot_msg = (
            "👋 Hello! I'm your Text Summarizer. "
            "Paste any long article or document (30+ words) and I'll summarize it."
        )
        new_html = chat_html + build_message(user_input, bot_msg)
        return new_html, ""

    try:
        inputs = tokenizer(
            user_input,
            return_tensors="pt",
            max_length=1024,
            truncation=True
        ).to(device)

        summary_ids = model.generate(
            inputs["input_ids"],
            max_new_tokens=int(max_len),
            min_new_tokens=int(min_len),
            num_beams=4,
            length_penalty=2.0,
            early_stopping=True,
            no_repeat_ngram_size=3
        )

        summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)

        if bullet_mode:
            sentences = summary.replace("?", ".").replace("!", ".").split(". ")
            bullets   = "".join(
                f"<li>{s.strip().capitalize()}</li>"
                for s in sentences if s.strip()
            )
            out = f"📌 <b>Summary (Bullet Points):</b><br><ul>{bullets}</ul>"
        else:
            out = f"📌 <b>Summary:</b><br><br>{summary}"

        orig_words = len(user_input.split())
        summ_words = len(summary.split())
        reduction  = round((1 - summ_words / orig_words) * 100, 1)
        out += (
            f"<br><br><hr>"
            f"<small>📊 Original: {orig_words} words → "
            f"Summary: {summ_words} words | "
            f"Reduced by {reduction}%</small>"
        )

        label = user_input[:100] + "..." if len(user_input) > 100 else user_input
        new_html = chat_html + build_message(label, out)

    except Exception as e:
        new_html = chat_html + build_message(user_input[:60], f"❌ Error: {str(e)}")

    return new_html, ""


def build_message(user_text, bot_text):
    return f"""
    <div style='margin:10px 0; padding:10px; background:#1e1e2e; border-radius:8px;'>
        <div style='background:#2a2a3e; padding:8px 12px; border-radius:6px; margin-bottom:8px;'>
            <b style='color:#a0a0ff;'>You:</b><br>
            <span style='color:#e0e0e0;'>{user_text}</span>
        </div>
        <div style='background:#1a3a2a; padding:8px 12px; border-radius:6px;'>
            <b style='color:#80ff80;'>Summarizer:</b><br>
            <span style='color:#e0e0e0;'>{bot_text}</span>
        </div>
    </div>
    """


def clear_chat():
    return "", ""


# ── 3. Gradio UI ───────────────────────────────────────────────────
with gr.Blocks(title="Text Summarizer") as demo:

    gr.Markdown("""
    # 📝 Text Summarization Chatbox
    ### Powered by sshleifer/distilbart-cnn-12-6
    Paste any long text and get an instant summary!
    """)

    with gr.Row():

        # Left — Chat
        with gr.Column(scale=7):

            # Chat display using HTML (works in ALL gradio versions)
            chat_display = gr.HTML(
                value="<div style='height:400px; overflow-y:auto; padding:10px; "
                      "background:#12121f; border-radius:8px; color:#ccc;'>"
                      "💬 Your summaries will appear here...</div>"
            )
            chat_state = gr.State("")

            txt_input = gr.Textbox(
                placeholder="Paste your article, report, or any long text here...",
                show_label=False,
                lines=4
            )
            with gr.Row():
                submit_btn = gr.Button("✨ Summarize", variant="primary")
                clear_btn  = gr.Button("🗑️ Clear",    variant="secondary")

        # Right — Settings
        with gr.Column(scale=3):
            gr.Markdown("### ⚙️ Settings")
            max_length = gr.Slider(
                minimum=50, maximum=300,
                value=130, step=10,
                label="Max Summary Length (tokens)"
            )
            min_length = gr.Slider(
                minimum=10, maximum=100,
                value=30, step=5,
                label="Min Summary Length (tokens)"
            )
            bullet_mode = gr.Checkbox(
                label="🔵 Bullet Point Mode",
                value=False
            )
            gr.Markdown("""
            ---
            **💡 Tips:**
            - Works best with 100–1000 word inputs
            - Articles, news, reports, essays
            - Toggle Bullet Mode for point wise output
            ---
            **Model:** distilbart-cnn-12-6
            **Device:** CPU
            """)

    submit_btn.click(
        summarize_text,
        inputs=[txt_input, max_length, min_length, bullet_mode, chat_state],
        outputs=[chat_display, txt_input]
    )
    txt_input.submit(
        summarize_text,
        inputs=[txt_input, max_length, min_length, bullet_mode, chat_state],
        outputs=[chat_display, txt_input]
    )
    clear_btn.click(clear_chat, outputs=[chat_display, chat_state])

# ── 4. Launch ──────────────────────────────────────────────────────
demo.launch(server_name="0.0.0.0", server_port=7860)