app.py

import os
import io
import re
import base64
import gc
import tempfile
import uuid
import asyncio
from typing import List, Dict, Optional, Tuple
from collections import Counter
from concurrent.futures import ThreadPoolExecutor

from fastapi import FastAPI, File, UploadFile, Form, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from starlette.requests import Request
import fitz  # PyMuPDF
import google.generativeai as genai
from PIL import Image
from fastapi import Query


# Azure Blob Storage
try:
    from azure.storage.blob import (
        BlobServiceClient,
        generate_blob_sas,
        BlobSasPermissions,
        ContentSettings
    )
    AZURE_AVAILABLE = True
except ImportError:
    AZURE_AVAILABLE = False
    print("Warning: azure-storage-blob not installed. Run: pip install azure-storage-blob")

# Google Gemini - optional import
try:

    GEMINI_AVAILABLE = True
except ImportError:
    GEMINI_AVAILABLE = False
    print("Warning: google-generativeai not installed. Image-based PDFs won't be supported.")

from datetime import datetime, timedelta

app = FastAPI(title="Invoice Splitter API with Azure Blob Storage - Optimized")

# Increase request body size limit
Request.max_body_size = 200 * 1024 * 1024  # 200MB

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ============================================================================
# ⭐ CONFIGURATION FROM ENVIRONMENT VARIABLES (Hugging Face Secrets)
# ============================================================================

# Gemini API Key - REQUIRED for image-based PDFs
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")

# Azure Blob Storage Configuration - REQUIRED for blob storage
AZURE_STORAGE_CONNECTION_STRING = os.environ.get(
    "AZURE_STORAGE_CONNECTION_STRING", "")
AZURE_STORAGE_ACCOUNT_NAME = os.environ.get("AZURE_STORAGE_ACCOUNT_NAME", "")
AZURE_STORAGE_ACCOUNT_KEY = os.environ.get("AZURE_STORAGE_ACCOUNT_KEY", "")

# Container name - can be configured or use default
AZURE_CONTAINER_NAME = os.environ.get("AZURE_CONTAINER_NAME", "invoice-splits")

# ⭐ FOLDER STRUCTURE CONFIGURATION
ROOT_FOLDER = os.environ.get("ROOT_FOLDER", "POD")  # Root folder name

# ⭐ PERFORMANCE CONFIGURATION
MAX_PARALLEL_GEMINI_CALLS = int(
    os.environ.get("MAX_PARALLEL_GEMINI_CALLS", "5"))
GEMINI_IMAGE_RESOLUTION = float(
    os.environ.get("GEMINI_IMAGE_RESOLUTION", "1.2"))
USE_SMART_SAMPLING = os.environ.get(
    "USE_SMART_SAMPLING", "false").lower() == "true"

# ⭐ SERVER CONFIGURATION
HOST = os.environ.get("HOST", "0.0.0.0")  # Hugging Face uses 0.0.0.0
PORT = int(os.environ.get("PORT", "7860"))  # Hugging Face default port

# ============================================================================
# GLOBAL VARIABLES
# ============================================================================

gemini_model = None
blob_service_client = None

# ============================================================================
# STARTUP VALIDATION
# ============================================================================


def validate_configuration():
    """Validate configuration and warn about missing credentials."""
    warnings = []
    errors = []

    # Check Gemini API Key
    if not GEMINI_API_KEY:
        warnings.append(
            "⚠️  GEMINI_API_KEY not set - image-based PDFs will not work")
    else:
        print(f"✅ GEMINI_API_KEY configured ({len(GEMINI_API_KEY)} chars)")

    # Check Azure credentials
    if not AZURE_STORAGE_CONNECTION_STRING:
        if not (AZURE_STORAGE_ACCOUNT_NAME and AZURE_STORAGE_ACCOUNT_KEY):
            errors.append(
                "❌ Azure credentials missing - set AZURE_STORAGE_CONNECTION_STRING or both AZURE_STORAGE_ACCOUNT_NAME and AZURE_STORAGE_ACCOUNT_KEY")
        else:
            print(
                f"✅ Azure credentials configured (account: {AZURE_STORAGE_ACCOUNT_NAME})")
    else:
        print(f"✅ Azure connection string configured")

    # Print all warnings
    for warning in warnings:
        print(warning)

    # Print all errors
    for error in errors:
        print(error)

    if errors:
        print("\n⚠️  WARNING: Some required credentials are missing!")
        print("   Set them in Hugging Face Spaces Settings > Repository secrets")

    return len(errors) == 0


# ============================================================================
# AZURE BLOB STORAGE FUNCTIONS
# ============================================================================


def get_blob_service_client():
    """Get or create Azure Blob Service Client."""
    global blob_service_client

    if not AZURE_AVAILABLE:
        print("❌ Azure SDK not available")
        return None

    if blob_service_client is None:
        try:
            if AZURE_STORAGE_CONNECTION_STRING:
                blob_service_client = BlobServiceClient.from_connection_string(
                    AZURE_STORAGE_CONNECTION_STRING
                )
                print("✅ Azure Blob Storage initialized with connection string")
            elif AZURE_STORAGE_ACCOUNT_NAME and AZURE_STORAGE_ACCOUNT_KEY:
                account_url = f"https://{AZURE_STORAGE_ACCOUNT_NAME}.blob.core.windows.net"
                blob_service_client = BlobServiceClient(
                    account_url=account_url,
                    credential=AZURE_STORAGE_ACCOUNT_KEY
                )
                print("✅ Azure Blob Storage initialized with account key")
            else:
                print("⚠️ WARNING: No Azure credentials configured")
                return None
        except Exception as e:
            print(f"❌ Failed to initialize Azure Blob Storage: {e}")
            return None

    return blob_service_client


