from fastapi import FastAPI, UploadFile, File
from fastapi.responses import RedirectResponse
import gradio as gr
from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
import tempfile
import os
from PIL import Image
import fitz  # PyMuPDF
import docx
import easyocr

app = FastAPI()

# Lightweight model choices
SUMMARIZATION_MODEL = "facebook/bart-large-cnn"  # 500MB
IMAGE_CAPTIONING_MODEL = "Salesforce/blip-image-captioning-base"  # 300MB

# Initialize models
try:
    summarizer = pipeline(
        "summarization",
        model=SUMMARIZATION_MODEL,
        device="cpu"
    )
except Exception as e:
    print(f"Error loading summarizer: {e}")
    summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")  # Fallback 250MB model

captioner = pipeline(
    "image-to-text",
    model=IMAGE_CAPTIONING_MODEL,
    device="cpu"
)

reader = easyocr.Reader(['en'])  # Lightweight OCR

def extract_text_from_file(file_path: str, file_type: str):
    """Extract text from different document formats"""
    try:
        if file_type == "pdf":
            with fitz.open(file_path) as doc:
                return "\n".join(page.get_text() for page in doc)
        elif file_type == "docx":
            doc = docx.Document(file_path)
            return "\n".join(p.text for p in doc.paragraphs)
        else:
            return "Unsupported file format (only PDF/DOCX supported in lightweight version)"
    except Exception as e:
        return f"Error reading file: {str(e)}"

def process_document(file):
    """Handle document summarization"""
    try:
        file_ext = os.path.splitext(file.name)[1][1:].lower()
        if file_ext not in ["pdf", "docx"]:
            return "Lightweight version only supports PDF and DOCX"
            
        with tempfile.NamedTemporaryFile(delete=False, suffix=f".{file_ext}") as tmp:
            tmp.write(file.read())
            tmp_path = tmp.name
        
        text = extract_text_from_file(tmp_path, file_ext)
        summary = summarizer(text, max_length=130, min_length=30, do_sample=False)[0]['summary_text']
        
        os.unlink(tmp_path)
        return summary
    except Exception as e:
        return f"Processing error: {str(e)}"

def process_image(image):
    """Handle image captioning and OCR"""
    try:
        img = Image.open(image)
        
        # Get caption
        caption = captioner(img)[0]['generated_text']
        
        # Get OCR text
        ocr_result = reader.readtext(img)
        ocr_text = " ".join([res[1] for res in ocr_result])
        
        return {
            "caption": caption,
            "ocr_text": ocr_text if ocr_text else "No readable text found"
        }
    except Exception as e:
        return {"error": str(e)}

# Gradio Interface
with gr.Blocks(title="Lightweight Document & Image Analysis") as demo:
    gr.Markdown("## 📄 Lightweight Document & Image Analysis")
    
    with gr.Tab("Document Summarization"):
        gr.Markdown("Supports PDF and DOCX files (max 10MB)")
        doc_input = gr.File(label="Upload Document", file_types=[".pdf", ".docx"])
        doc_output = gr.Textbox(label="Summary")
        doc_button = gr.Button("Summarize")
        
    with gr.Tab("Image Analysis"):
        gr.Markdown("Get captions and extracted text from images")
        img_input = gr.Image(type="filepath", label="Upload Image")
        with gr.Accordion("Results", open=False):
            caption_output = gr.Textbox(label="Image Caption")
            ocr_output = gr.Textbox(label="Extracted Text")
        img_button = gr.Button("Analyze")
    
    doc_button.click(process_document, inputs=doc_input, outputs=doc_output)
    img_button.click(process_image, inputs=img_input, outputs=[caption_output, ocr_output])

# Mount Gradio app
app = gr.mount_gradio_app(app, demo, path="/")

@app.get("/")
def redirect_to_interface():
    return RedirectResponse(url="/")