import gradio as gr
import pdfplumber
import re
from transformers import pipeline

# Model do rozpoznawania nazw organizacji i wartości numerycznych
extractor = pipeline("ner", model="dbmdz/bert-large-cased-finetuned-conll03-english", aggregation_strategy="simple")

# Reguły do wykrywania NIP, kwot, dat
nip_pattern = re.compile(r'\bPL\s?\d{10}\b|\b\d{10}\b')  # Polski NIP (z "PL" lub bez)
kwota_pattern = re.compile(r'\b\d+[\.,]?\d*\b')  # Kwoty: np. 123.45 lub 123
data_pattern = re.compile(r'\b\d{2}\.\d{2}\.\d{4}\b')  # Daty w formacie DD.MM.YYYY
payment_keywords = ["data płatności", "termin płatności", "zapłata", "płatność"]

def extract_invoice_data(pdf_file):
    with pdfplumber.open(pdf_file) as pdf:
        full_text = "\n".join(page.extract_text() for page in pdf.pages if page.extract_text())

    # Znalezienie nazw organizacji
    entities = extractor(full_text)
    seller_name = []
    
    for entity in entities:
        if "ORG" in entity["entity_group"]:
            seller_name.append(entity["word"])

    # Znajdujemy wartości numeryczne dla NIP, kwot, dat
    seller_nip = nip_pattern.search(full_text)
    kwoty = kwota_pattern.findall(full_text)
    kwoty = [float(k.replace(",", ".")) for k in kwoty if k.replace(",", ".").replace(".", "").isdigit()]
    total_amount = max(kwoty) if kwoty else None

    # Szukamy daty płatności na podstawie kontekstu
    payment_date = None
    for line in full_text.split("\n"):
        if any(keyword in line.lower() for keyword in payment_keywords):
            date_match = data_pattern.search(line)
            if date_match:
                payment_date = date_match.group()
                break

    return {
        "Sprzedawca": " ".join(seller_name) if seller_name else "Nie znaleziono",
        "NIP": seller_nip.group() if seller_nip else "Nie znaleziono",
        "Kwota całkowita": total_amount if total_amount else "Nie znaleziono",
        "Data płatności": payment_date if payment_date else "Nie znaleziono"
    }

# Interfejs użytkownika w Hugging Face Spaces
iface = gr.Interface(
    fn=extract_invoice_data,
    inputs=gr.File(label="Wybierz plik PDF"),
    outputs="json",
    title="Ekstrakcja danych z faktury",
    description="Prześlij plik PDF, a model zwróci dane sprzedawcy, NIP, kwotę i datę płatności."
)

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