def ensure_container_exists(container_name: str = None):
    """Create container if it doesn't exist."""
    if container_name is None:
        container_name = AZURE_CONTAINER_NAME

    try:
        client = get_blob_service_client()
        if client:
            container_client = client.get_container_client(container_name)
            if not container_client.exists():
                container_client.create_container()
                print(f"✅ Created container: {container_name}")
            else:
                print(f"✅ Container exists: {container_name}")
    except Exception as e:
        print(f"⚠️ Container check error: {e}")


def upload_raw_pdf_to_blob(
    pdf_bytes: bytes,
    filename: str,
    batch_id: str,
    container_name: str = None
) -> dict:
    """
    Upload original/raw PDF to Azure Blob Storage.

    Path structure: POD/{batch_id}/{filename}/Raw/{filename}
    """
    if container_name is None:
        container_name = AZURE_CONTAINER_NAME

    try:
        client = get_blob_service_client()
        if not client:
            raise HTTPException(
                status_code=500,
                detail="Azure Blob Storage not configured"
            )

        # Clean filename for folder name
        base_filename = os.path.splitext(filename)[0]
        safe_folder_name = re.sub(r'[<>:"/\\|?*]', '_', base_filename)

        blob_name = f"{ROOT_FOLDER}/{batch_id}/{safe_folder_name}/Raw/{filename}"

        # Get blob client
        blob_client = client.get_blob_client(
            container=container_name,
            blob=blob_name
        )

        # Upload PDF
        print(f"📤 Uploading raw PDF to: {blob_name}")
        blob_client.upload_blob(
            pdf_bytes,
            overwrite=True,
            content_settings=ContentSettings(content_type='application/pdf'),
            metadata={
                'batch_id': batch_id,
                'file_type': 'raw',
                'uploaded_at': datetime.now().isoformat(),
                'original_filename': filename
            }
        )

        # Generate SAS URL (valid for 24 hours)
        expiry_hours = 24
        sas_token = generate_blob_sas(
            account_name=AZURE_STORAGE_ACCOUNT_NAME,
            container_name=container_name,
            blob_name=blob_name,
            account_key=AZURE_STORAGE_ACCOUNT_KEY,
            permission=BlobSasPermissions(read=True),
            expiry=datetime.utcnow() + timedelta(hours=expiry_hours)
        )

        # Construct URLs
        blob_url = blob_client.url
        download_url = f"{blob_url}?{sas_token}"
        expires_at = (datetime.utcnow() +
                      timedelta(hours=expiry_hours)).isoformat() + "Z"

        print(f"✅ Uploaded raw PDF: {blob_name}")

        return {
            "blob_name": blob_name,
            "blob_url": blob_url,
            "download_url": download_url,
            "expires_at": expires_at,
            "expires_in_hours": expiry_hours,
            "storage": "azure_blob",
            "folder_type": "raw",
            "container": container_name,
            "size_bytes": len(pdf_bytes),
            "size_mb": round(len(pdf_bytes) / (1024 * 1024), 2)
        }

    except Exception as e:
        print(f"❌ Raw PDF upload failed: {e}")
        raise HTTPException(
            status_code=500,
            detail=f"Azure Blob upload failed: {str(e)}"
        )


def upload_split_pdf_to_blob(
    pdf_bytes: bytes,
    invoice_filename: str,
    original_filename: str,
    batch_id: str,
    container_name: str = None
) -> dict:
    """
    Upload split invoice PDF to Azure Blob Storage.

    Path structure: POD/{batch_id}/{original_filename}/Splitted/{invoice_filename}
    """
    if container_name is None:
        container_name = AZURE_CONTAINER_NAME

    try:
        client = get_blob_service_client()
        if not client:
            raise HTTPException(
                status_code=500,
                detail="Azure Blob Storage not configured"
            )

        # Clean original filename for folder name
        base_filename = os.path.splitext(original_filename)[0]
        safe_folder_name = re.sub(r'[<>:"/\\|?*]', '_', base_filename)

        blob_name = f"{ROOT_FOLDER}/{batch_id}/{safe_folder_name}/Splitted/{invoice_filename}"

        # Get blob client
        blob_client = client.get_blob_client(
            container=container_name,
            blob=blob_name
        )

        # Upload PDF
        blob_client.upload_blob(
            pdf_bytes,
            overwrite=True,
            content_settings=ContentSettings(content_type='application/pdf'),
            metadata={
                'batch_id': batch_id,
                'file_type': 'split',
                'uploaded_at': datetime.now().isoformat(),
                'original_filename': original_filename,
                'invoice_filename': invoice_filename
            }
        )

        # Generate SAS URL (valid for 24 hours)
        expiry_hours = 24
        sas_token = generate_blob_sas(
            account_name=AZURE_STORAGE_ACCOUNT_NAME,
            container_name=container_name,
            blob_name=blob_name,
            account_key=AZURE_STORAGE_ACCOUNT_KEY,
            permission=BlobSasPermissions(read=True),
            expiry=datetime.utcnow() + timedelta(hours=expiry_hours)
        )

        # Construct URLs
        blob_url = blob_client.url
        download_url = f"{blob_url}?{sas_token}"
        expires_at = (datetime.utcnow() +
                      timedelta(hours=expiry_hours)).isoformat() + "Z"

        return {
            "blob_name": blob_name,
            "blob_url": blob_url,
            "download_url": download_url,
            "expires_at": expires_at,
            "expires_in_hours": expiry_hours,
            "storage": "azure_blob",
            "folder_type": "split",
            "container": container_name,
            "size_bytes": len(pdf_bytes),
            "size_mb": round(len(pdf_bytes) / (1024 * 1024), 2)
        }

    except Exception as e:
        print(f"❌ Split PDF upload failed: {e}")
        raise HTTPException(
            status_code=500,
            detail=f"Azure Blob upload failed: {str(e)}"
        )


