from transformers import (
    BartTokenizer, BartForConditionalGeneration,
    T5Tokenizer, T5ForConditionalGeneration,
    PegasusTokenizer, PegasusForConditionalGeneration,
)
import gradio as gr
import re
from collections import Counter

# MODEL REGISTRY
MODEL_OPTIONS = {
    "BART  (facebook/bart-large-cnn)":    "bart",
    "T5    (t5-small)":                   "t5",
    "Pegasus (google/pegasus-xsum)":      "pegasus",
}

# Lazy-load cache so we only download what the user picks
_cache = {}

def load_model(key):
    if key in _cache:
        return _cache[key]

    if key == "bart":
        tok = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
        mdl = BartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")
    elif key == "t5":
        tok = T5Tokenizer.from_pretrained("t5-small")
        mdl = T5ForConditionalGeneration.from_pretrained("t5-small")
    elif key == "pegasus":
        tok = PegasusTokenizer.from_pretrained("google/pegasus-xsum")
        mdl = PegasusForConditionalGeneration.from_pretrained("google/pegasus-xsum")
    else:
        raise ValueError(f"Unknown model key: {key}")

    _cache[key] = (tok, mdl)
    return tok, mdl


# KEYWORD EXTRACTION (TF-IDF style, no extra library needed)
STOPWORDS = {
    "the","a","an","and","or","but","in","on","at","to","for","of","with",
    "is","are","was","were","be","been","being","have","has","had","do","does",
    "did","will","would","could","should","may","might","this","that","these",
    "those","it","its","i","we","you","he","she","they","them","their","our",
    "your","his","her","from","by","as","not","no","so","if","about","which",
    "who","whom","when","where","how","what","all","also","just","more","than",
}

def extract_keywords(text, top_n=5):
    words = re.findall(r'\b[a-zA-Z]{4,}\b', text.lower())
    filtered = [w for w in words if w not in STOPWORDS]
    freq = Counter(filtered)
    top = freq.most_common(top_n)
    if not top:
        return "No keywords found."
    return "  ·  ".join([f"🏷️ {w}" for w, _ in top])


# READABILITY SCORE (Flesch Reading Ease)
def count_syllables(word):
    word = word.lower().strip(".,!?;:")
    vowels = "aeiouy"
    count = 0
    prev_vowel = False
    for ch in word:
        is_v = ch in vowels
        if is_v and not prev_vowel:
            count += 1
        prev_vowel = is_v
    if word.endswith("e") and count > 1:
        count -= 1
    return max(1, count)

def flesch_score(text):
    sentences = [s.strip() for s in re.split(r'[.!?]+', text) if s.strip()]
    words     = re.findall(r'\b\w+\b', text)
    if not sentences or not words:
        return 0.0
    syllables = sum(count_syllables(w) for w in words)
    score = (
        206.835
        - 1.015  * (len(words) / len(sentences))
        - 84.6   * (syllables / len(words))
    )
    return round(max(0, min(100, score)), 1)

def grade_label(score):
    if score >= 90: return "Very Easy (Grade 5)"
    if score >= 80: return "Easy (Grade 6)"
    if score >= 70: return "Fairly Easy (Grade 7)"
    if score >= 60: return "Standard (Grade 8–9)"
    if score >= 50: return "Fairly Difficult (Grade 10–12)"
    if score >= 30: return "Difficult (College)"
    return "Very Difficult (Professional)"

def readability_report(original, summarized):
    if not original.strip() or not summarized.strip():
        return "—"
    o_score = flesch_score(original)
    s_score = flesch_score(summarized)
    arrow   = "⬆️ Easier" if s_score > o_score else ("⬇️ Harder" if s_score < o_score else "➡️ Same")
    return (
        f"📄 Original  : {o_score} / 100  →  {grade_label(o_score)}\n"
        f"📝 Summary   : {s_score} / 100  →  {grade_label(s_score)}\n"
        f"📊 Change    : {arrow}"
    )


# CORE SUMMARIZE FUNCTION
def run_summary(text, model_label, max_length):
    if not text.strip():
        return "⚠️ Please enter some text.", "—", "—"

    key = MODEL_OPTIONS[model_label]
    tokenizer, model = load_model(key)

    input_text = ("summarize: " + text) if key == "t5" else text

    inputs = tokenizer.encode(
        input_text,
        return_tensors="pt",
        max_length=1024,
        truncation=True
    )
    ids = model.generate(
        inputs,
        max_length=int(max_length),
        min_length=30,
        do_sample=False
    )
    result = tokenizer.decode(ids[0], skip_special_tokens=True)

    keywords    = extract_keywords(result)
    readability = readability_report(text, result)

    return result, keywords, readability


def run_file_summary(file, model_label, max_length):
    if file is None:
        return "⚠️ Please upload a .txt file.", "—", "—"
    with open(file.name, "r", encoding="utf-8") as f:
        content = f.read()
    if not content.strip():
        return "⚠️ The uploaded file appears to be empty.", "—", "—"
    return run_summary(content, model_label, max_length)


