import re
import json
import pdfplumber
import pytesseract
from pdf2image import convert_from_path
from google import genai

# ==============================
# HARD-CODED GEMINI KEY
# ==============================

GEMINI_API_KEY = "AIzaSyApRTnUARTDwULx12FPiNf-T-MlrgLbnP0"

client = genai.Client(api_key=GEMINI_API_KEY)

# ==============================
# OCR CONFIG
# ==============================

pytesseract.pytesseract.tesseract_cmd = r"C:\Program Files\Tesseract-OCR\tesseract.exe"

# ==============================
# RAW TEXT EXTRACTION
# ==============================

def extract_raw_text(pdf_path):

    full_text = ""

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

    if not full_text.strip():
        images = convert_from_path(
            pdf_path,
            dpi=300,
            poppler_path=r"C:\poppler\Library\bin"
        )

        for img in images:
            ocr_text = pytesseract.image_to_string(img)
            full_text += ocr_text + "\n"

    return full_text


# ==============================
# FINAL BALANCE DUE
# ==============================

def extract_final_balance_due(raw_text):

    pages = raw_text.strip().split("\f")
    last_page_text = pages[-1] if pages else raw_text

    matches = re.findall(
        r"Balance\s*Due.*?\$?\s*((?:\d[\d,]*)\.\d{2})",
        last_page_text,
        re.IGNORECASE | re.DOTALL
    )

    if matches:
        return matches[-1].replace(",", "")

    return ""


# ==============================
# GEMINI FALLBACK
# ==============================

def extract_with_gemini(raw_text):

    prompt = f"""
You are an AI that reads invoice documents.

Extract:
- invoice_number
- purchase_order
- total_amount_usd (final payable amount only)

Return ONLY valid JSON:

{{
  "invoice_number": "",
  "purchase_order": "",
  "total_amount_usd": ""
}}

Invoice text:
\"\"\"
{raw_text[:10000]}
\"\"\"
"""

    try:
        response = client.models.generate_content(
            model="gemini-2.5-flash",
            contents=prompt
        )

        text = response.text.strip()
        text = re.sub(r"```json|```", "", text).strip()

        data = json.loads(text)

        if isinstance(data, dict):
            return data

        return {}

    except Exception:
        return {}