async def cleanup_old_blobs(batch_id: str, container_name: str = None):
    """Delete all blobs for a specific batch_id."""
    if container_name is None:
        container_name = AZURE_CONTAINER_NAME

    try:
        client = get_blob_service_client()
        if not client:
            return

        container_client = client.get_container_client(container_name)

        prefix = f"{ROOT_FOLDER}/{batch_id}/"
        blobs = container_client.list_blobs(name_starts_with=prefix)

        deleted_count = 0
        for blob in blobs:
            blob_client = container_client.get_blob_client(blob.name)
            blob_client.delete_blob()
            deleted_count += 1

        print(f"🧹 Cleaned up {deleted_count} blobs for batch {batch_id}")

    except Exception as e:
        print(f"⚠️ Cleanup error: {e}")


# ============================================================================
# OPTIMIZED GEMINI FUNCTIONS WITH ASYNC PROCESSING
# ============================================================================

def get_gemini_model():
    """Get or create Gemini model instance."""
    global gemini_model

    if not GEMINI_AVAILABLE:
        return None

    if gemini_model is None:
        if not GEMINI_API_KEY:
            return None

        try:
            genai.configure(api_key=GEMINI_API_KEY)
            # Use Gemini 2.5 Flash
            gemini_model = genai.GenerativeModel('gemini-2.5-flash')
            print("✅ Google Gemini 2.5 Flash initialized")
        except Exception as e:
            print(f"❌ Failed to initialize Gemini: {e}")
            return None

    return gemini_model


def _clean_gemini_invoice_text(text: str) -> Optional[str]:
    if not text:
        return None

    cleaned = text.strip()
    cleaned = cleaned.replace("*", "").replace("#", "")
    cleaned = re.sub(
        r"(?i)\b(invoice|inv|bill|document|doc|tax\s*invoice)\s*(no|number)?\b",
        "",
        cleaned
    )
    cleaned = re.sub(r"[:\-\(\)\[\]]", " ", cleaned)
    cleaned = re.sub(r"\s+", " ", cleaned).strip()

    # Extract candidates
    tokens = re.findall(r"[A-Z0-9][A-Z0-9\-\/]{2,}", cleaned.upper())

    # Prefer alphanumeric invoice IDs first
    for token in tokens:
        if any(c.isalpha() for c in token) and any(c.isdigit() for c in token):
            return token

    # Fallback to numeric-only (6-15 digits)
    for token in tokens:
        if token.isdigit() and 6 <= len(token) <= 15:
            return token

    return None


def extract_invoice_gemini_sync(page: fitz.Page) -> Optional[str]:
    """
    Optimized synchronous Gemini extraction for thread pool execution.
    - Reduced image resolution for faster processing
    - Simplified prompt for quicker responses
    - OCR fallback for better accuracy
    """
    model = get_gemini_model()
    if not model:
        return None

    img = None
    try:
        # Reduced resolution for faster processing
        pix = page.get_pixmap(matrix=fitz.Matrix(
            GEMINI_IMAGE_RESOLUTION, GEMINI_IMAGE_RESOLUTION))
        img_bytes = pix.tobytes("png")
        pix = None
        img = Image.open(io.BytesIO(img_bytes))

        # Updated prompt to prioritize labeled alphanumeric invoice numbers
        prompt = """Extract ONLY the invoice number from this image.
Prefer the value next to labels like: Invoice No, Invoice Number, Bill No, Document No.
Return ONLY the identifier (keep letters, e.g., A07966). If not found, return: NONE."""

        response = model.generate_content([prompt, img])
        if response and response.text:
            extracted_text = response.text.strip()
            candidate = _clean_gemini_invoice_text(extracted_text)
            if candidate and len(candidate) > 2:
                img.close()
                return candidate

        # OCR Fallback: Extract full text then run regex
        ocr_prompt = "Extract all text from this invoice image. Return the complete text content."
        ocr_response = model.generate_content([ocr_prompt, img])
        if ocr_response and ocr_response.text:
            inv = try_extract_invoice_from_text(ocr_response.text)
            if inv:
                img.close()
                return inv

        if img:
            img.close()
        return None

    except Exception as e:
        print(f"Gemini error: {e}")
        if img:
            img.close()
        return None


async def extract_invoices_batch_async(
    doc: fitz.Document,
    is_image_pdf: bool,
    batch_size: int = MAX_PARALLEL_GEMINI_CALLS
) -> List[Optional[str]]:
    """
    🚀 OPTIMIZED: Extract invoice numbers with parallel processing.

    For text PDFs: Fast sequential processing
    For image PDFs: Parallel Gemini API calls (5-10x faster)
    """
    page_invoice_nos = []

    if not is_image_pdf:
        # Fast text-based extraction (no parallelization needed)
        print(f"  📝 Text-based extraction (sequential)")
        for i in range(doc.page_count):
            if i % 50 == 0:
                print(f"  Extracting... Page {i+1}/{doc.page_count}")
            page = doc.load_page(i)
            inv = extract_invoice_text_based(page)
            page_invoice_nos.append(inv)
            page = None
            if i % 100 == 0:
                gc.collect()
        return page_invoice_nos

    # Image-based PDF: Use parallel Gemini processing
    print(f"  🚀 Image-based extraction (parallel, batch_size={batch_size})")

    # Use ThreadPoolExecutor for parallel API calls
    with ThreadPoolExecutor(max_workers=batch_size) as executor:
        futures = []

        # Submit all pages to thread pool
        for i in range(doc.page_count):
            page = doc.load_page(i)
            # First try text extraction (fast)
            text_result = extract_invoice_text_based(page)
            if text_result:
                futures.append((i, None, text_result))
            else:
                # Submit to Gemini thread pool
                future = executor.submit(extract_invoice_gemini_sync, page)
                futures.append((i, future, None))

        # Collect results in order
        page_invoice_nos = [None] * doc.page_count
        completed = 0

        for i, future, text_result in futures:
            try:
                if text_result:
                    # Already extracted from text
                    page_invoice_nos[i] = text_result
                    completed += 1
                else:
                    # Wait for Gemini result
                    result = future.result(timeout=30)
                    page_invoice_nos[i] = result
                    completed += 1

                if completed % 5 == 0:
                    print(
                        f"  ✓ Processed {completed}/{doc.page_count} pages...")

            except Exception as e:
                print(f"  ⚠️ Page {i+1} failed: {e}")
                page_invoice_nos[i] = None

            if completed % 20 == 0:
                gc.collect()

    print(f"  ✅ Extraction complete: {completed}/{doc.page_count} pages")
    return page_invoice_nos


