import re
import json
import sys
import os
import time
import pdfplumber
import pytesseract
import traceback
from pdf2image import convert_from_path
from google import genai

# ==============================
# CONFIGURATION
# ==============================
GEMINI_API_KEY = "AIzaSyDCYE2mBGGW9uHIu3cWkiumAUXvqQvCBzo"
pytesseract.pytesseract.tesseract_cmd = r"C:\Program Files\Tesseract-OCR\tesseract.exe"
POPPLER_PATH = r"C:\poppler\Library\bin"

client = genai.Client(api_key=GEMINI_API_KEY)

SCHEMAS = {
    'Rebar': ["vendor_name", "po", "invoice_date", "invoice_number", "qty", "bol", "branch", "amount"],
    'Non Rebar': ["vendor_name", "po", "invoice_date", "invoice_number", "bol", "branch", "amount"],
    'Inbound': ["vendor_name", "bol_number", "invoice_date", "invoice_number", "ship_to", "amount"],
    'Outbound': ["vendor_name", "load_number_or_bol", "invoice_date", "invoice_number", "ship_from", "amount"],
    'Ap Invoice': ["vendor_name", "invoice_date", "invoice_number", "amount"]
}

def extract_raw_text(pdf_path):
    full_text = ""
    try:
        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=POPPLER_PATH)
            for img in images:
                full_text += pytesseract.image_to_string(img) + "\n"
    except Exception as e:
        raise Exception(f"PDF Extraction failed: {str(e)}")
    return full_text

def extract_with_gemini(raw_text, invoice_type):
    fields = SCHEMAS.get(invoice_type, ["vendor_name", "invoice_date", "invoice_number", "amount"])
    
    prompt = f"""
    ROLE: You are an expert Accounts Payable Specialist specializing in high-accuracy data extraction from diverse invoice formats.

    TASK: Analyze the provided INVOICE TEXT and extract data for the specific fields required for the invoice category: {invoice_type}.

    INVOICE TYPE: {invoice_type}
    
    FIELDS TO EXTRACT:
    - {', '.join(fields)}

    EXTRACTION GUIDELINES:
    1. VENDOR IDENTIFICATION: The 'vendor_name' is typically found at the very top of the document or near the remit-to address. 
    2. FIELD MAPPING:
       - For 'po', if an explicit "P.O. Number" is missing, search for "Purchase Order", "Ref #", "Job #", or "Customer Order #".
       - For 'bol', search for "Bill of Lading", "Shipping Ticket", or "Tracking Number".
       - For 'amount', identify the "Total", "Balance Due", or "Amount Due".
       - For 'invoice_number', search for "Invoice no.:", "Inv #", "Invoice #", or simply a 5-6 digit number located in a "Details" or "Summary" block.
       - For 'Branch' look for branch name on the 'Ship to' address and extract the City or Branch name (e.g., 'Oakland', 'Fresno').
    3. DATA NORMALIZATION:
       - Dates: Convert all dates to MM/DD/YYYY format.
       - Currency: Remove all symbols ($, USD) and thousands-separator commas from numeric values. Return only decimal numbers (e.g., 12,960.00).
       - Quantity: Return as a clean integer or decimal.
    4. NULL VALUES: If a field is not present in the text, return an empty string "" for that key.


    OUTPUT INSTRUCTIONS:
    - Return ONLY a valid JSON object.
    - Do not use markdown blocks (```json).
    - Ensure the keys in your JSON exactly match the REQUIRED FIELDS list provided above.

    INVOICE TEXT:
    \"\"\"
    {raw_text[:12000]}
    \"\"\"
    """

    # --- PROFESSIONAL QUOTA HANDLING (RETRIES) ---
    max_retries = 3
    retry_delay = 10  # Initial wait of 10 seconds

    for attempt in range(max_retries):
        try:
            response = client.models.generate_content(
                model="gemini-3.1-flash-lite-preview", 
                contents=prompt
            )
            
            res_text = response.text.strip()
            res_text = re.sub(r"```json|```", "", res_text).strip()
            
            start = res_text.find('{')
            end = res_text.rfind('}') + 1
            if start != -1 and end != 0:
                return json.loads(res_text[start:end])
                
            raise Exception("Could not find JSON object in AI response")

        except Exception as e:
            err_msg = str(e).upper()
            # If it's a quota error (429 or RESOURCE_EXHAUSTED)
            if "429" in err_msg or "RESOURCE" in err_msg or "QUOTA" in err_msg:
                if attempt < max_retries - 1:
                    # Log to stderr for debugging (won't affect PHP's json_decode)
                    sys.stderr.write(f"Quota hit. Retrying in {retry_delay}s...\n")
                    time.sleep(retry_delay)
                    retry_delay *= 2 
                    continue
                else:
                    raise Exception("AI Quota Exhausted. The system is busy, please try again in 1 minute.")
            
            # For any other error, stop immediately
            raise Exception(f"Extraction Logic Error: {str(e)}")

# ==============================
# MAIN EXECUTION BLOCK
# ==============================
if __name__ == "__main__":
    try:
        if len(sys.argv) < 3:
            raise Exception("Missing CLI arguments (Path and Type)")

        pdf_path, inv_type = sys.argv[1], sys.argv[2]
        
        if not os.path.exists(pdf_path):
            raise Exception(f"File not found at: {pdf_path}")

        raw_data = extract_raw_text(pdf_path)
        if not raw_data.strip():
            raise Exception("No text content found in document")

        final_result = extract_with_gemini(raw_data, inv_type)
        print(json.dumps(final_result))

    except Exception as e:
        # ALWAYS return a JSON object so PHP json_decode works
        print(json.dumps({
            "error": "Python Error",
            "message": str(e)
        }))