import gradio as gr
import pdfplumber
import re
from transformers import LayoutLMForTokenClassification, AutoTokenizer
import torch

model_name = "kryman27/layoutlmv3-finetuned"
model = LayoutLMForTokenClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Wzorce regex dla nowych pól
nip_pattern = re.compile(r'\b(?:PL\s?)?\d{10}\b')
invoice_number_pattern = re.compile(r'Faktura\s*(?:VAT)?\s*(?:nr\.?|#)\s*([\w\-/]+)', re.IGNORECASE)
sale_date_pattern = re.compile(r'Data\s+wystawienia[:\s]*([\d]{2}[.\-/][\d]{2}[.\-/][\d]{4})', re.IGNORECASE)
delivery_date_pattern = re.compile(r'Data\s+dostawy[:\s]*([\d]{2}[.\-/][\d]{2}[.\-/][\d]{4})', re.IGNORECASE)
payment_date_pattern = re.compile(r'(?:Termin\s+płatności|Data\s+płatności)[:\s]*([\d]{2}[.\-/][\d]{2}[.\-/][\d]{4})', re.IGNORECASE)
order_number_pattern = re.compile(r'Zamówienie\s*Nr[:\s]*([\w\-/]+)', re.IGNORECASE)
order_date_pattern = re.compile(r'Data\s+zamówienia[:\s]*([\d]{2}[.\-/][\d]{2}[.\-/][\d]{4})', re.IGNORECASE)
sale_order_pattern = re.compile(r'Zlecenie\s+sprzedaży\s*Nr[:\s]*([\w\-/]+)', re.IGNORECASE)
payment_amount_pattern = re.compile(r'(?:Kwota\s+zapłacona)[:\s]*([\d.,]+)', re.IGNORECASE)
payment_method_pattern = re.compile(r'(?:Forma\s+płatności)[:\s]*([\w/]+)', re.IGNORECASE)

def extract_section(text, section_title):
    pattern = re.compile(rf'{section_title}:(.*?)(?=\n\S|$)', re.IGNORECASE | re.DOTALL)
    match = pattern.search(text)
    return match.group(1).strip() if match else None

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())
    
    # Wyodrębnienie sekcji na podstawie tytułów
    sprzedawca_section = extract_section(full_text, "Sprzedawca")
    nabywca_section = extract_section(full_text, "Nabywca")
    
    sprzedawca = {}
    nabywca = {}
    faktura = {}
    platnosc = {}
    pozycje = []  # Do implementacji ekstrakcji tabelarycznej
    podsumowanie = {}

    # Ekstrakcja danych Sprzedawcy
    if sprzedawca_section:
        lines = sprzedawca_section.splitlines()
        sprzedawca['Nazwa'] = lines[0].strip() if lines else "Nie znaleziono"
        nip_match = nip_pattern.search(sprzedawca_section)
        sprzedawca['NIP'] = nip_match.group() if nip_match else "Nie znaleziono"
        bdo_match = re.search(r'BDO[:\s]*([\w\d]+)', sprzedawca_section, re.IGNORECASE)
        sprzedawca['Numer Rejestracyjny BDO'] = bdo_match.group(1) if bdo_match else "Nie znaleziono"
        sprzedawca['Adres'] = lines[1].strip() if len(lines) > 1 else "Nie znaleziono"
        telefon_match = re.search(r'tel\.?[:\s]*([\+\d\s()-]+)', sprzedawca_section, re.IGNORECASE)
        sprzedawca['Telefon'] = telefon_match.group(1).strip() if telefon_match else "Nie znaleziono"
        fax_match = re.search(r'fax\.?[:\s]*([\+\d\s()-]+)', sprzedawca_section, re.IGNORECASE)
        sprzedawca['Fax'] = fax_match.group(1).strip() if fax_match else "Nie znaleziono"
    else:
        sprzedawca = {
            "Nazwa": "Nie znaleziono",
            "NIP": "Nie znaleziono",
            "Numer Rejestracyjny BDO": "Nie znaleziono",
            "Adres": "Nie znaleziono",
            "Telefon": "Nie znaleziono",
            "Fax": "Nie znaleziono"
        }

    # Ekstrakcja danych Nabywcy
    if nabywca_section:
        lines = nabywca_section.splitlines()
        nabywca['Nazwa'] = lines[0].strip() if lines else "Nie znaleziono"
        nip_match = nip_pattern.search(nabywca_section)
        nabywca['NIP'] = nip_match.group() if nip_match else "Nie podano"
        nabywca['Adres'] = lines[1].strip() if len(lines) > 1 else "Nie znaleziono"
        klient_match = re.search(r'Nr\s+Klienta[:\s]*([\w\d]+)', nabywca_section, re.IGNORECASE)
        nabywca['Nr Klienta'] = klient_match.group(1) if klient_match else "Nie znaleziono"
    else:
        nabywca = {
            "Nazwa": "Nie znaleziono",
            "NIP": "Nie podano",
            "Adres": "Nie znaleziono",
            "Nr Klienta": "Nie znaleziono"
        }

    # Ekstrakcja danych faktury
    invoice_number_match = invoice_number_pattern.search(full_text)
    faktura['Numer'] = invoice_number_match.group(1) if invoice_number_match else "Nie znaleziono"
    sale_date_match = sale_date_pattern.search(full_text)
    faktura['Data Wystawienia'] = sale_date_match.group(1) if sale_date_match else "Nie znaleziono"
    delivery_date_match = delivery_date_pattern.search(full_text)
    faktura['Data Dostawy'] = delivery_date_match.group(1) if delivery_date_match else "Nie znaleziono"
    order_number_match = order_number_pattern.search(full_text)
    faktura['Zamówienie Nr'] = order_number_match.group(1) if order_number_match else "Nie znaleziono"
    order_date_match = order_date_pattern.search(full_text)
    faktura['Data Zamówienia'] = order_date_match.group(1) if order_date_match else "Nie znaleziono"
    sale_order_match = sale_order_pattern.search(full_text)
    faktura['Zlecenie Sprzedaży Nr'] = sale_order_match.group(1) if sale_order_match else "Nie znaleziono"

    # Ekstrakcja danych płatności
    payment_date_match = payment_date_pattern.search(full_text)
    platnosc['Termin Zapłaty'] = payment_date_match.group(1) if payment_date_match else "Nie znaleziono"
    payment_method_match = payment_method_pattern.search(full_text)
    platnosc['Forma Zapłaty'] = payment_method_match.group(1) if payment_method_match else "Nie znaleziono"
    payment_amount_match = payment_amount_pattern.search(full_text)
    platnosc['Kwota Zapłacona'] = float(payment_amount_match.group(1).replace(',', '.')) if payment_amount_match else "Nie znaleziono"

    # Ekstrakcja podsumowania (przykładowo)
    podsumowanie_match = re.search(r'Razem[:\s]*([\d.,]+)', full_text)
    podsumowanie['Suma Brutto'] = float(podsumowanie_match.group(1).replace(',', '.')) if podsumowanie_match else "Nie znaleziono"

    result = {
        "Sprzedawca": sprzedawca,
        "Nabywca": nabywca,
        "Faktura": faktura,
        "Płatność": platnosc,
        "Pozycje": pozycje,
        "Podsumowanie": podsumowanie
    }
    return result

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 narzędzie zwróci szczegółowe dane faktury."
)

#UI start
if __name__ == "__main__":
    iface.launch()