def extract_invoices_smart_sampling(doc: fitz.Document, is_image_pdf: bool) -> List[Optional[str]]:
    """
    ⚡ FASTEST: Smart sampling strategy for large PDFs.
    """
    print(f"  ⚡ Smart sampling mode (faster, ~95% accurate)")

    page_invoice_nos = [None] * doc.page_count

    # Always extract from first page
    page = doc.load_page(0)
    page_invoice_nos[0] = extract_invoice_no_from_page(page, is_image_pdf)
    print(f"  ✓ Page 1: {page_invoice_nos[0]}")

    # Sample every Nth page to detect changes
    sample_interval = max(3, doc.page_count // 20)
    print(f"  Sampling interval: every {sample_interval} pages")

    for i in range(sample_interval, doc.page_count, sample_interval):
        page = doc.load_page(i)
        inv = extract_invoice_no_from_page(page, is_image_pdf)
        page_invoice_nos[i] = inv

        if i % 10 == 0:
            print(f"  Sampling page {i+1}/{doc.page_count}...")

        # If invoice changed, extract nearby pages to find exact boundary
        prev_known_idx = i - sample_interval
        while prev_known_idx >= 0 and page_invoice_nos[prev_known_idx] is None:
            prev_known_idx -= 1

        if prev_known_idx >= 0 and inv != page_invoice_nos[prev_known_idx]:
            print(f"  🔍 Boundary detected near page {i+1}, refining...")
            for offset in range(-3, 4):
                idx = i + offset
                if 0 <= idx < doc.page_count and page_invoice_nos[idx] is None:
                    page = doc.load_page(idx)
                    page_invoice_nos[idx] = extract_invoice_no_from_page(
                        page, is_image_pdf)

    # Also check last page
    if page_invoice_nos[-1] is None:
        page = doc.load_page(doc.page_count - 1)
        page_invoice_nos[-1] = extract_invoice_no_from_page(page, is_image_pdf)
        print(f"  ✓ Last page: {page_invoice_nos[-1]}")

    # Forward-fill gaps
    last_known = page_invoice_nos[0]
    filled = 0
    for i in range(len(page_invoice_nos)):
        if page_invoice_nos[i] is not None:
            last_known = page_invoice_nos[i]
        else:
            page_invoice_nos[i] = last_known
            filled += 1

    print(f"  ✅ Smart sampling complete: forward-filled {filled} pages")
    return page_invoice_nos


# ============================================================================
# PDF PROCESSING FUNCTIONS
# ============================================================================

def is_image_based_pdf(doc: fitz.Document, sample_pages: int = 3) -> Tuple[bool, float]:
    total_text_length = 0
    pages_to_check = min(sample_pages, doc.page_count)

    for i in range(pages_to_check):
        text = doc.load_page(i).get_text("text") or ""
        total_text_length += len(text.strip())

    avg_text_length = total_text_length / pages_to_check
    is_image_based = avg_text_length < 50

    print(f"  PDF Type: {'IMAGE-BASED' if is_image_based else 'TEXT-BASED'}")
    print(f"  Avg text per page: {avg_text_length:.1f} chars")
    return is_image_based, avg_text_length


def normalize_text_for_search(s: str) -> str:
    if not s:
        return s
    s = s.replace("\u00A0", " ")
    s = re.sub(r"[\r\n\t]+", " ", s)
    s = re.sub(r"[ ]{2,}", " ", s).strip()
    return s


def is_valid_invoice_number(candidate: str) -> bool:
    if not candidate or len(candidate) < 3:
        return False
    if len(candidate) == 15 and re.match(r'^[0-9A-Z]{15}$', candidate.upper()):
        return False
    if re.match(r'^\d+$', candidate):
        return 6 <= len(candidate) <= 15
    if re.match(r'^\d+\.\d{2,}$', candidate):
        return False
    has_letter = any(c.isalpha() for c in candidate)
    has_digit = any(c.isdigit() for c in candidate)
    return has_letter and has_digit


def try_extract_invoice_from_text(text: str) -> Optional[str]:
    if not text:
        return None
    text_norm = normalize_text_for_search(text)
    
    # DEBUG: Print first 600 chars
    print(f"\n🔍 DEBUG - Extracted text (first 600 chars):\n{text_norm[:600]}\n")

    # PRIORITY 1: Look for CREDIT number (14 digits, common in pharma invoices)
    credit_match = re.search(
        r"CREDIT\s*(?:NO|NUMBER|#)?\s*[:\-]?\s*(\d{12,20})",
        text_norm, re.IGNORECASE
    )
    if credit_match:
        credit_num = credit_match.group(1).strip()
        print(f"✓ Found CREDIT number: {credit_num}")
        if 12 <= len(credit_num) <= 20:
            return credit_num.upper()

    # PRIORITY 2: Look for "Invoice No" or "Bill No" followed by long numeric (12-20 digits)
    invoice_patterns = [
        r"Invoice\s*(?:No|Number)\.?\s*[:\-]?\s*(\d{12,20})",
        r"Bill\s*(?:No|Number)\.?\s*[:\-]?\s*(\d{12,20})",
        r"Tax\s*Invoice\s*(?:No|Number)\.?\s*[:\-]?\s*(\d{12,20})",
    ]
    
    for pattern in invoice_patterns:
        match = re.search(pattern, text_norm, re.IGNORECASE)
        if match:
            num = match.group(1).strip()
            print(f"✓ Found labeled long numeric invoice: {num}")
            return num.upper()

    # PRIORITY 3: Look for "Invoice No" with alphanumeric (but EXCLUDE batch patterns)
    label_patterns = [
        r"Invoice\s*No\.?\s*[:\-]\s*([A-Z][A-Z0-9\-\/]{2,20})",
        r"Bill\s*No\.?\s*[:\-]\s*([A-Z][A-Z0-9\-\/]{2,20})",
    ]
    
    for pattern in label_patterns:
        match = re.search(pattern, text_norm, re.IGNORECASE)
        if match:
            invoice_num = match.group(1).strip()
            
            # EXCLUDE batch number patterns (single letter + 6 digits: F500256, I500734, etc.)
            if re.match(r'^[A-Z]\d{6}$', invoice_num, re.IGNORECASE):
                print(f"⚠️ Skipping (batch pattern): {invoice_num}")
                continue
            
            # EXCLUDE license patterns (KA-MY2-157424)
            if re.match(r'^[A-Z]{2,3}-[A-Z0-9]+-\d+$', invoice_num, re.IGNORECASE):
                print(f"⚠️ Skipping (license pattern): {invoice_num}")
                continue
                
            print(f"✓ Found labeled alphanumeric: {invoice_num}")
            if any(c.isalpha() for c in invoice_num) and any(c.isdigit() for c in invoice_num):
                if 3 <= len(invoice_num) <= 20:
                    return invoice_num.upper()

    # PRIORITY 4: Look for long numeric values (12-20 digits) in top area
    top_text = text_norm[:1000]
    long_numerics = re.findall(r'\b(\d{12,20})\b', top_text)
    
    if long_numerics:
        # Take the longest one (most likely to be invoice number)
        longest = max(long_numerics, key=len)
        print(f"✓ Found long numeric value: {longest}")
        return longest.upper()

    # PRIORITY 5: Look near "Invoice" label for tokens, EXCLUDE batch patterns
    label_match = re.search(
        r"(?:Invoice|Bill|Tax\s*Invoice)\s*(?:No|Number|#|\.|:\s*)",
        text_norm, re.IGNORECASE
    )
    if label_match:
        start_idx = label_match.end()
        candidate_text = text_norm[start_idx:start_idx + 100]
        print(f"🔍 Text after label: '{candidate_text[:50]}...'")
        
        tokens = re.findall(r"\b([A-Z0-9][A-Z0-9\-\/]{2,20})\b", candidate_text, re.IGNORECASE)
        print(f"🔍 Tokens found: {tokens}")
        
        for token in tokens:
            token = token.strip(".,;:-*")
            
            # Skip common words
            if token.upper() in ("ORDER", "REF", "NO", "DATE", "DT", "INV", "BILL", "ACCOUNT", "PO", "COPY", "OF"):
                continue
            
            # EXCLUDE batch patterns (F500256, I500734)
            if re.match(r'^[A-Z]\d{6}$', token, re.IGNORECASE):
                print(f"⚠️ Skipping (batch pattern): {token}")
                continue
            
            # EXCLUDE license patterns
            if re.match(r'^[A-Z]{2,3}-[A-Z0-9]+-\d+$', token, re.IGNORECASE):
                print(f"⚠️ Skipping (license pattern): {token}")
                continue
                
            if any(c.isalpha() for c in token) and any(c.isdigit() for c in token):
                if 3 <= len(token) <= 20:
                    print(f"✓ Selected token: {token}")
                    return token.upper()

    # PRIORITY 6: Medium-length numeric (10-15 digits)
    medium_numerics = re.findall(r'\b(\d{10,15})\b', top_text)
    for num in medium_numerics:
        # Exclude phone numbers (10 digits starting with 6-9)
        if len(num) == 10 and num[0] in '6789':
            continue
        # Exclude dates (8 digits starting with 20)
        if len(num) == 8 and num.startswith('20'):
            continue
        print(f"✓ Found medium numeric value: {num}")
        return num.upper()

    print("✗ No invoice number found")
    return None

def extract_invoice_text_based(page: fitz.Page) -> Optional[str]:
    text = page.get_text("text") or ""
    inv = try_extract_invoice_from_text(text)
    if inv:
        return inv
    for block in (page.get_text("blocks") or []):
        block_text = block[4] if len(block) > 4 else ""
        if block_text:
            inv = try_extract_invoice_from_text(block_text)
            if inv:
                return inv
    return None


def extract_invoice_no_from_page(page: fitz.Page, is_image_pdf: bool) -> Optional[str]:
    """Extract invoice number from a single page (used by smart sampling)."""
    text_result = extract_invoice_text_based(page)
    if text_result:
        return text_result
    if is_image_pdf:
        return extract_invoice_gemini_sync(page)
    return None


def build_pdf_from_pages(src_doc: fitz.Document, page_indices: List[int]) -> bytes:
    out = fitz.open()
    try:
        for i in page_indices:
            out.insert_pdf(src_doc, from_page=i, to_page=i)
        pdf_bytes = out.tobytes(garbage=4, deflate=True)
        return pdf_bytes
    finally:
        out.close()


def remove_file(path: str):
    try:
        if os.path.exists(path):
            os.remove(path)
    except Exception as e:
        print(f"⚠️ Cleanup warning: {e}")


# ============================================================================
# API ENDPOINTS
# ============================================================================

@app.post("/split-invoices")
async def split_invoices(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),

    # ⭐ REQUIRED: Batch ID
    batch_id: str = Form(...,
                         description="Batch ID (required) - used for folder structure"),

    # Blob Storage options
    use_blob_storage: bool = Form(
        True, description="Upload PDFs to Azure Blob Storage"),
    blob_container: Optional[str] = Form(
        None, description="Custom Azure container (optional)"),

    # Response options
    include_base64: bool = Form(
        False, description="Include base64 in response"),

    # Performance options
    parallel_batch_size: int = Form(
        MAX_PARALLEL_GEMINI_CALLS, description="Parallel Gemini API calls (1-10)"),
    use_smart_sampling: bool = Form(
        USE_SMART_SAMPLING, description="Use smart sampling (faster, ~95% accurate)"),

    # File size limit
    max_file_size_mb: int = Form(200, description="Maximum file size in MB"),
):
    """
    ⭐ OPTIMIZED INVOICE SPLITTER - SUPPORTS PDF AND IMAGES

    Performance Improvements:
    - Parallel Gemini API calls (5-10x faster for image PDFs)
    - Smart sampling option for large PDFs
    - Reduced image resolution for faster processing
    - Optimized prompts for quicker responses

    File Support:
    - PDF files (text-based or image-based)
    - Image files (PNG, JPG, JPEG, TIFF, BMP) - auto-converted to PDF

    Folder Structure in Blob Storage:
    POD/
      └── {batch_id}/
           └── {filename}/
                ├── Raw/ (original uploaded file)
                └── Splitted/ (individual split invoice PDFs)

    Required Parameters:
    - file: PDF or image file to upload
    - batch_id: Batch identifier (used for folder structure)

    Returns:
    - All invoice URLs with proper folder paths
    """

    # ============================================================================
    # ENHANCED VALIDATION - ACCEPT PDF AND IMAGES
    # ============================================================================
    
    if not file.filename:
        raise HTTPException(status_code=400, detail="No filename provided")
    
    filename_lower = file.filename.lower()
    
    # Supported formats
    SUPPORTED_EXTENSIONS = ['.pdf', '.png', '.jpg', '.jpeg', '.tiff', '.tif', '.bmp']
    
    file_extension = None
    for ext in SUPPORTED_EXTENSIONS:
        if filename_lower.endswith(ext):
            file_extension = ext
            break
    
    if not file_extension:
        raise HTTPException(
            status_code=400, 
            detail=f"Unsupported file format. Supported: PDF, PNG, JPG, JPEG, TIFF, BMP"
        )
    
    is_image_file = file_extension in ['.png', '.jpg', '.jpeg', '.tiff', '.tif', '.bmp']
    
    # Check PIL availability for image files
    if is_image_file and not GEMINI_AVAILABLE:
        raise HTTPException(
            status_code=500,
            detail="Image processing requires PIL. Install: pip install Pillow"
        )

    # Check blob storage
    if use_blob_storage and not get_blob_service_client():
        raise HTTPException(
            status_code=500, detail="Azure Blob Storage not configured")

    # Container
    container_name = blob_container if blob_container else AZURE_CONTAINER_NAME

    # Ensure container exists
    if use_blob_storage:
        ensure_container_exists(container_name)

    # Stream upload to temp file
    max_size_bytes = max_file_size_mb * 1024 * 1024
    fd, temp_path = tempfile.mkstemp(suffix=file_extension)
    os.close(fd)

    doc = None
    original_pdf_bytes = None
    start_time = datetime.now()
    pdf_path = temp_path
    original_filename = file.filename

    try:
        print(f"\n{'='*70}")
        print(f"📥 Processing: {file.filename}")
        print(f"   File Type: {'IMAGE' if is_image_file else 'PDF'}")
        print(f"   Batch ID: {batch_id}")
        print(
            f"   Performance Mode: {'Smart Sampling' if use_smart_sampling else f'Parallel ({parallel_batch_size} workers)'}")
        print(f"{'='*70}")

        total_size = 0
        with open(temp_path, "wb") as buffer:
            chunk_read_size = 5 * 1024 * 1024
            while content := await file.read(chunk_read_size):
                total_size += len(content)
                if total_size > max_size_bytes:
                    remove_file(temp_path)
                    raise HTTPException(
                        status_code=413, detail=f"File too large. Max: {max_file_size_mb}MB")
                buffer.write(content)

        file_size_mb = total_size / (1024 * 1024)
        print(f"💾 File size: {file_size_mb:.2f}MB")

        # ============================================================================
        # IMAGE TO PDF CONVERSION
        # ============================================================================
        
        if is_image_file:
            print(f"🖼️  Converting image to PDF...")
            try:
                from PIL import Image as PILImage
                
                # Open image and convert to PDF
                img = PILImage.open(temp_path)
                
                # Convert to RGB if necessary (for RGBA, grayscale, etc.)
                if img.mode != 'RGB':
                    img = img.convert('RGB')
                
                # Create PDF path
                pdf_path = temp_path.replace(file_extension, '.pdf')
                
                # Save as PDF
                img.save(pdf_path, 'PDF', resolution=100.0)
                img.close()
                
                print(f"✅ Image converted to PDF")
                
                # Update filename for storage
                file.filename = file.filename.replace(file_extension, '.pdf')
                
            except Exception as e:
                print(f"❌ Image conversion failed: {e}")
                raise HTTPException(
                    status_code=500, 
                    detail=f"Failed to convert image to PDF: {str(e)}"
                )

        # Read PDF bytes (either original or converted)
        with open(pdf_path, "rb") as f:
            original_pdf_bytes = f.read()

        # Upload original PDF to Raw folder
        raw_pdf_info = None
        if use_blob_storage:
            try:
                print(f"\n📤 Uploading original {'PDF' if not is_image_file else 'converted PDF'} to Raw folder...")
                raw_pdf_info = upload_raw_pdf_to_blob(
                    original_pdf_bytes,
                    file.filename,
                    batch_id,
                    container_name
                )
                print(f"✅ Original PDF uploaded: {raw_pdf_info['blob_name']}")
            except Exception as e:
                print(f"⚠️ Failed to upload raw PDF: {e}")

        # Open PDF for processing
        doc = fitz.open(pdf_path)
        if doc.page_count == 0:
            raise HTTPException(status_code=400, detail="Empty PDF")

        print(f"📄 Total pages: {doc.page_count}")

        # Detect PDF type
        is_image_pdf, _ = is_image_based_pdf(doc)
        if is_image_pdf and not get_gemini_model():
            raise HTTPException(
                status_code=500, detail="Image PDF detected but Gemini not configured")

        # ⚡ OPTIMIZED EXTRACTION
        print(f"\n📊 Extracting invoice numbers...")
        extraction_start = datetime.now()

        if use_smart_sampling and doc.page_count > 10:
            # Smart sampling for large PDFs
            page_invoice_nos = extract_invoices_smart_sampling(
                doc, is_image_pdf)
        else:
            # Parallel extraction (async batch processing)
            page_invoice_nos = await extract_invoices_batch_async(
                doc,
                is_image_pdf,
                batch_size=parallel_batch_size
            )

        extraction_time = (datetime.now() - extraction_start).total_seconds()
        print(f"✅ Extraction completed in {extraction_time:.1f} seconds")
        print(f"   Speed: {doc.page_count / extraction_time:.1f} pages/second")

        # ============================================================================
        # 🔧 CORRECTED GROUPING LOGIC - NO AGGRESSIVE FILTERING
        # ============================================================================

        print(f"\n🔧 Grouping invoices...")

        # DEBUG: Show raw extraction results
        print(f"\n🔍 DEBUG - Raw extraction results:")
        for idx, inv in enumerate(page_invoice_nos[:min(10, len(page_invoice_nos))]):
            print(f"   Page {idx+1}: {inv if inv else '(not found)'}")
        if len(page_invoice_nos) > 10:
            print(
                f"   ... (showing first 10 of {len(page_invoice_nos)} pages)")

        # Step 1: Normalize extracted invoice numbers (only filter GST numbers)
        page_invoice_nos_normalized = []
        for v in page_invoice_nos:
            if v and v.upper().startswith("GST"):
                # Filter out GST numbers (not invoice numbers)
                page_invoice_nos_normalized.append(None)
            elif v:
                # Normalize: uppercase, remove spaces/underscores
                normalized = v.upper().strip().replace(" ", "").replace("_", "")
                page_invoice_nos_normalized.append(normalized)
            else:
                page_invoice_nos_normalized.append(None)

        # Step 2: Smart forward-fill for failed extractions
        # Only fill None values, DON'T remove any extracted invoice numbers
        page_invoice_nos_filled = []
        last_known_invoice = None

        for idx, inv in enumerate(page_invoice_nos_normalized):
            if inv is not None:
                # Valid invoice number found
                last_known_invoice = inv
                page_invoice_nos_filled.append(inv)
            else:
                # Extraction failed - use last known invoice
                page_invoice_nos_filled.append(last_known_invoice)

        # Count how many pages were forward-filled
        filled_count = sum(1 for i in range(len(page_invoice_nos_normalized))
                           if page_invoice_nos_normalized[i] is None and page_invoice_nos_filled[i] is not None)

        # Debug: Count unique invoice numbers
        unique_invoices = set(
            [v for v in page_invoice_nos_filled if v is not None])
        print(f"\n   📊 Found {len(unique_invoices)} unique invoice numbers:")
        for inv_no in sorted(unique_invoices) if unique_invoices else []:
            page_count = sum(1 for v in page_invoice_nos_filled if v == inv_no)
            print(f"      • {inv_no}: {page_count} pages")

        # Step 3: Group consecutive pages by invoice number
        groups = []
        current_group = []
        current_invoice = None

        for idx, inv in enumerate(page_invoice_nos_filled):
            if idx == 0:
                # First page
                current_invoice = inv
                current_group = [idx]
            else:
                if inv != current_invoice:
                    # Invoice number changed - save current group and start new one
                    groups.append({
                        "invoice_no": current_invoice,
                        "pages": current_group[:]
                    })
                    print(
                        f"   📄 Group {len(groups)}: Invoice {current_invoice or 'UNKNOWN'} - Pages {current_group[0]+1}-{current_group[-1]+1} ({len(current_group)} pages)")
                    current_invoice = inv
                    current_group = [idx]
                else:
                    # Same invoice - add to current group
                    current_group.append(idx)

        # Don't forget the last group
        if current_group:
            groups.append({
                "invoice_no": current_invoice,
                "pages": current_group[:]
            })
            print(
                f"   📄 Group {len(groups)}: Invoice {current_invoice or 'UNKNOWN'} - Pages {current_group[0]+1}-{current_group[-1]+1} ({len(current_group)} pages)")

        # Handle edge case: entire PDF has no invoice numbers
        if len(groups) == 1 and groups[0]["invoice_no"] is None:
            groups = [{
                "invoice_no": None,
                "pages": list(range(doc.page_count))
            }]

        print(f"\n✅ Created {len(groups)} invoice groups")
        print(
            f"   Forward-filled {filled_count} pages with missing invoice numbers")

        # Build and upload split PDFs
        print(f"\n🔨 Building and uploading split invoices...")
        all_parts = []

        for idx, g in enumerate(groups):
            if (idx + 1) % 20 == 0:
                print(f"  Processing {idx + 1}/{len(groups)} invoices...")

            # Build PDF
            part_bytes = build_pdf_from_pages(doc, g["pages"])

            # Generate filename
            invoice_no = g["invoice_no"] if g["invoice_no"] else f"NO_NUMBER_{idx + 1}"
            safe_invoice_no = re.sub(r'[<>:"/\\|?*]', '_', invoice_no)
            invoice_filename = f"invoice_{safe_invoice_no}.pdf"

            # Prepare invoice info
            invoice_info = {
                "invoice_no": g["invoice_no"],
                "pages": [p + 1 for p in g["pages"]],
                "page_range": f"{g['pages'][0]+1}-{g['pages'][-1]+1}" if len(g['pages']) > 1 else f"{g['pages'][0]+1}",
                "num_pages": len(g["pages"]),
                "size_bytes": len(part_bytes),
                "size_mb": round(len(part_bytes) / (1024 * 1024), 2),
            }

            # Upload to Splitted folder
            if use_blob_storage:
                try:
                    blob_info = upload_split_pdf_to_blob(
                        part_bytes,
                        invoice_filename,
                        file.filename,
                        batch_id,
                        container_name
                    )
                    invoice_info["storage"] = blob_info
                    invoice_info["pdf_url"] = blob_info["download_url"]
                    invoice_info["blob_name"] = blob_info["blob_name"]
                    invoice_info["expires_at"] = blob_info["expires_at"]
                except Exception as e:
                    print(f"  ⚠️ Failed to upload invoice {idx+1}: {e}")
                    invoice_info["upload_error"] = str(e)

            # Include base64 if requested
            if include_base64:
                invoice_info["pdf_base64"] = base64.b64encode(
                    part_bytes).decode("ascii")

            all_parts.append(invoice_info)
            del part_bytes

            if idx % 50 == 0:
                gc.collect()

        print(f"✅ Processed all {len(all_parts)} invoices")

        # ⭐ SAVE VALUES BEFORE CLOSING DOCUMENT
        total_pages_count = doc.page_count

        # Close document
        doc.close()
        doc = None
        
        # Clean up temp files
        remove_file(temp_path)
        if pdf_path != temp_path:
            remove_file(pdf_path)
        
        gc.collect()

        # Calculate total processing time
        total_time = (datetime.now() - start_time).total_seconds()

        # Return response
        response_data = {
            "success": True,
            "batch_id": batch_id,
            "folder_structure": {
                "root": ROOT_FOLDER,
                "path": f"{ROOT_FOLDER}/{batch_id}/{os.path.splitext(file.filename)[0]}",
                "raw_folder": f"{ROOT_FOLDER}/{batch_id}/{os.path.splitext(file.filename)[0]}/Raw",
                "split_folder": f"{ROOT_FOLDER}/{batch_id}/{os.path.splitext(file.filename)[0]}/Splitted"
            },
            "source_file": {
                "name": file.filename,
                "original_name": original_filename,
                "size_mb": round(file_size_mb, 2),
                "total_pages": total_pages_count,
                "pdf_type": "image-based" if is_image_pdf else "text-based",
                "was_converted": is_image_file,
                "raw_pdf": raw_pdf_info
            },
            "summary": {
                "total_invoices": len(all_parts),
                "unique_invoice_numbers": len(unique_invoices),
                "extraction_method": "gemini" if is_image_pdf else "text",
                "pages_forward_filled": filled_count,
                "storage_type": "azure_blob" if use_blob_storage else "base64"
            },
            "performance": {
                "total_time_seconds": round(total_time, 2),
                "extraction_time_seconds": round(extraction_time, 2),
                "pages_per_second": round(total_pages_count / extraction_time, 2) if extraction_time > 0 else 0,
                "parallel_batch_size": parallel_batch_size,
                "smart_sampling_used": use_smart_sampling and total_pages_count > 10
            },
            "invoices": all_parts
        }

        print(f"\n{'='*70}")
        print(f"✅ SUCCESS!")
        print(f"   Batch ID: {batch_id}")
        print(f"   Original File: {original_filename}")
        if is_image_file:
            print(f"   ✓ Image converted to PDF")
        print(
            f"   Raw PDF: {raw_pdf_info['blob_name'] if raw_pdf_info else 'Not uploaded'}")
        print(f"   Split invoices: {len(all_parts)}")
        print(f"   Unique invoice numbers: {len(unique_invoices)}")
        print(f"   Total time: {total_time:.1f}s")
        print(
            f"   Extraction time: {extraction_time:.1f}s ({total_pages_count / extraction_time:.1f} pages/sec)")
        print(f"{'='*70}\n")

        return JSONResponse(response_data)

    except HTTPException:
        raise
    except Exception as e:
        print(f"\n✗ Error: {str(e)}")
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        if doc:
            doc.close()
        remove_file(temp_path)
        if pdf_path != temp_path:
            remove_file(pdf_path)
        gc.collect()


