import gradio as gr
import pdfplumber
import pandas as pd
import re


def extract_invoices(file):

    text = ""

    with pdfplumber.open(file.name) as pdf:
        for page in pdf.pages:
            t = page.extract_text()
            if t:
                text += t + "\n"

    # ---------------- INVOICE NUMBER ---------------- #
    invoice_pattern = r'(?:Invoice\s*No|Invoice\s*Number|InvoiceNo|Inv\.?\s*No|Invoice\s*#|Bill\s*No|Bill\s*Number|Bill\s*of\s*Supply\s*No|Bill\s*of\s*Supply\s*Number)\s*[:#\-]?\s*([A-Za-z0-9\-\/]+)'

    invoice_matches = list(re.finditer(invoice_pattern, text, re.IGNORECASE))

    rows = []
    processed_invoices = set()

    # ---------------- GSTIN PATTERN ---------------- #
    gstin_pattern = r'\b\d{2}[A-Z]{5}\d{4}[A-Z][A-Z\d]Z[A-Z\d]\b'

    for i, match in enumerate(invoice_matches):

        invoice_number = match.group(1)

        if not re.search(r'\d', invoice_number) or len(invoice_number) < 4:
            continue

        if invoice_number in processed_invoices:
            continue

        processed_invoices.add(invoice_number)

        start = match.start()

        if i + 1 < len(invoice_matches):
            end = invoice_matches[i + 1].start()
        else:
            end = len(text)

        block = text[start:end]

        # ---------------- DATE EXTRACTION ---------------- #
        date_pattern = r'(?:\d{1,2}[\/\-.]\d{1,2}[\/\-.]\d{2,4}|\d{4}[\/\-.]\d{1,2}[\/\-.]\d{1,2}|(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)[a-z]*\s+\d{1,2},?\s*\d{2,4}|\d{1,2}\s+(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)[a-z]*\s*\d{2,4})'

        date_match = re.search(
            r'(Invoice\s*Date|Bill\s*Date|Date\s*of\s*Invoice)\s*[:\-]?\s*\n?\s*(' + date_pattern + ')',
            block,
            re.IGNORECASE
        )

        if date_match:
            invoice_date = date_match.group(2)
        else:
            fallback_date = re.search(date_pattern, block, re.IGNORECASE)
            invoice_date = fallback_date.group(0) if fallback_date else ""

        # ---------------- GSTIN EXTRACTION (BEFORE + AFTER) ---------------- #

        # Look BEFORE invoice (important fix)
        pre_start = max(0, start - 500)
        pre_block = text[pre_start:start]

        # Combine both
        search_area = pre_block + block

        # Try keyword-based extraction
        gstin_match = re.search(
            r'(GSTIN|GST\s*No|Supplier\s*GSTIN|Vendor\s*GSTIN)\s*[:\-]?\s*(' + gstin_pattern + ')',
            search_area,
            re.IGNORECASE
        )

        if gstin_match:
            gstin = gstin_match.group(2)
        else:
            # fallback: get all GSTINs
            gstin_list = re.findall(gstin_pattern, search_area, re.IGNORECASE)

            if len(gstin_list) > 0:
                gstin = gstin_list[0]   # usually seller GSTIN
            else:
                gstin = ""

        # ---------------- FLOAT DETECTION ---------------- #
        float_pattern = r'(?:\d{1,3}(?:,\d{3})+|\d{1,2}(?:,\d{2})+(?:,\d{3}))\.\d{2}|\d+\.\d{2}'

        float_numbers = re.findall(float_pattern, block)
        float_numbers = [float(x.replace(",", "")) for x in float_numbers]

        total_amount = max(float_numbers) if float_numbers else ""

        # ---------------- ROW ---------------- #
        row = {
            "Invoice Number": invoice_number,
            "Invoice Date": invoice_date,
            "GSTIN": gstin,
            "Narration": "Invoice processed",
            "Debit Account": "Expense",
            "Credit Account": "Vendor",
            "Total Amount": total_amount
        }

        rows.append(row)

    df = pd.DataFrame(rows)

    excel_file = "invoice_output.xlsx"
    df.to_excel(excel_file, index=False)

    return df, excel_file


# ---------------- GRADIO UI ---------------- #

with gr.Blocks() as demo:

    gr.Markdown("# Multiple Invoice PDF → Excel Extractor")

    file_input = gr.File(label="Upload Invoice PDF")

    extract_btn = gr.Button("Extract Invoices")

    output_table = gr.Dataframe(label="Extracted Invoice Data")

    download_file = gr.File(label="Download Excel")

    extract_btn.click(
        extract_invoices,
        inputs=file_input,
        outputs=[output_table, download_file]
    )

demo.launch()