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

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

def extract_seller(pdf_file):
    with pdfplumber.open(pdf_file) as pdf:
        # Pobranie tekstu z PDF
        full_text = "\n".join(page.extract_text() for page in pdf.pages if page.extract_text())

    # Podział tekstu na krótkie fragmenty (maks. 512 znaków, aby model działał szybciej)
    chunks = [full_text[i:i+512] for i in range(0, len(full_text), 512)]

    seller_tokens = []

    for chunk in chunks:
        entities = extractor(chunk)

        for entity in entities:
            if "ORG" in entity["entity_group"]:  # Szukamy nazw organizacji
                word = entity["word"]

                # Usuwamy błędne tokeny (np. "Faktura", "Z", itp.)
                if not re.match(r"^(Faktura|Z|##|I|KRZYSZTOF|ŻARNOWIECKA)$", word, re.IGNORECASE):
                    seller_tokens.append(word)

        if seller_tokens:  # Jeśli znaleziono organizację, przerywamy pętlę
            break

    # Łączymy tokeny w pełną nazwę organizacji
    seller_name = " ".join(seller_tokens)

    # Usuwamy ewentualne powtórzone litery lub spacje między literami np. "S A" → "S.A."
    seller_name = re.sub(r"\bS A\b", "S.A.", seller_name)

    return {"Sprzedawca": seller_name if seller_name else "Nie znaleziono"}

# Interfejs użytkownika w Hugging Face Spaces
iface = gr.Interface(
    fn=extract_seller,
    inputs=gr.File(label="Wybierz plik PDF"),
    outputs="json",
    title="Ekstrakcja Sprzedawcy z Faktury",
    description="Prześlij plik PDF, aby wydobyć nazwę sprzedawcy."
)

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