@app.post("/cleanup-batch/{batch_id}")
async def cleanup_batch(
    batch_id: str,
    background_tasks: BackgroundTasks,
    container_name: Optional[str] = Form(None)
):
    """Delete all blobs for a specific batch (entire POD/{batch_id}/ folder)."""
    if container_name is None:
        container_name = AZURE_CONTAINER_NAME

    background_tasks.add_task(cleanup_old_blobs, batch_id, container_name)

    return JSONResponse({
        "success": True,
        "message": f"Cleanup started for batch {batch_id}",
        "batch_id": batch_id,
        "folder_path": f"{ROOT_FOLDER}/{batch_id}/",
        "container": container_name
    })


    # ============================================================
# GENERATE DOWNLOAD URL FROM BLOB PATH (PERMANENT ACCESS)
# ============================================================


@app.get("/blob/download-url")
def generate_download_url(
    blob_name: str = Query(..., description="Blob path like POD/BATCH.../file.pdf"),
    container: str = Query(AZURE_CONTAINER_NAME),
    expiry_minutes: int = Query(60)
):
    """
    Returns fresh SAS download URL.
    This prevents expired links problem.
    """

    try:
        sas_token = generate_blob_sas(
            account_name=AZURE_STORAGE_ACCOUNT_NAME,
            container_name=container,
            blob_name=blob_name,
            account_key=AZURE_STORAGE_ACCOUNT_KEY,
            permission=BlobSasPermissions(read=True),
            expiry=datetime.utcnow() + timedelta(minutes=expiry_minutes)
        )

        download_url = f"https://{AZURE_STORAGE_ACCOUNT_NAME}.blob.core.windows.net/{container}/{blob_name}?{sas_token}"

        return {
            "success": True,
            "download_url": download_url,
            "expires_in_minutes": expiry_minutes
        }

    except Exception as e:
        return {"success": False, "error": str(e)}