import gradio as gr
from transformers import DonutProcessor, VisionEncoderDecoderModel
import torch
from PIL import Image
import re
import sentencepiece

# Carica il modello con parametri specifici
processor = DonutProcessor.from_pretrained(
    "naver-clova-ix/donut-base",
    use_fast=True
)
model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base")
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
model.eval()

def estrai_titolo_autore(testo):
    titolo = None
    autore = None
    
    # Cerca titolo
    titolo_match = re.search(r"Title[:\-–]\s*(.*?)(?:,|$)", testo, re.IGNORECASE)
    if titolo_match:
        titolo = titolo_match.group(1).strip()
    
    # Cerca autore
    autore_match = re.search(r"Author[:\-–]\s*(.*?)(?:,|$)", testo, re.IGNORECASE)
    if autore_match:
        autore = autore_match.group(1).strip()
    
    # Fallback: prendi le prime due righe se non troviamo match
    if not titolo or not autore:
        linee = testo.split("\n")
        if len(linee) >= 2:
            titolo = titolo or linee[0].strip()
            autore = autore or linee[1].strip()
    
    return titolo or "", autore or ""

def analizza_copertina(img):
    try:
        # Converti l'immagine in RGB
        image = img.convert("RGB")
        
        # Prepara il prompt
        prompt = "<s_docvqa><s_question>Extract book title and author from this cover.</s_question><s_answer>"
        
        # Prepara l'input
        pixel_values = processor(image, return_tensors="pt").pixel_values.to(device)
        decoder_input_ids = processor.tokenizer(prompt, add_special_tokens=False, return_tensors="pt")["input_ids"].to(device)
        
        # Genera la risposta con parametri aggiornati
        outputs = model.generate(
            pixel_values,
            decoder_input_ids=decoder_input_ids,
            max_new_tokens=100,  # Aumentiamo il numero di token generati
            min_length=1,        # Lunghezza minima
            num_beams=4,         # Aumentiamo il beam search per risultati migliori
            early_stopping=True,
            pad_token_id=processor.tokenizer.pad_token_id,
            eos_token_id=processor.tokenizer.eos_token_id,
            use_cache=True,
            no_repeat_ngram_size=3  # Evita ripetizioni
        )
        
        # Decodifica il risultato
        testo = processor.batch_decode(outputs, skip_special_tokens=True)[0]
        
        # Estrai titolo e autore
        titolo, autore = estrai_titolo_autore(testo)
        
        return {
            "titolo": titolo,
            "autore": autore,
            "testo_completo": testo
        }
    except Exception as e:
        return {
            "error": str(e),
            "titolo": "",
            "autore": "",
            "testo_completo": ""
        }

# Crea l'interfaccia Gradio
iface = gr.Interface(
    fn=analizza_copertina,
    inputs=gr.Image(type="pil"),
    outputs=gr.JSON(),
    title="Estrai Titolo e Autore da una Copertina",
    description="Carica un'immagine di copertina per ottenere titolo e autore del libro."
)

# Lancia l'applicazione
iface.launch(server_name="0.0.0.0", share=True)