import gradio as gr
import fitz
import torch
from transformers import pipeline
import time, io

device = 0 if torch.cuda.is_available() else -1
if device == -1: print("⚠️ No GPU detected. Expect ~10–20s for 300,000 chars on CPU.")

summarizer = pipeline("summarization", model="google/pegasus-xsum", device=device, torch_dtype=torch.int8)

def extract_text(file_bytes):
    if file_bytes[:4].startswith(b'%PDF'):
        doc = fitz.open(stream=file_bytes, filetype="pdf")
        text = "".join(page.get_text("text", flags=16) for page in doc)
        doc.close()
        return text
    try: return file_bytes.decode("utf-8")
    except: return "❌ Unsupported format (PDF/TXT only)"

async def summarize_file(file_bytes):
    start = time.time()
    text = extract_text(file_bytes)[:300000] or "❌ No text found"
    if len(text.strip()) == 0: return text
    chunks = [text[i:i+15000] for i in range(0, len(text), 15000)]
    if not chunks: return "❌ No chunks to summarize"
    summaries = []
    batch_size = 2 if device == -1 else 10  # Smaller batch for CPU
    for i in range(0, len(chunks), batch_size):
        if time.time() - start > 9: 
            summaries.append("⚠️ Stopped early")
            break
        batch = chunks[i:i+batch_size]
        try:
            batch_summaries = summarizer(batch, max_length=40, min_length=10, do_sample=False, batch_size=batch_size)
            summaries.extend(f"**Chunk {i+j+1}**:\n{s['summary_text']}" for j, s in enumerate(batch_summaries))
        except: summaries.append(f"**Chunk {i+1}**: ❌ Error")
    return f"**Chars**: {len(text)}\n**Time**: {time.time()-start:.2f}s\n\n" + "\n\n".join(summaries)

demo = gr.Interface(
    fn=summarize_file, inputs=gr.File(label="📄 PDF/TXT Notes"), 
    outputs=gr.Textbox(label="📝 Summary"), 
    title="Fast Summarizer", description="300,000+ chars in ~5–10s (GPU) or ~10–20s (CPU)"
)

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