import os
import re
from functools import lru_cache

import gradio as gr
import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer


DEFAULT_MODEL = os.getenv("MODEL_ID", "sshleifer/distilbart-cnn-12-6")
MAX_INPUT_TOKENS = int(os.getenv("MAX_INPUT_TOKENS", "900"))


@lru_cache(maxsize=1)
def load_model():
    tokenizer = AutoTokenizer.from_pretrained(DEFAULT_MODEL)
    model = AutoModelForSeq2SeqLM.from_pretrained(DEFAULT_MODEL)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.eval()
    return tokenizer, model, device


def clean_text(text: str) -> str:
    return re.sub(r"\s+", " ", (text or "")).strip()


def split_into_chunks(text: str, tokenizer, max_tokens: int = MAX_INPUT_TOKENS):
    sentences = re.split(r"(?<=[.!?])\s+", clean_text(text))
    chunks, current = [], []

    for sentence in sentences:
        candidate = " ".join(current + [sentence]).strip()
        if current and len(tokenizer.encode(candidate, add_special_tokens=True)) > max_tokens:
            chunks.append(" ".join(current))
            current = [sentence]
        else:
            current.append(sentence)

    if current:
        chunks.append(" ".join(current))

    # Handle a single sentence that is longer than the model context.
    safe_chunks = []
    for chunk in chunks:
        token_ids = tokenizer.encode(chunk, add_special_tokens=False)
        for start in range(0, len(token_ids), max_tokens):
            safe_chunks.append(
                tokenizer.decode(token_ids[start : start + max_tokens], skip_special_tokens=True)
            )
    return [chunk for chunk in safe_chunks if chunk.strip()]


def generate_summary(text: str, length: str, progress=gr.Progress()):
    text = clean_text(text)
    if len(text) < 80:
        raise gr.Error("Please paste an article with at least 80 characters.")

    tokenizer, model, device = load_model()
    chunks = split_into_chunks(text, tokenizer)
    settings = {
        "Short": (35, 90),
        "Medium": (60, 150),
        "Detailed": (90, 220),
    }
    min_length, max_length = settings[length]

    summaries = []
    for index, chunk in enumerate(chunks):
        progress((index + 1) / len(chunks), desc=f"Summarizing section {index + 1}/{len(chunks)}")
        inputs = tokenizer(
            chunk,
            return_tensors="pt",
            truncation=True,
            max_length=MAX_INPUT_TOKENS,
        ).to(device)
        with torch.inference_mode():
            output = model.generate(
                **inputs,
                num_beams=4,
                length_penalty=1.8,
                no_repeat_ngram_size=3,
                min_length=min(min_length, max(12, len(inputs["input_ids"][0]) // 6)),
                max_length=max_length,
                early_stopping=True,
            )
        summaries.append(tokenizer.decode(output[0], skip_special_tokens=True))

    combined = " ".join(summaries)
    # A second pass makes multi-chunk article summaries read as one coherent abstract.
    if len(chunks) > 1 and len(tokenizer.encode(combined)) > max_length:
        inputs = tokenizer(
            combined,
            return_tensors="pt",
            truncation=True,
            max_length=MAX_INPUT_TOKENS,
        ).to(device)
        with torch.inference_mode():
            output = model.generate(
                **inputs,
                num_beams=4,
                length_penalty=1.5,
                no_repeat_ngram_size=3,
                min_length=min_length,
                max_length=max_length,
                early_stopping=True,
            )
        combined = tokenizer.decode(output[0], skip_special_tokens=True)

    return combined, f"{len(text.split()):,} words → {len(combined.split()):,} words"


EXAMPLE = (
    "Artificial intelligence is increasingly used in healthcare to help clinicians "
    "analyze medical images, predict patient risks, and reduce administrative work. "
    "Researchers say these systems can improve speed and consistency, but they also "
    "warn that models must be tested across diverse patient groups. Hospitals need "
    "strong privacy controls, human oversight, and clear processes for correcting "
    "errors. Regulators are developing standards intended to make clinical AI safer "
    "and more transparent while preserving room for useful innovation."
)

with gr.Blocks(theme=gr.themes.Soft(), title="Article Summarizer") as demo:
    gr.Markdown(
        """
        # AI Article Summarizer
        Paste a long article and generate a concise, readable abstractive summary.
        Powered by a BART model fine-tuned on CNN/DailyMail.
        """
    )
    with gr.Row():
        with gr.Column(scale=3):
            article = gr.Textbox(
                label="Article",
                placeholder="Paste an article here…",
                lines=16,
                value=EXAMPLE,
            )
            length = gr.Radio(
                ["Short", "Medium", "Detailed"],
                value="Medium",
                label="Summary length",
            )
            summarize = gr.Button("Generate summary", variant="primary")
        with gr.Column(scale=2):
            summary = gr.Textbox(label="Generated summary", lines=12, show_copy_button=True)
            stats = gr.Textbox(label="Compression", interactive=False)

    summarize.click(generate_summary, [article, length], [summary, stats])
    article.submit(generate_summary, [article, length], [summary, stats])

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