# HELPERS
def clear_text():
    return "", "", "—", "—"

def clear_file():
    return None, "", "—", "—"

def sync_style(choice):
    return {"Short (50)": 50, "Medium (150)": 150, "Detailed (300)": 300}.get(choice, 150)


# CSS (same as original)
css = """
@import url('https://fonts.googleapis.com/css2?family=Syne:wght@400;600;700&family=DM+Sans:wght@300;400;500&display=swap');

:root {
    --qd-bg:      #0a0a0f;
    --qd-surface: #13131a;
    --qd-card:    #1a1a24;
    --qd-accent:  #7c6af7;
    --qd-accent2: #4ecdc4;
    --qd-text:    #f0eff8;
    --qd-muted:   #8887a0;
    --qd-border:  #2a2a3a;
}

body, .gradio-container {
    background: var(--qd-bg) !important;
    font-family: 'DM Sans', sans-serif !important;
    color: var(--qd-text) !important;
}

#qd-header { text-align: center; padding: 2rem 1rem 1rem; }
#qd-logo {
    font-family: 'Syne', sans-serif;
    font-size: 2.4rem; font-weight: 700;
    letter-spacing: -1px; color: var(--qd-text);
}
#qd-logo span { color: var(--qd-accent); }
#qd-tagline {
    display: inline-block; margin-top: 6px;
    background: #1e1b3a; border: 1px solid var(--qd-accent);
    color: var(--qd-accent); font-size: 11px;
    padding: 3px 14px; border-radius: 20px;
    letter-spacing: 1.5px; font-weight: 500;
}

.tab-nav {
    background: var(--qd-surface) !important;
    border-radius: 12px !important; padding: 6px !important;
    border: none !important; gap: 6px !important;
}
.tab-nav button {
    background: transparent !important; border: none !important;
    border-radius: 8px !important; color: var(--qd-muted) !important;
    font-family: 'DM Sans', sans-serif !important;
    font-size: 13px !important; font-weight: 500 !important;
    padding: 10px 20px !important; transition: all 0.2s !important;
}
.tab-nav button.selected { background: var(--qd-accent) !important; color: #fff !important; }

label span, .label-wrap span {
    font-family: 'DM Sans', sans-serif !important;
    font-size: 12px !important; font-weight: 500 !important;
    color: var(--qd-muted) !important; letter-spacing: 0.8px !important;
    text-transform: uppercase !important;
}

textarea, input[type="text"] {
    background: var(--qd-surface) !important;
    border: 1px solid var(--qd-border) !important;
    border-radius: 10px !important; color: var(--qd-text) !important;
    font-family: 'DM Sans', sans-serif !important;
    font-size: 14px !important; padding: 12px !important;
    transition: border 0.2s !important;
}
textarea:focus, input[type="text"]:focus {
    border-color: var(--qd-accent) !important;
    outline: none !important;
    box-shadow: 0 0 0 2px rgba(124,106,247,0.15) !important;
}

input[type="range"] { accent-color: var(--qd-accent) !important; }

#summarize-btn, #file-summarize-btn {
    background: var(--qd-accent) !important; border: none !important;
    border-radius: 10px !important; color: #fff !important;
    font-family: 'Syne', sans-serif !important;
    font-size: 14px !important; font-weight: 600 !important;
    letter-spacing: 0.5px !important; padding: 13px !important;
    transition: opacity 0.2s !important; width: 100% !important;
}
#summarize-btn:hover, #file-summarize-btn:hover { opacity: 0.85 !important; }

#copy-btn, #file-copy-btn {
    background: #1e1b3a !important;
    border: 1px solid var(--qd-accent) !important;
    color: var(--qd-accent) !important;
    border-radius: 8px !important; font-size: 13px !important;
}
#clear-btn, #file-clear-btn {
    background: #1a1015 !important; border: 1px solid #3a2a2a !important;
    color: #c07070 !important; border-radius: 8px !important;
    font-size: 13px !important;
}

#kw-box, #file-kw-box {
    background: #0f1a1a !important;
    border: 1px solid var(--qd-accent2) !important;
    border-radius: 10px !important; color: var(--qd-accent2) !important;
    font-size: 13px !important; font-weight: 500 !important;
}
#read-box, #file-read-box {
    background: #0f0f1a !important;
    border: 1px solid #534AB7 !important;
    border-radius: 10px !important; color: #b0a8f7 !important;
    font-size: 13px !important; font-family: monospace !important;
}

.gr-box, .gr-form, .gr-panel, .block {
    background: var(--qd-card) !important;
    border: 1px solid var(--qd-border) !important;
    border-radius: 16px !important;
}

footer { display: none !important; }
"""


# UI
with gr.Blocks(css=css, theme=gr.themes.Base(), title="⚡ QuickDigest") as demo:

    gr.HTML("""
    <div id="qd-header">
        <div id="qd-logo">⚡ Quick<span>Digest</span></div>
        <div id="qd-tagline">AI POWERED &nbsp;·&nbsp; BART · T5 · PEGASUS</div>
    </div>
    """)

    with gr.Tabs():

        # Tab 1: Text Summarizer
        with gr.Tab("✍️ Text Summarizer"):
            with gr.Row():
                with gr.Column():
                    text_input = gr.Textbox(
                        label="Input Text",
                        placeholder="Paste your article, paragraph, or any text here...",
                        lines=10
                    )
                    model_selector = gr.Radio(
                        choices=list(MODEL_OPTIONS.keys()),
                        value=list(MODEL_OPTIONS.keys())[0],
                        label="🤖 Choose Model",
                        interactive=True
                    )
                    style_radio = gr.Radio(
                        choices=["Short (50)", "Medium (150)", "Detailed (300)"],
                        value="Medium (150)",
                        label="Summary Style",
                        interactive=True
                    )
                    length_slider = gr.Slider(
                        minimum=50, maximum=300, value=150, step=50,
                        label="Token Length", interactive=True
                    )
                    summarize_btn = gr.Button(
                        "⚡ Summarize Now",
                        elem_id="summarize-btn", variant="primary"
                    )

                with gr.Column():
                    text_output = gr.Textbox(
                        label="📄 Summarized Output",
                        lines=7, interactive=False
                    )
                    with gr.Row():
                        copy_btn  = gr.Button("📋 Copy Output", elem_id="copy-btn",  size="sm")
                        clear_btn = gr.Button("🗑️ Clear All",   elem_id="clear-btn", size="sm")
                    kw_output = gr.Textbox(
                        label="🏷️ Top Keywords",
                        lines=2, interactive=False, elem_id="kw-box"
                    )
                    read_output = gr.Textbox(
                        label="📊 Readability Score",
                        lines=3, interactive=False, elem_id="read-box"
                    )

            style_radio.change(fn=sync_style, inputs=style_radio, outputs=length_slider)
            summarize_btn.click(
                fn=run_summary,
                inputs=[text_input, model_selector, length_slider],
                outputs=[text_output, kw_output, read_output]
            )
            clear_btn.click(fn=clear_text, inputs=[], outputs=[text_input, text_output, kw_output, read_output])
            copy_btn.click(fn=lambda x: x, inputs=[text_output], outputs=[text_output])

        # Tab 2: File Summarizer
        with gr.Tab("📂 File Summarizer"):
            with gr.Row():
                with gr.Column():
                    file_input = gr.File(
                        label="Upload Text File (.txt)",
                        file_types=[".txt"]
                    )
                    file_model = gr.Radio(
                        choices=list(MODEL_OPTIONS.keys()),
                        value=list(MODEL_OPTIONS.keys())[0],
                        label="🤖 Choose Model",
                        interactive=True
                    )
                    file_style = gr.Radio(
                        choices=["Short (50)", "Medium (150)", "Detailed (300)"],
                        value="Medium (150)",
                        label="Summary Style",
                        interactive=True
                    )
                    file_slider = gr.Slider(
                        minimum=50, maximum=300, value=150, step=50,
                        label="Token Length", interactive=True
                    )
                    file_btn = gr.Button(
                        "📂 Upload & Summarize",
                        elem_id="file-summarize-btn", variant="primary"
                    )

                with gr.Column():
                    file_output = gr.Textbox(
                        label="📄 File Summary Output",
                        lines=7, interactive=False
                    )
                    with gr.Row():
                        file_copy_btn  = gr.Button("📋 Copy Output", elem_id="file-copy-btn",  size="sm")
                        file_clear_btn = gr.Button("🗑️ Clear All",   elem_id="file-clear-btn", size="sm")
                    file_kw_output = gr.Textbox(
                        label="🏷️ Top Keywords",
                        lines=2, interactive=False, elem_id="file-kw-box"
                    )
                    file_read_output = gr.Textbox(
                        label="📊 Readability Score",
                        lines=3, interactive=False, elem_id="file-read-box"
                    )

            file_style.change(fn=sync_style, inputs=file_style, outputs=file_slider)
            file_btn.click(
                fn=run_file_summary,
                inputs=[file_input, file_model, file_slider],
                outputs=[file_output, file_kw_output, file_read_output]
            )
            file_clear_btn.click(fn=clear_file, inputs=[], outputs=[file_input, file_output, file_kw_output, file_read_output])
            file_copy_btn.click(fn=lambda x: x, inputs=[file_output], outputs=[file_output])

    gr.HTML("""
    <div style="text-align:center;margin-top:2rem;font-size:11px;color:#8887a0;letter-spacing:1px;">
        QUICKDIGEST &nbsp;·&nbsp; BART · T5 · PEGASUS &nbsp;·&nbsp; BY ABU SHADAB KHAN
    </div>
    """)

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