src/streamlit_app.py

# =========================
# Invoice Extractor (Qwen3-VL via RunPod vLLM) - Batch Mode with Tax Validation
# Updated code #2 with code #1 postprocessing logic
# =========================
import os
from pathlib import Path

# -----------------------------
# Environment hardening (HF Spaces, /.cache issue)
# -----------------------------
_home = os.environ.get("HOME", "")
if _home in ("", "/", None):
    repo_dir = os.getcwd()
    safe_home = repo_dir if os.access(repo_dir, os.W_OK) else "/tmp"
    os.environ["HOME"] = safe_home
    print(f"[startup] HOME not set or unwritable — setting HOME={safe_home}")

streamlit_dir = Path(os.environ["HOME"]) / ".streamlit"
try:
    streamlit_dir.mkdir(parents=True, exist_ok=True)
    print(f"[startup] ensured {streamlit_dir}")
except Exception as e:
    print(f"[startup] WARNING: could not create {streamlit_dir}: {e}")

# -----------------------------
# Imports
# -----------------------------
import json
from io import BytesIO
import hashlib
from typing import Dict, Any, List
from datetime import datetime, timedelta
from copy import deepcopy

import streamlit as st
import pandas as pd
from PIL import Image

try:
    from pdf2image import convert_from_bytes
except Exception:
    convert_from_bytes = None

import requests
import base64
import re

POD_URL      = os.getenv("POD_URL", "https://5y05rd290uuemp-8000.proxy.runpod.net")
VLLM_API_KEY = os.getenv("VLLM_API_KEY", "9386aa7335c2072882ad367791f1dc863c63b4be88eda14c2dea0408f9d0105e")
MODEL_NAME   = "phase2-v1-merged"

MAX_PAGES_PER_REQUEST = 10
MAX_DIMENSION_PER_PAGE = 1560
MAX_TOTAL_PAYLOAD_MB   = 40.0

if not POD_URL or not VLLM_API_KEY:
    st.error("⚠️ API credentials not configured. Please set POD_URL and VLLM_API_KEY in Space settings.")
    st.stop()

# -----------------------------
# Page config & CSS
# -----------------------------
st.set_page_config(page_title="Invoice Extractor (Qwen3-VL) - Batch Mode", layout="wide")
st.title("Invoice Extraction")

st.markdown(
    """
    <style>
        .stApp { background-color: #ECECEC !important; }
        div.block-container { padding-top: 3rem; padding-bottom: 1rem; }
        [data-testid="stSidebar"] { background-color: #F7F7F7 !important; }
        div[data-testid="stTabs"] > div > div { padding-bottom: 6px !important; }
        [data-testid="column"]:nth-of-type(2) { min-height: 780px; }
        [data-testid="stImage"] img { image-rendering: auto; }
    </style>
    <script>
        // Push a dummy state so there's always something to pop back to
        history.pushState(null, null, window.location.href);
        window.addEventListener('popstate', function(event) {
            history.pushState(null, null, window.location.href);
        });
    </script>
    """,
    unsafe_allow_html=True
)

DATA_EDITOR_HEIGHT = 380

# -----------------------------
# Helpers — numeric/date cleaning from code #1
# -----------------------------
def ensure_state(k: str, default):
    """Initialize a session_state key once, then let widgets bind to it via key=... (no value=...)."""
    if k not in st.session_state:
        st.session_state[k] = default

def parse_time_to_minutes(x) -> float:
    """
    Parse time format quantities to minutes.
    
    Examples:
        "0:35"  → 35.0   (0 hours, 35 minutes = 35 minutes)
        "1:30"  → 90.0   (1 hour, 30 minutes = 90 minutes)
        "2:15"  → 135.0  (2 hours, 15 minutes = 135 minutes)
        "0:05"  → 5.0    (5 minutes)
        "123"   → 123.0  (regular number, not time format)
        "1.5"   → 1.5    (regular decimal, not time format)
    """
    if x is None:
        return 0.0
    if isinstance(x, (int, float)):
        return float(x)
    
    s = str(x).strip()
    if s == "":
        return 0.0
    
    # Check if it's in time format (H:MM or HH:MM)
    time_pattern = r'^(\d+):(\d{1,2})$'
    match = re.match(time_pattern, s)
    
    if match:
        hours = int(match.group(1))
        minutes = int(match.group(2))
        total_minutes = (hours * 60) + minutes
        return float(total_minutes)
    
    # Not time format, treat as regular number
    return 0.0

def clean_quantity(x, currency=None) -> float:
    """
    Parse quantity - handles both time format (H:MM) and regular numbers.
    Currency-aware the same way clean_float is: for EUR, an ambiguous single-comma
    value like "1,000" is read as the European decimal 1.0, not the US thousands 1000.

    Examples:
        "0:35"     → 35.0   (time format: 35 minutes)
        "1:30"     → 90.0   (time format: 90 minutes)
        "123"      → 123.0  (regular number)
        "1,234.56" → 1234.56 (US format with decimals)
        "1.234,56" → 1234.56 (EUR format with decimals)
        "1,000" with currency="EUR" → 1.0 (EUR decimal, not US thousands)
    """
    if x is None:
        return 0.0
    if isinstance(x, (int, float)):
        return float(x)

    s = str(x).strip()
    if s == "":
        return 0.0

    # First check if it's time format (H:MM or HH:MM)
    time_value = parse_time_to_minutes(s)
    if time_value > 0.0:
        return time_value

    # Not time format, use regular number parsing
    return clean_float(s, currency)

def clean_tax_percentage(x) -> float:
    """
    Parse tax percentage - ALWAYS treats periods as decimals, never as thousands.
    
    Tax percentages are typically: 8.875, 19.5, 2.75, 0.0875, etc.
    We should NEVER interpret "8.875" as "8875"
    
    Examples:
        "8.875"    → 8.875   (decimal percentage)
        "8,875"    → 8.875   (European decimal)
        "19.5"     → 19.5    (standard decimal)
        "2,75"     → 2.75    (European decimal)
    """
    if x is None:
        return 0.0
    if isinstance(x, (int, float)):
        return float(x)
    
    s = str(x).strip()
    if s == "":
        return 0.0
    
    # Handle negative signs (could be leading or trailing)
    is_negative = False
    if s.startswith('-'):
        is_negative = True
        s = s[1:].strip()
    elif s.endswith('-'):
        is_negative = True
        s = s[:-1].strip()
    elif s.startswith('(') and s.endswith(')'):
        is_negative = True
        s = s[1:-1].strip()
    
    # Remove percent signs, currency symbols and spaces
    s = re.sub(r'[%€$£¥₹\s]', '', s)
    
    if s == "":
        return 0.0
    
    # Count occurrences
    comma_count = s.count(',')
    period_count = s.count('.')
    
    # For tax percentages, logic is simpler:
    # - If both comma and period exist, the LAST one is decimal separator
    # - If only comma exists, it's decimal separator (European style)
    # - If only period exists, it's ALWAYS decimal separator (no thousands logic)
    
    if comma_count > 0 and period_count > 0:
        # Both present - LAST one is decimal
        last_comma = s.rfind(',')
        last_period = s.rfind('.')
        
        if last_comma > last_period:
            # European: 1.234,56 → comma is decimal
            s = s.replace('.', '').replace(',', '.')
        else:
            # US: 1,234.56 → period is decimal
            s = s.replace(',', '')
    
    elif comma_count > 0:
        # Only comma - treat as European decimal separator
        s = s.replace(',', '.')
    
    # If only period(s) exist, keep as-is (always decimal for tax percentages)
    
    # Clean any remaining non-numeric characters except period and minus
    s = re.sub(r'[^\d.]', '', s)
    
    if s == "" or s == ".":
        return 0.0
    
    try:
        result = float(s)
        return -result if is_negative else result
    except ValueError:
        return 0.0

def clean_float(x, currency=None) -> float:
    """
    Parse a number string handling both US and European formats.
    Use this for MONETARY AMOUNTS only, NOT for tax percentages.
    US Format:      1,234,567.89  (comma = thousands, period = decimal)
    European:       1.234.567,89  (period = thousands, comma = decimal)
    Currency-aware for ambiguous cases (3 digits after comma):
        - EUR: "10,000" → 10.0 (European decimal)
        - USD/other: "10,000" → 10000 (thousands separator)
    Examples:
        "1,234.56"    → 1234.56  (US)
        "1.234,56"    → 1234.56  (European)
        "3.000,2234"  → 3000.2234 (European with 4 decimal places)
        "261,49"      → 261.49   (European decimal only)
        "39,22-"      → -39.22   (European with trailing minus)
    """
    if x is None:
        return 0.0
    if isinstance(x, (int, float)):
        return float(x)
    
    s = str(x).strip()
    if s == "":
        return 0.0
    
    # Handle negative signs (could be leading or trailing)
    is_negative = False
    if s.startswith('-'):
        is_negative = True
        s = s[1:].strip()
    elif s.endswith('-'):
        is_negative = True
        s = s[:-1].strip()
    elif s.startswith('(') and s.endswith(')'):
        # Accounting format: (123.45) means negative
        is_negative = True
        s = s[1:-1].strip()
    
    # Remove currency symbols and spaces
    s = re.sub(r'[€$£¥₹\s]', '', s)
    
    if s == "":
        return 0.0
    
    # Count occurrences
    comma_count = s.count(',')
    period_count = s.count('.')
    
    # Find positions of last comma and last period
    last_comma = s.rfind(',')
    last_period = s.rfind('.')
    
    # Determine format based on which separator comes last
    if comma_count > 0 and period_count > 0:
        # Both separators present - the LAST one is the decimal separator
        if last_comma > last_period:
            # European format: 1.234,56 → comma is decimal
            # Remove periods (thousands), replace comma with period
            s = s.replace('.', '').replace(',', '.')
        else:
            # US format: 1,234.56 → period is decimal
            # Remove commas (thousands)
            s = s.replace(',', '')
    
    elif comma_count > 0 and period_count == 0:
        # Only commas present
        # Check what comes after the LAST comma
        after_last_comma = s[last_comma + 1:] if last_comma < len(s) - 1 else ""

        if comma_count == 1 and len(after_last_comma) == 3 and after_last_comma.isdigit():
            # AMBIGUOUS CASE: Single comma with exactly 3 digits after
            # Could be US thousands ("10,000" = 10000) or European decimal ("10,000" = 10.0)
            # Use currency to decide:
            if currency and currency.upper() == 'EUR':
                # European: treat comma as decimal separator
                # "10,000" → 10.000 → 10.0
                s = s.replace(',', '.')
            else:
                # USD/other: treat comma as thousands separator
                # "10,000" → 10000
                s = s.replace(',', '')
        elif comma_count == 1 and len(after_last_comma) <= 4 and after_last_comma.isdigit():
            # Single comma with 1, 2, or 4 digits after → European decimal
            # "261,49" → 261.49, "1234,5678" → 1234.5678
            s = s.replace(',', '.')
        elif len(after_last_comma) == 3 and comma_count >= 1:
            # Multiple commas with 3 digits after last → thousands separator
            # "1,234,567" → 1234567
            s = s.replace(',', '')
        else:
            # Multiple commas → thousands separator
            # "1,234,567" → 1234567
            s = s.replace(',', '')
    
    elif period_count > 0 and comma_count == 0:
        # Only periods present
        # Check what comes after the LAST period
        after_last_period = s[last_period + 1:] if last_period < len(s) - 1 else ""
        
        if period_count > 1:
            # Multiple periods → definitely thousands separator (European: "1.234.567")
            s = s.replace('.', '')
        elif len(after_last_period) == 3 and after_last_period.isdigit():
            # Single period with exactly 3 digits after
            before_period = s[:last_period]
            # Only treat as thousands if there are 2+ digits before period
            # "10.000" or "123.000" → thousands
            # "8.875" or "9.123" → decimal
            if before_period.isdigit() and len(before_period) >= 2 and len(before_period) <= 3:
                s = s.replace('.', '')  # European thousands: "10.000" → 10000
            # Otherwise keep as decimal: "8.875" → 8.875
        # Otherwise keep as is (standard decimal like "1.50", "123.45")
    
    # Clean any remaining non-numeric characters except period and minus
    s = re.sub(r'[^\d.]', '', s)
    
    if s == "" or s == ".":
        return 0.0
    
    try:
        result = float(s)
        return -result if is_negative else result
    except ValueError:
        return 0.0

DECIMAL_DISPLAY_PLACES = 3
TAX_RATE_CLOSE_CALL_TOLERANCE = 0.1
DISCOUNT_RATE_CLOSE_CALL_TOLERANCE = 0.1
SUMMARY_NUMERIC_FIELDS = {
    "Subtotal",
    "Tax Percentage",
    "Total Tax",
    "Discount Rate",
    "Total Discount Amount",
    "Total Amount",
}
TAX_RATE_FIELDS = {"Tax Percentage", "Tax Rate Per Item"}
DISCOUNT_RATE_FIELDS = {"Discount Rate", "Discount Rate Per Item"}
ITEM_NUMERIC_FIELDS = {
    "Quantity",
    "Unit Price",
    "Amount",
    "Discount Rate Per Item",
    "Discount Amount Per Item",
    "Tax Rate Per Item",
    "Tax Amount Per Item",
    "Line Total",
    "Tax",
    "line_total",
}
EDITOR_NUMERIC_COLUMNS = [
    "Quantity",
    "Unit Price",
    "Amount",
    "Discount Rate Per Item",
    "Discount Amount Per Item",
    "Tax Rate Per Item",
    "Tax Amount Per Item",
    "Line Total",
]

def is_blank_numeric_value(value) -> bool:
    if value is None:
        return True
    try:
        if bool(pd.isna(value)):
            return True
    except (TypeError, ValueError):
        pass
    return isinstance(value, str) and value.strip() == ""

def limit_decimal_places(value, places: int = DECIMAL_DISPLAY_PLACES):
    """Round numeric display/export values to at most `places` decimals."""
    if is_blank_numeric_value(value):
        return ""
    rounded = round(clean_float(value), places)
    if rounded == 0:
        return 0.0
    return rounded

def normalize_rate_close_call(value, tolerance: float):
    """Snap a near-whole percentage rate to the nearest whole number."""
    if is_blank_numeric_value(value):
        return ""
    rate = clean_tax_percentage(value)
    if is_blank_numeric_value(rate):
        return ""
    nearest_whole = round(rate)
    if abs(rate - nearest_whole) <= tolerance + 1e-9:
        return float(nearest_whole)
    return limit_decimal_places(rate)

def normalize_tax_rate_close_call(value):
    """Snap near-whole tax rates like 19.99/20.002/20.01 to 20."""
    return normalize_rate_close_call(value, TAX_RATE_CLOSE_CALL_TOLERANCE)

def normalize_discount_rate_close_call(value):
    """Snap near-whole discount rates like 14.99/20.08 to the nearest whole number."""
    return normalize_rate_close_call(value, DISCOUNT_RATE_CLOSE_CALL_TOLERANCE)

def editable_numeric_value(field: str, value):
    if field in TAX_RATE_FIELDS:
        normalized = normalize_tax_rate_close_call(value)
    elif field in DISCOUNT_RATE_FIELDS:
        normalized = normalize_discount_rate_close_call(value)
    else:
        normalized = limit_decimal_places(value)
    if normalized == "":
        return 0.0
    return float(normalized)

def coerce_numeric_items_df(df: pd.DataFrame) -> pd.DataFrame:
    """Keep editable item numeric columns numeric before save/export."""
    out = df.copy() if isinstance(df, pd.DataFrame) else pd.DataFrame()
    for col in EDITOR_NUMERIC_COLUMNS:
        if col not in out.columns:
            out[col] = 0.0
        if col in TAX_RATE_FIELDS:
            out[col] = out[col].apply(normalize_tax_rate_close_call)
        elif col in DISCOUNT_RATE_FIELDS:
            out[col] = out[col].apply(normalize_discount_rate_close_call)
        else:
            out[col] = out[col].apply(limit_decimal_places)
        out[col] = pd.to_numeric(out[col], errors="coerce").fillna(0.0).astype(float)
    return out

def editor_dataframes_equal(left: pd.DataFrame, right: pd.DataFrame) -> bool:
    """Stable equality check for deciding whether to commit data_editor output."""
    if not isinstance(left, pd.DataFrame) or not isinstance(right, pd.DataFrame):
        return False
    left_norm = left.reset_index(drop=True).fillna("").astype(str)
    right_norm = right.reset_index(drop=True).fillna("").astype(str)
    return left_norm.equals(right_norm)

def apply_data_editor_state(base_df: pd.DataFrame, editor_state) -> pd.DataFrame:
    """Apply Streamlit data_editor edited/added/deleted row state to a base df."""
    if isinstance(editor_state, pd.DataFrame):
        return coerce_numeric_items_df(editor_state)
    if not isinstance(editor_state, dict):
        return coerce_numeric_items_df(base_df)

    out = base_df.copy() if isinstance(base_df, pd.DataFrame) else pd.DataFrame()

    for row_idx, updates in (editor_state.get("edited_rows") or {}).items():
        try:
            idx = int(row_idx)
        except (TypeError, ValueError):
            continue
        if idx < 0 or idx >= len(out) or not isinstance(updates, dict):
            continue
        for col, value in updates.items():
            if col in out.columns:
                out.at[out.index[idx], col] = value

    deleted_rows = sorted(
        [int(i) for i in (editor_state.get("deleted_rows") or []) if str(i).lstrip("-").isdigit()],
        reverse=True,
    )
    if deleted_rows:
        drop_labels = [out.index[i] for i in deleted_rows if 0 <= i < len(out)]
        out = out.drop(index=drop_labels)

    added_rows = editor_state.get("added_rows") or []
    if added_rows:
        out = pd.concat([out, pd.DataFrame(added_rows)], ignore_index=True)

    return coerce_numeric_items_df(out.reset_index(drop=True))

def commit_data_editor_state(items_state_key: str, editor_key: str, selected_hash: str) -> None:
    """Commit data_editor deltas to durable session state before rerun."""
    committed = apply_data_editor_state(
        st.session_state.get(items_state_key, pd.DataFrame()),
        st.session_state.get(editor_key),
    )
    st.session_state[items_state_key] = committed
    live_edited_data = deepcopy(st.session_state.batch_results[selected_hash].get("edited_data", {}) or {})
    live_edited_data["Itemized Data"] = committed.to_dict("records")
    st.session_state.batch_results[selected_hash]["edited_data"] = live_edited_data

def limit_invoice_decimals(data: dict) -> dict:
    """Apply the 3-decimal display cap only to known numeric invoice fields."""
    for field in SUMMARY_NUMERIC_FIELDS:
        if field in data:
            if field in TAX_RATE_FIELDS:
                data[field] = normalize_tax_rate_close_call(data.get(field))
            elif field in DISCOUNT_RATE_FIELDS:
                data[field] = normalize_discount_rate_close_call(data.get(field))
            else:
                data[field] = limit_decimal_places(data.get(field))

    for item in data.get("Itemized Data", []) or []:
        if not isinstance(item, dict):
            continue
        for field in ITEM_NUMERIC_FIELDS:
            if field in item:
                if field in TAX_RATE_FIELDS:
                    item[field] = normalize_tax_rate_close_call(item.get(field))
                elif field in DISCOUNT_RATE_FIELDS:
                    item[field] = normalize_discount_rate_close_call(item.get(field))
                else:
                    item[field] = limit_decimal_places(item.get(field))
    return data

def normalize_date(date_str, currency=None) -> str:
    """
    Normalize various date formats:
    - Full dates (day-month-year) → dd-MMM-yyyy (e.g., 01-Jan-2025)
    - Month-year only → MMM-yyyy (e.g., Aug-2025)
    
    Currency-aware parsing:
    - If currency is USD and date is numeric format (11/09/2025, 11-09-2025), 
      treat as MM/DD/YYYY
    - For text formats (06-Nov-2025, December 6, 2025), parse normally
    
    Returns empty string if date cannot be parsed
    """
    if not date_str or date_str == "":
        return ""

    if isinstance(date_str, str):
        date_str = date_str.strip()
        if date_str == "":
            return ""
        
        # EXTRA CLEANING: Replace various unicode spaces and clean up
        # Non-breaking space, thin space, etc. → regular space
        date_str = re.sub(r'[\u00A0\u2000-\u200B\u202F\u205F\u3000]', ' ', date_str)
        # Remove zero-width characters
        date_str = re.sub(r'[\u200B-\u200D\uFEFF]', '', date_str)
        # Normalize multiple spaces to single space
        date_str = re.sub(r'\s+', ' ', date_str).strip()

    # Clean ordinal suffixes FIRST (1st, 2nd, 3rd, 4th, 06th, etc.)
    cleaned_date = date_str
    if isinstance(date_str, str):
        # Handle ordinals: "06th December 2025" → "06 December 2025"
        # Also handles: "December 6th, 2025" → "December 6, 2025"
        cleaned_date = re.sub(r'(\d+)(st|nd|rd|th)\b', r'\1', date_str, flags=re.IGNORECASE)
        # Strip a leading/trailing weekday name: "Friday, 26 June 2026" → "26 June 2026"
        weekday_pattern = r'\b(?:Mon(?:day)?|Tue(?:s(?:day)?)?|Wed(?:nesday)?|Thu(?:rs(?:day)?)?|Fri(?:day)?|Sat(?:urday)?|Sun(?:day)?)\b\.?,?\s*'
        cleaned_date = re.sub(rf'^{weekday_pattern}', '', cleaned_date, flags=re.IGNORECASE).strip()
        cleaned_date = re.sub(rf',?\s*{weekday_pattern}$', '', cleaned_date, flags=re.IGNORECASE).strip()
        # Strip filler words: "12th of June 2026" → "12 June 2026", "the 12 June 2026" → "12 June 2026"
        cleaned_date = re.sub(r'^\bthe\b\s+', '', cleaned_date, flags=re.IGNORECASE).strip()
        cleaned_date = re.sub(r'\s+\bof\b\s+', ' ', cleaned_date, flags=re.IGNORECASE).strip()

        # Normalize mixed numeric separators: "08.06 2026" → "08-06-2026"
        # (invoice may use a different separator between day/month than between month/year)
        mixed_sep_match = re.match(r'^(\d{1,4})[.\-/\s]+(\d{1,2})[.\-/\s]+(\d{1,4})$', cleaned_date)
        if mixed_sep_match:
            cleaned_date = '-'.join(mixed_sep_match.groups())

    # Check if date is NUMERIC format (contains only digits and separators)
    # Pattern: XX/XX/XXXX, XX-XX-XXXX, XX.XX.XXXX (with 2 or 4 digit year)
    is_numeric_format = bool(re.match(r'^\d{1,2}[/\-\.]\d{1,2}[/\-\.]\d{2,4}$', cleaned_date))
    
    # US FORMAT PRIORITY: If currency is USD and date is numeric, try MM/DD/YYYY first
    if currency and currency.upper() == 'USD' and is_numeric_format:
        us_formats = [
            "%m/%d/%Y",           # 01/15/2025
            "%m-%d-%Y",           # 01-15-2025
            "%m.%d.%Y",           # 01.15.2025
            "%m/%d/%y",           # 01/15/25
            "%m-%d-%y",           # 01-15-25
            "%m.%d.%y",           # 01.15.25
        ]
        for fmt in us_formats:
            try:
                parsed_date = datetime.strptime(cleaned_date, fmt)
                return parsed_date.strftime("%d-%b-%Y")
            except (ValueError, TypeError):
                continue

    # FULL DATE FORMATS (day-month-year) - standard parsing
    full_date_formats = [
        # ISO formats (4-digit year) - these are unambiguous
        "%Y-%m-%d",           # 2025-01-15
        "%Y/%m/%d",           # 2025/01/15
        "%Y.%m.%d",           # 2025.01.15
        "%Y %m %d",           # 2025 01 15
        "%Y%m%d",             # 20250115 (compact)
        
        # European formats with full month names (4-digit year) - UNAMBIGUOUS
        "%d %B, %Y",          # 15 December, 2025 (with comma)
        "%d %b, %Y",          # 15 Dec, 2025 (with comma)
        "%d %B %Y",           # 15 January 2025
        "%d %b %Y",           # 15 Jan 2025
        "%d-%B-%Y",           # 15-January-2025
        "%d-%b-%Y",           # 15-Jan-2025
        "%d.%B.%Y",           # 15.January.2025
        "%d.%b.%Y",           # 15.Jan.2025
        "%d/%B/%Y",           # 15/January/2025
        "%d/%b/%Y",           # 15/Jan/2025
        
        # US formats with full month names (4-digit year) - UNAMBIGUOUS
        "%B %d, %Y",          # January 15, 2025
        "%b %d, %Y",          # Jan 15, 2025
        "%B %d %Y",           # January 15 2025
        "%b %d %Y",           # Jan 15 2025
        "%B-%d-%Y",           # January-15-2025
        "%b-%d-%Y",           # Jan-15-2025
        "%B %d,%Y",           # January 15,2025 (no space after comma)
        "%b %d,%Y",           # Jan 15,2025
        
        # European formats - Day first (4-digit year)
        "%d-%m-%Y",           # 15-01-2025
        "%d/%m/%Y",           # 15/01/2025
        "%d.%m.%Y",           # 15.01.2025
        "%d %m %Y",           # 15 01 2025
        
        # US formats - Month first (4-digit year) - only if not USD or not numeric
        "%m-%d-%Y",           # 01-15-2025
        "%m/%d/%Y",           # 01/15/2025
        "%m.%d.%Y",           # 01.15.2025
        "%m %d %Y",           # 01 15 2025
        
        # European formats with 2-digit year - Day first
        "%d-%m-%y",           # 15-01-25
        "%d/%m/%y",           # 15/01/25
        "%d.%m.%y",           # 15.01.25
        "%d %m %y",           # 15 01 25
        
        # US formats with 2-digit year - Month first
        "%m-%d-%y",           # 01-15-25
        "%m/%d/%y",           # 01/15/25
        "%m.%d.%y",           # 01.15.25
        "%m %d %y",           # 01 15 25
        
        # ISO with 2-digit year
        "%y-%m-%d",           # 25-01-15
        "%y/%m/%d",           # 25/01/15
        "%y.%m.%d",           # 25.01.15
        "%y %m %d",           # 25 01 15
        
        # Compact formats with 2-digit year
        "%y%m%d",             # 250115
        "%d%m%y",             # 150125
        "%m%d%y",             # 011525
        
        # European formats with abbreviated month (2-digit year) - UNAMBIGUOUS
        "%d %B, %y",          # 15 December, 25 (with comma)
        "%d %b, %y",          # 15 Dec, 25 (with comma)
        "%d-%b-%y",           # 15-Jan-25
        "%d/%b/%y",           # 15/Jan/25
        "%d.%b.%y",           # 15.Jan.25
        "%d %b %y",           # 15 Jan 25
        "%d-%B-%y",           # 15-January-25
        "%d/%B/%y",           # 15/January/25
        
        # US formats with abbreviated month (2-digit year) - UNAMBIGUOUS
        "%b %d, %y",          # Jan 15, 25
        "%b %d %y",           # Jan 15 25
        "%B %d, %y",          # January 15, 25
        "%B %d %y",           # January 15 25
        "%b-%d-%y",           # Jan-15-25
        "%B-%d-%y",           # January-15-25
        
        # Compact 8-digit formats
        "%d%m%Y",             # 15012025
        "%m%d%Y",             # 01152025
        "%Y%d%m",             # 20251501
    ]

    # Try full date formats → output as dd-MMM-yyyy
    for fmt in full_date_formats:
        try:
            parsed_date = datetime.strptime(cleaned_date, fmt)
            return parsed_date.strftime("%d-%b-%Y")
        except (ValueError, TypeError):
            continue

    # MONTH-YEAR ONLY FORMATS - output as MMM-yyyy
    month_year_formats = [
        # Full month name with year
        "%B %Y",              # August 2025
        "%b %Y",              # Aug 2025
        "%B, %Y",             # August, 2025
        "%b, %Y",             # Aug, 2025
        "%B-%Y",              # August-2025
        "%b-%Y",              # Aug-2025
        "%B/%Y",              # August/2025
        "%b/%Y",              # Aug/2025
        
        # Numeric month-year (4-digit year)
        "%m/%Y",              # 08/2025
        "%m-%Y",              # 08-2025
        "%m.%Y",              # 08.2025
        "%m %Y",              # 08 2025
        "%Y-%m",              # 2025-08
        "%Y/%m",              # 2025/08
        "%Y.%m",              # 2025.08
        "%Y %m",              # 2025 08
        
        # Numeric month-year (2-digit year)
        "%m/%y",              # 08/25
        "%m-%y",              # 08-25
        "%m.%y",              # 08.25
        "%m %y",              # 08 25
        "%y-%m",              # 25-08
        "%y/%m",              # 25/08
        
        # Full month name with 2-digit year
        "%B %y",              # August 25
        "%b %y",              # Aug 25
        "%B-%y",              # August-25
        "%b-%y",              # Aug-25
    ]

    # Try month-year formats → output as MMM-yyyy (no day)
    for fmt in month_year_formats:
        try:
            parsed_date = datetime.strptime(cleaned_date, fmt)
            return parsed_date.strftime("%b-%Y")  # Aug-2025 format
        except (ValueError, TypeError):
            continue

    # DAY-MONTH ONLY (no year printed, e.g. "8th July" → "8 July") — assume current year
    day_month_formats = [
        "%d %B", "%d %b", "%d-%B", "%d-%b", "%d/%B", "%d/%b", "%d.%B", "%d.%b",
        "%B %d", "%b %d", "%B-%d", "%b-%d", "%B/%d", "%b/%d",
    ]
    current_year = datetime.now().year
    for fmt in day_month_formats:
        try:
            parsed_date = datetime.strptime(cleaned_date, fmt).replace(year=current_year)
            return parsed_date.strftime("%d-%b-%Y")
        except (ValueError, TypeError):
            continue

    # If no format matched, return empty string
    return ""

def parse_date_to_object(date_str, currency=None):
    """
    Parse a date string to a datetime.date object for date_input widget
    Currency-aware: If USD and numeric format, treat as MM/DD/YYYY
    Returns None if date cannot be parsed
    """
    if not date_str or date_str == "":
        return None

    if isinstance(date_str, str):
        date_str = date_str.strip()
        if date_str == "":
            return None
        
        # EXTRA CLEANING: Replace various unicode spaces and clean up
        date_str = re.sub(r'[\u00A0\u2000-\u200B\u202F\u205F\u3000]', ' ', date_str)
        date_str = re.sub(r'[\u200B-\u200D\uFEFF]', '', date_str)
        date_str = re.sub(r'\s+', ' ', date_str).strip()

    # Clean ordinal suffixes FIRST (1st, 2nd, 3rd, 4th, 06th, etc.)
    cleaned_date = str(date_str)
    if isinstance(date_str, str):
        cleaned_date = re.sub(r'(\d+)(st|nd|rd|th)\b', r'\1', date_str, flags=re.IGNORECASE)
        # Strip a leading/trailing weekday name: "Friday, 26 June 2026" → "26 June 2026"
        weekday_pattern = r'\b(?:Mon(?:day)?|Tue(?:s(?:day)?)?|Wed(?:nesday)?|Thu(?:rs(?:day)?)?|Fri(?:day)?|Sat(?:urday)?|Sun(?:day)?)\b\.?,?\s*'
        cleaned_date = re.sub(rf'^{weekday_pattern}', '', cleaned_date, flags=re.IGNORECASE).strip()
        cleaned_date = re.sub(rf',?\s*{weekday_pattern}$', '', cleaned_date, flags=re.IGNORECASE).strip()
        # Strip filler words: "12th of June 2026" → "12 June 2026", "the 12 June 2026" → "12 June 2026"
        cleaned_date = re.sub(r'^\bthe\b\s+', '', cleaned_date, flags=re.IGNORECASE).strip()
        cleaned_date = re.sub(r'\s+\bof\b\s+', ' ', cleaned_date, flags=re.IGNORECASE).strip()

        # Normalize mixed numeric separators: "08.06 2026" → "08-06-2026"
        # (invoice may use a different separator between day/month than between month/year)
        mixed_sep_match = re.match(r'^(\d{1,4})[.\-/\s]+(\d{1,2})[.\-/\s]+(\d{1,4})$', cleaned_date)
        if mixed_sep_match:
            cleaned_date = '-'.join(mixed_sep_match.groups())

    # Check if date is NUMERIC format (contains only digits and separators)
    is_numeric_format = bool(re.match(r'^\d{1,2}[/\-\.]\d{1,2}[/\-\.]\d{2,4}$', cleaned_date))
    
    # US FORMAT PRIORITY: If currency is USD and date is numeric, try MM/DD/YYYY first
    if currency and currency.upper() == 'USD' and is_numeric_format:
        us_formats = [
            "%m/%d/%Y",           # 01/15/2025
            "%m-%d-%Y",           # 01-15-2025
            "%m.%d.%Y",           # 01.15.2025
            "%m/%d/%y",           # 01/15/25
            "%m-%d-%y",           # 01-15-25
            "%m.%d.%y",           # 01.15.25
        ]
        for fmt in us_formats:
            try:
                parsed_date = datetime.strptime(cleaned_date, fmt)
                return parsed_date.date()
            except (ValueError, TypeError):
                continue

    # Standard formats
    formats = [
        # ISO formats (4-digit year)
        "%Y-%m-%d", "%Y/%m/%d", "%Y.%m.%d", "%Y %m %d", "%Y%m%d",
        
        # Text month formats with comma - MUST BE FIRST for "06 December, 2025"
        "%d %B, %Y", "%d %b, %Y",          # 06 December, 2025 / 06 Dec, 2025
        
        # Text month formats - UNAMBIGUOUS
        "%d %B %Y", "%d %b %Y", "%d-%B-%Y", "%d-%b-%Y",
        "%d.%B.%Y", "%d.%b.%Y", "%d/%B/%Y", "%d/%b/%Y",
        "%B %d, %Y", "%b %d, %Y", "%B %d %Y", "%b %d %Y",
        "%B-%d-%Y", "%b-%d-%Y", "%B %d,%Y", "%b %d,%Y",
        
        # European formats - Day first
        "%d-%m-%Y", "%d/%m/%Y", "%d.%m.%Y", "%d %m %Y",
        "%d-%m-%y", "%d/%m/%y", "%d.%m.%y", "%d %m %y",
        
        # US formats - Month first
        "%m-%d-%Y", "%m/%d/%Y", "%m.%d.%Y", "%m %d %Y",
        "%m-%d-%y", "%m/%d/%y", "%m.%d.%y", "%m %d %y",
        
        # ISO with 2-digit year
        "%y-%m-%d", "%y/%m/%d", "%y.%m.%d", "%y %m %d",
        
        # Compact formats
        "%y%m%d", "%d%m%y", "%m%d%y", "%d%m%Y", "%m%d%Y", "%Y%d%m",
        
        # Text month with 2-digit year (with comma)
        "%d %B, %y", "%d %b, %y",          # 06 December, 25 / 06 Dec, 25
        "%d-%b-%y", "%d/%b/%y", "%d.%b.%y", "%d %b %y",
        "%d-%B-%y", "%d/%B/%y",
        "%b %d, %y", "%b %d %y", "%B %d, %y", "%B %d %y",
        "%b-%d-%y", "%B-%d-%y",
        
        # Month-year only
        "%B %Y", "%b %Y", "%B, %Y", "%b, %Y",
        "%B-%Y", "%b-%Y", "%B/%Y", "%b/%Y",
        "%m/%Y", "%m-%Y", "%m.%Y", "%m %Y",
        "%Y-%m", "%Y/%m", "%Y.%m", "%Y %m",
        "%m/%y", "%m-%y", "%m.%y", "%m %y",
        "%y-%m", "%y/%m",
        "%B %y", "%b %y", "%B-%y", "%b-%y",
    ]

    for fmt in formats:
        try:
            parsed_date = datetime.strptime(cleaned_date, fmt)
            return parsed_date.date()
        except (ValueError, TypeError):
            continue

    # DAY-MONTH ONLY (no year printed, e.g. "8th July" → "8 July") — assume current year
    day_month_formats = [
        "%d %B", "%d %b", "%d-%B", "%d-%b", "%d/%B", "%d/%b", "%d.%B", "%d.%b",
        "%B %d", "%b %d", "%B-%d", "%b-%d", "%B/%d", "%b/%d",
    ]
    current_year = datetime.now().year
    for fmt in day_month_formats:
        try:
            return datetime.strptime(cleaned_date, fmt).replace(year=current_year).date()
        except (ValueError, TypeError):
            continue

    return None


# =============================================================================
# vLLM Inference — MULTI-IMAGE SINGLE REQUEST
# =============================================================================
def _encode_page(page: Image.Image, page_idx: int, max_dim: int = MAX_DIMENSION_PER_PAGE):
    orig_w, orig_h = page.size

    api_img = page.convert("RGB")
    api_img = api_img.resize((max_dim, max_dim), Image.Resampling.LANCZOS)

    buf = BytesIO()
    api_img.save(buf, format="PNG", optimize=True)

    b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
    data_url = f"data:image/png;base64,{b64}"

    return data_url, (orig_w, orig_h), (max_dim, max_dim)


def run_inference_vllm(images: List[Image.Image]) -> str | None:
    EXTRACTION_PROMPT = """Extract structured data from the provided invoice image(s). Return a single JSON object with three sections: header, items, summary.
## Output Schema
Return this exact structure. Use empty string "" for any field not found.
json
{
  "header": {
    "invoice_no": "",
    "po_no": "",
    "invoice_date": "",
    "payment_terms": "",
    "due_date": "",
    "sender_name": "",
    "sender_addr": "",
    "tax_id": "",
    "rcpt_name": "",
    "rcpt_addr": "",
    "bank_iban": "",
    "bank_name": "",
    "bank_acc_no": "",
    "bank_routing": "",
    "bank_swift": "",
    "bank_acc_name": "",
    "bank_branch": ""
  },
  "items": [
    {
      "service_name": "",
      "service_start_date": "",
      "service_end_date": "",
      "descriptions": "",
      "SKU": "",
      "quantity": "",
      "unit_price": "",
      "amount": "",
      "discount_rate_per_item": "",
      "discount_amount_per_item": "",
      "tax_rate_per_item": "",
      "tax_amount_per_item": "",
      "line_total": ""
    }
  ],
  "summary": {
    "subtotal": "",
    "tax_rate": "",
    "tax_amount": "",
    "discount_rate": "",
    "total_discount_amount": "",
    "total_amount": "",
    "currency": ""
  }
}
## Core Rules
1. EXTRACT ONLY WHAT IS EXPLICITLY VISIBLE. Never compute, derive, infer totals, or fill missing values via arithmetic. The following format normalizations and explicit-data operations ARE permitted:
   - Bracketed amounts → negative: (100.00) → "-100.00".
   - Date ranges → expanded to full start/end dates using the invoice year.
   - Language normalization of payment terms → English.
   - Summing multiple tax AMOUNTS shown separately into one combined amount.
   - Summing multiple discount AMOUNTS shown separately into one combined amount.
   - Default quantity to "1.00" when unit_price and amount are both shown and equal, with no quantity printed.
   - Setting tax to "0.00" when "No VAT" / "VAT Exempt" is explicitly stated, OR when a tax rate is shown but tax amount is blank AND subtotal equals total as printed.
2. Return "" for any field not present or not clearly identifiable.
3. PRESERVE ORIGINAL FORMATTING for numbers (commas, decimals) and dates (except for date-range expansion noted above). For tax_rate and discount_rate, preserve exactly as printed including any % symbol.
4. STRIP CURRENCY SYMBOLS from all amount/price fields. Place currency code (USD, EUR, GBP, INR, etc.) only in summary.currency. If multiple currencies appear, use the currency of total_amount. If no currency is shown anywhere on the invoice, set summary.currency to "".
5. MULTI-PAGE: Merge all pages into one JSON object. Extract every line item as printed — do not deduplicate. Only skip a row if it is clearly a continuation header (e.g., the row repeats the table column titles like "Description / Qty / Price").
6. NO DUPLICATION ACROSS SCHEMA SECTIONS: If a value appears in the line items area of the invoice, place it only in the items array of the schema. If it appears in the summary/totals area (the section at the bottom showing subtotal, tax, discount, and total — separate from the itemized table), place it only in the summary object of the schema. Do not copy the same value into both unless the invoice explicitly prints it in both areas.
7. OUTPUT ONLY THE JSON. No preamble, no markdown fences, no explanation.
8. INCLUDE ALL LINE ITEMS, even those with a zero amount, zero quantity, or negative amount. Never skip a line item because its value is 0 or negative. This applies equally to credit notes and adjustments. Negative quantities, if printed, are extracted as-is (preserve the negative sign).
9. NO LINE ITEMS: If the invoice shows no individual line items (e.g., a flat-fee invoice with only a total), return an empty array: "items": [].
10. CONFLICTING VALUES: If the invoice shows two different totals (e.g., "Total" and "Amount Due"), prefer the value labeled as the final payable amount ("Amount Due", "Balance Due", "Total Payable", "Grand Total").
## Field Rules
### Sender & Recipient
1.sender_name: The invoicing company/person. Prefer company name over person name.
2.sender_addr: The full address of the invoicing entity, wherever it appears on the invoice.
3.tax_id: Always the SENDER's tax ID / VAT number. Never the recipient's.
4.rcpt_name: The name of the recipient, customer, or entity being billed.
5.rcpt_addr: The address of the recipient, customer, or entity being billed.
### PO Number
1.po_no: Extract if a single PO number is shown. If the same PO number appears multiple times on the invoice (e.g., header and footer), extract it once. If multiple DISTINCT PO numbers are present, return "".
### Bank Details
1.Prefer the bank account matching the invoice currency (USD invoice → USD account).
2.If no currency-matching account exists, extract the only available account.
3.If multiple non-matching accounts exist with no currency-matching account, return "" for all bank fields.
4.bank_acc_name: The name of the account holder. It may also appear in phrases such as “Cheques should be payable to” or “Remittance should be payable to.” Extract the mentioned company/entity name as the bank_acc_name.
### Dates
5.Preserve the original date format exactly as printed.
6.Date range exception: When a service period is shown as a range (e.g., "Feb 1-3"), expand to two full dates: service_start_date: "Feb 1 2026", service_end_date: "Feb 3 2026". Use the invoice year if the range omits a year.
7.due_date: If multiple due dates, extract only the first. If payment terms are "due on receipt" / "due on presentation" / "immediate", set due_date = invoice_date.
8.payment_terms: Normalize to English ("30 jours" → "30 days", "60 gg df" → "60 days").
### Descriptions
1.service_name: A short label, product name, or service title (e.g., "Consulting", "Web Hosting", "Premium Plan").
2.descriptions: The full detailed description text of the line item.
3.If the line item shows both a short label and a longer expanded text → put the short label in service_name and the longer text in descriptions.
4.If only one text is available for the line item → place it in descriptions and leave service_name as "".
5.If a SKU/code appears together with descriptive text (e.g., "SKU-12345 — Premium Web Hosting") → put the SKU in SKU, the descriptive portion in descriptions, and leave service_name as "" unless a separate short label also exists.
6.Do not extract handwritten annotations or stamps unless they are clearly legible and printed-quality. When in doubt, extract only the printed text.
### Amounts, Tax & Discounts
1.Preserve original number formatting (commas, decimals) exactly as printed. Do not reinterpret decimal separators — preserve as written (European 1.234,56 stays 1.234,56; US 1,234.56 stays 1,234.56).
2.Bracketed amounts are negative: (100.00) → "-100.00".
3."No VAT" / "VAT Exempt" → set relevant tax field to "0.00".
4.Tax rate shown but tax amount blank AND subtotal equals total as printed → tax amount = "0.00".
5.Tax-inclusive pricing: If the invoice states "all prices include VAT" or "tax-inclusive pricing", extract amount as printed without separating tax. Extract summary.tax_amount, summary.tax_rate, tax_rate_per_item, and tax_amount_per_item exactly as printed on the invoice — do not attempt to back-calculate.
6.Summing multiple values:
  - Multiple tax AMOUNTS shown separately → sum into one combined amount in the appropriate field.
  - Multiple DISCOUNT AMOUNTS shown separately → sum into one combined amount in the appropriate field.
  - Multiple tax RATES (e.g., 5% CGST + 12% SGST) → do NOT sum. Concatenate as printed (e.g., "5% + 12%") or extract the dominant rate as printed. Never produce a summed rate like "17%".
7.0Gross vs. Net amount columns: If both gross and net per-line amounts are shown, use the NET (pre-tax) amount as amount.
8.Discount/tax as a line item: If discount or tax appears as a row in the line items table — identified by labels such as "Discount", "Rebate", "Tax", "VAT", "Adjustment", or by being a negative-value row with no quantity/unit price — create a separate entry in the items array. Populate ONLY descriptions, the relevant discount/tax fields (discount_rate_per_item, discount_amount_per_item, tax_rate_per_item, tax_amount_per_item), and line_total. Leave quantity, unit_price, and amount as "".
9.Discount/tax in the totals/summary area: Place only in the summary object of the schema. Do not duplicate into the items array.
### Quantity & Unit Price
1.If unit_price and amount are shown and equal but no quantity is printed → quantity: "1.00".
2.If the invoice uses alternate labels for quantity (e.g., impressions, units, hours, sessions, clicks) and unit price (e.g., CPM, rate, cost per unit), map them to quantity and unit_price respectively.
3.If both a standard quantity and an alternate metric are shown, extract the standard quantity.
4.Negative quantities (e.g., on credit notes) are preserved as-is with the negative sign.
### Line Total
line_total: Extract ONLY if ALL these conditions are met:
  1. The invoice explicitly prints a dedicated line total column or value for the row.
  2. The invoice has NO separate per-line tax column or per-line tax amount anywhere.
  3. The invoice has NO separate per-line discount column or per-line discount amount anywhere.
If any per-line tax or discount column exists anywhere on the invoice → set line_total to "" for ALL rows.
If no line total column exists on the invoice → line_total = ""."""

    if not images:
        st.error("No images provided to run_inference_vllm.")
        return None

    if len(images) > MAX_PAGES_PER_REQUEST:
        st.warning(
            f"Invoice has {len(images)} pages — only the first "
            f"{MAX_PAGES_PER_REQUEST} will be sent per the request limit."
        )
        images = images[:MAX_PAGES_PER_REQUEST]

    try:
        image_content_blocks: List[dict] = []
        total_b64_bytes = 0
        resize_info     = []

        for idx, page in enumerate(images):
            data_url, orig_size, new_size = _encode_page(page, idx)
            total_b64_bytes += len(data_url)
            if orig_size != new_size:
                resize_info.append(
                    f"Page {idx+1}: {orig_size[0]}×{orig_size[1]} → {new_size[0]}×{new_size[1]}"
                )
            image_content_blocks.append({"type": "image_url", "image_url": {"url": data_url}})

        if resize_info:
            st.info("Resized for API payload:\n" + "\n".join(resize_info))

        total_payload_mb = total_b64_bytes / (1024 * 1024)
        if total_payload_mb > MAX_TOTAL_PAYLOAD_MB:
            st.warning(
                f"Large multi-page payload ({total_payload_mb:.1f} MB across "
                f"{len(images)} page(s)). This may be slow or time-out."
            )
        else:
            st.info(f"Sending {len(images)} page(s) in one request (payload ≈ {total_payload_mb:.2f} MB).")

        user_content = image_content_blocks + [
            {
                "type": "text",
                "text": (
                    f"The {len(images)} image(s) above are all pages of the same invoice "
                    f"(page 1 through page {len(images)}).\n"
                    "Extract invoice data from ALL pages combined into a single JSON object."
                )
            }
        ]

        payload = {
            "model": MODEL_NAME,
            "messages": [
                {"role": "system", "content": EXTRACTION_PROMPT},
                {"role": "user",   "content": user_content}
            ],
            "temperature": 0,
            "max_tokens":  10000
        }

        headers = {
            "Authorization": f"Bearer {VLLM_API_KEY}",
            "Content-Type":  "application/json"
        }

        response = requests.post(
            f"{POD_URL}/v1/chat/completions",
            headers=headers,
            json=payload,
            timeout=600
        )

        if response.status_code == 200:
            result = response.json()
            return result["choices"][0]["message"]["content"]
        else:
            st.error(f"❌ API Error {response.status_code}")
            try:
                st.json(response.json())
            except Exception:
                st.code(response.text)
            return None

    except Exception as e:
        st.error(f"Error calling vLLM: {str(e)}")
        return None


# -----------------------------
# JSON Parser + postprocessing mapper from code #1
# -----------------------------
def parse_vllm_json(raw_json_text):
    """Parse vLLM JSON output into:
    1) raw_data: untouched parsed JSON object
    2) mapped_data: cleaned structured format for validation
    """
    try:
        text_to_parse = raw_json_text.strip()

        if text_to_parse.startswith("```json"):
            text_to_parse = text_to_parse[7:]
        elif text_to_parse.startswith("```"):
            text_to_parse = text_to_parse[3:]
        if text_to_parse.endswith("```"):
            text_to_parse = text_to_parse[:-3]
        text_to_parse = text_to_parse.strip()

        raw_data = json.loads(text_to_parse)

        # Support both flat output and nested output
        header = raw_data.get("header", raw_data)
        summary = raw_data.get("summary", raw_data)
        items = raw_data.get("items", [])

        currency = summary.get("currency", header.get("currency", ""))

        def clean_amount(value):
            return clean_float(value, currency)

        mapped = {
            "Invoice Number": header.get("invoice_no", ""),
            "PO Number": header.get("po_no", ""),
            "Invoice Date": normalize_date(header.get("invoice_date", ""), currency),
            "Payment Terms": header.get("payment_terms", ""),
            "Due Date": normalize_date(header.get("due_date", ""), currency),
            "Sender Name": header.get("sender_name", ""),
            "Sender Address": header.get("sender_addr", ""),
            "Tax ID": header.get("tax_id", ""),
            "Sender": {
                "Name": header.get("sender_name", ""),
                "Address": header.get("sender_addr", "")
            },
            "Recipient Name": header.get("rcpt_name", ""),
            "Recipient Address": header.get("rcpt_addr", ""),
            "Recipient": {
                "Name": header.get("rcpt_name", ""),
                "Address": header.get("rcpt_addr", "")
            },
            "Bank Details": {
                "bank_iban": header.get("bank_iban", ""),
                "bank_name": header.get("bank_name", ""),
                "bank_account_number": header.get("bank_acc_no", ""),
                "bank_routing": header.get("bank_routing", ""),
                "bank_swift": header.get("bank_swift", ""),
                "bank_acc_name": header.get("bank_acc_name", ""),
                "bank_branch": header.get("bank_branch", "")
            },
            "Subtotal": clean_amount(summary.get("subtotal", "0")),
            "Tax Percentage": clean_tax_percentage(summary.get("tax_rate", "0")),
            "Total Tax": clean_amount(summary.get("tax_amount", "0")),
            "Discount Rate": clean_float(summary.get("discount_rate", "0"), currency),
            "Total Discount Amount": clean_amount(summary.get("total_discount_amount", "0")),
            "Total Amount": clean_amount(summary.get("total_amount", "0")),
            "Currency": currency,
            "Itemized Data": []
        }

        for item in items:
            raw_tax = item.get("tax_amount_per_item", item.get("tax", ""))
            raw_line_total = item.get("line_total", item.get("Line_total", ""))
            tax_amount = clean_amount(raw_tax) if raw_tax != "" else 0.0
            line_total = clean_amount(raw_line_total) if raw_line_total != "" else ""

            mapped["Itemized Data"].append({
                "Service Name": item.get("service_name", "") or item.get("Service_name", ""),
                "Service Start Date": item.get("service_start_date", ""),
                "Service End Date": item.get("service_end_date", ""),
                "Description": item.get("descriptions", ""),
                "SKU": item.get("SKU", ""),
                "Quantity": clean_quantity(item.get("quantity", "0"), currency),
                "Unit Price": clean_amount(item.get("unit_price", "0")),
                "Amount": clean_amount(item.get("amount", "0")),
                "Discount Rate Per Item": clean_float(item.get("discount_rate_per_item", "0"), currency),
                "Discount Amount Per Item": clean_amount(item.get("discount_amount_per_item", "0")),
                "Tax Rate Per Item": clean_tax_percentage(item.get("tax_rate_per_item", "0")),
                "Tax Amount Per Item": tax_amount,
                "Tax_Raw": raw_tax,
                "Tax": tax_amount,
                "Line Total": line_total,
                "line_total": raw_line_total
            })

        return raw_data, mapped

    except Exception as e:
        st.error(f"JSON parse error: {str(e)}")
        return None, None




# -----------------------------
# Postprocessing logic from code #1
# -----------------------------
def validate_and_calculate_taxes(structured_data):
    """
    Validate and, when needed, calculate per-item tax values.
    Supports both the legacy item keys:
      - Tax
      - Line Total
    and the new item keys:
      - Tax Amount Per Item
      - Tax Rate Per Item
      - Discount Amount Per Item
      - Discount Rate Per Item
      - line_total
    """

    subtotal = structured_data.get("Subtotal", 0.0)
    total_amount = structured_data.get("Total Amount", 0.0)
    model_tax_rate = structured_data.get("Tax Percentage", 0.0)
    model_tax_amount = structured_data.get("Total Tax", 0.0)
    items = structured_data.get("Itemized Data", [])

    def get_raw_tax_value(item):
        if not isinstance(item, dict):
            return ""
        if "Tax_Raw" in item:
            # Tax_Raw reflects exactly what the raw model output had for this
            # item's tax_amount_per_item — trust it even when blank, since blank
            # means "not given" (eligible for later summary-level distribution),
            # not "explicitly zero". Falling back to "Tax Amount Per Item" here
            # would wrongly read back the 0.0 default parse_vllm_json already
            # assigned for a blank field, misclassifying it as explicit zero.
            return item.get("Tax_Raw", "")
        # Legacy items with no Tax_Raw key at all — fall back to whatever tax value exists.
        val = item.get("Tax Amount Per Item", item.get("Tax", ""))
        return "" if val in ("", None) else val

    def item_discount(item):
        if not isinstance(item, dict):
            return 0.0
        try:
            return abs(clean_float(item.get("Discount Amount Per Item", 0.0)))
        except (ValueError, TypeError):
            return 0.0

    def set_item_tax_value(item, value):
        item["Tax Amount Per Item"] = value
        item["Tax"] = value

    def set_item_line_total(item, value):
        item["Line Total"] = value
        item["line_total"] = value

    # Discount to net against subtotal when reconciling the total — otherwise
    # "total = subtotal + tax" silently drops any discount the invoice applied.
    # Prefer the model's own summary figure; fall back to summing per-item
    # discounts when the summary total is blank/zero.
    model_discount_amount = abs(clean_float(structured_data.get("Total Discount Amount", 0.0)))
    items_discount_sum = sum(item_discount(item) for item in items if isinstance(item, dict))
    total_discount = model_discount_amount if model_discount_amount else items_discount_sum

    if subtotal >= total_amount or subtotal <= 0:
        structured_data["tax_validated"] = False
        structured_data["tax_skip_reason"] = "No tax detected"
        return structured_data

    if model_tax_rate > 0 and model_tax_amount == 0.0:
        structured_data["tax_validated"] = False
        structured_data["tax_skip_reason"] = "Tax rate exists but tax amount is 0"
        return structured_data

    taxable_items = []
    non_taxable_items = []
    has_blank_tax_items = False

    for item in items:
        amount = item.get("Amount", 0.0)
        raw_tax_value = get_raw_tax_value(item)
        taxable_base = amount - item_discount(item)

        if amount == 0.0:
            set_item_tax_value(item, 0.0)
            set_item_line_total(item, 0.0)
            non_taxable_items.append(item)
            continue

        is_empty = False
        is_explicitly_zero = False

        if isinstance(raw_tax_value, str):
            cleaned = raw_tax_value.strip()
            if cleaned == "":
                is_empty = True
            else:
                try:
                    cleaned_value = float(re.sub(r"[^\d\.-]", "", cleaned) or "0")
                    if cleaned_value == 0.0:
                        is_explicitly_zero = True
                except (ValueError, TypeError):
                    pass
        elif raw_tax_value is None or raw_tax_value == "":
            is_empty = True
        elif raw_tax_value == 0 or raw_tax_value == 0.0:
            is_explicitly_zero = True

        if is_empty:
            # Genuinely blank — not zero, just not given. Leave it eligible for
            # distribute_summary_tax_and_discount to fill in later, and don't let
            # its absence from taxable_items shrink the summary Total Tax below.
            has_blank_tax_items = True

        if is_empty or is_explicitly_zero:
            set_item_tax_value(item, 0.0)
            set_item_line_total(item, round(taxable_base, 2))
            non_taxable_items.append(item)
            continue

        taxable_items.append(item)

    authoritative_rate = None
    authority_source = None

    if taxable_items:
        total_taxable_amount = sum(item.get("Amount", 0.0) - item_discount(item) for item in taxable_items)
        if total_taxable_amount <= 0:
            structured_data["tax_validated"] = False
            structured_data["tax_skip_reason"] = "No taxable items with valid amounts"
            return structured_data

        if model_tax_rate > 0:
            expected_tax_from_rate = total_taxable_amount * (model_tax_rate / 100)
            expected_total_from_rate = subtotal - total_discount + expected_tax_from_rate
            error_from_rate = abs(expected_total_from_rate - total_amount)
        else:
            error_from_rate = float("inf")

        if model_tax_amount > 0:
            calculated_rate_from_amount = (model_tax_amount / total_taxable_amount) * 100
            expected_total_from_amount = subtotal - total_discount + model_tax_amount
            error_from_amount = abs(expected_total_from_amount - total_amount)
        else:
            error_from_amount = float("inf")

        if model_tax_rate > 0 or model_tax_amount > 0:
            if error_from_rate <= error_from_amount:
                authoritative_rate = round(model_tax_rate, 4)
                authority_source = "tax_rate"
            else:
                authoritative_rate = round(calculated_rate_from_amount, 4)
                authority_source = "tax_amount"
        else:
            structured_data["tax_validated"] = False
            structured_data["tax_skip_reason"] = "No tax rate or amount provided"
            return structured_data
    else:
        structured_data["tax_validated"] = False
        structured_data["tax_skip_reason"] = "No taxable items found"
        return structured_data

    calculated_total_tax = 0.0
    if authoritative_rate is not None:
        for item in taxable_items:
            amount = item.get("Amount", 0.0)
            taxable_base = amount - item_discount(item)
            # Trust the model's own per-item tax figure — authoritative_rate is
            # only used above to validate/reconcile the invoice-level totals,
            # never to overwrite a tax amount the model actually printed.
            own_tax = clean_float(item.get("Tax Amount Per Item", 0.0))
            set_item_tax_value(item, own_tax)
            calculated_total_tax += own_tax
            set_item_line_total(item, round(taxable_base + own_tax, 2))

    structured_data["Tax Percentage"] = authoritative_rate
    if not has_blank_tax_items:
        # Only finalize Total Tax / Total Amount here when every item's tax is
        # fully known. If some items are still blank, distribute_summary_tax_and_discount
        # needs the model's original total intact to compute the correct residual —
        # overwriting it here with the partial calculated_total_tax would silently
        # drop the blank items' share of the tax.
        structured_data["Total Tax"] = round(calculated_total_tax, 2)
        structured_data["Total Amount"] = round(subtotal - total_discount + calculated_total_tax, 2)
    structured_data["tax_validated"] = True
    structured_data["tax_authority_source"] = authority_source
    structured_data["original_tax_rate"] = model_tax_rate
    structured_data["original_tax_amount"] = model_tax_amount

    return structured_data

from copy import deepcopy

def recalc_display_totals_from_items(invoice_data):
    """
    UI/download-only recalculation.
    subtotal = sum(Amount)
    Line Total = Amount + Discount Amount Per Item + Tax Amount Per Item
    """
    data = deepcopy(invoice_data or {})
    items = data.get("Itemized Data", []) or []

    subtotal = 0.0
    for item in items:
        if not isinstance(item, dict):
            continue

        amount = clean_float(item.get("Amount", 0))
        discount = clean_float(item.get("Discount Amount Per Item", 0))
        tax = clean_float(item.get("Tax Amount Per Item", 0))

        item["Line Total"] = round(amount - abs(discount) + tax, 2)
        subtotal += amount

    data["Subtotal"] = round(subtotal, 2)
    return data

from datetime import datetime, timedelta
import re
def derive_due_date_ui(invoice_date_str, payment_terms, due_date_str):
    """
    UI-only due date post-processing.
    Rules:
    1) If due date already exists, keep it as-is.
    2) If only due date is missing and payment terms exist, derive due date from payment terms.
    3) For immediate/receipt/presentation terms, due date = invoice date.
    4) Never modifies raw model output.
    """

    # Rule 1: if due date is already present, do not touch it
    if due_date_str is not None and str(due_date_str).strip() != "":
        return due_date_str

    # Rule 2: if payment terms are missing, no derivation
    if payment_terms is None or str(payment_terms).strip() == "":
        return ""

    # Parse invoice date using your existing parser
    invoice_date_obj = parse_date_to_object(invoice_date_str)
    if not invoice_date_obj:
        return ""

    terms = str(payment_terms).lower().strip()

    # Immediate / payable-on-receipt style terms
    immediate_phrases = [
        "due upon receipt",
        "upon receipt",
        "immediate",
        "immediately",
        "due on presentation",
        "payable on receipt",
        "payment on receipt",
        "cash on delivery",
        "cod",
        "on receipt",
    ]

    if any(phrase in terms for phrase in immediate_phrases):
        return invoice_date_obj.strftime("%d-%b-%Y")

    # Extract day count from common term patterns
    day_patterns = [
        r"\bnet\s*(\d+)\b",              # Net 30
        r"\bwithin\s*(\d+)\s*days?\b",   # within 30 days
        r"\bdue\s*(\d+)\s*days?\b",      # due 30 days
        r"\b(\d+)\s*days?\b",            # 30 days
        r"\b(\d+)\s*d\b",                # 30d
    ]

    for pattern in day_patterns:
        match = re.search(pattern, terms)
        if match:
            days = int(match.group(1))
            calculated = invoice_date_obj + timedelta(days=days)
            return calculated.strftime("%d-%b-%Y")

    return ""

from copy import deepcopy

def build_display_data_from_raw(raw_data: dict, mapped_data: dict) -> dict:
    """
    UI-only post-processing based on RAW model output.
    - Does not modify raw_data
    - Fills only if the raw field is truly empty
    - Treats 0 / "0" / "0.0" as real values, not blank
    """
    data = deepcopy(mapped_data or {})
    raw_data = raw_data or {}

    def is_blank_raw(v):
        return v is None or (isinstance(v, str) and v.strip() == "")

    def num(v):
        try:
            if v is None or (isinstance(v, str) and v.strip() == ""):
                return None
            return float(clean_float(v))
        except Exception:
            return None

    header = raw_data.get("header", raw_data)
    summary = raw_data.get("summary", raw_data)
    raw_items = raw_data.get("items", []) or []
    display_items = deepcopy(data.get("Itemized Data", []) or [])

    if is_blank_raw(data.get("Invoice Number")):
        data["Invoice Number"] = (
            header.get("invoice_no")
            or raw_data.get("invoice_no")
            or raw_data.get("invoice_no".lower())
            or data.get("Invoice Number", "")
        )

    if is_blank_raw(data.get("PO Number")):
        data["PO Number"] = header.get("po_no") or raw_data.get("po_no") or data.get("PO Number", "")

    # ---- Header-level ----
    if is_blank_raw(header.get("due_date")):
        if is_blank_raw(data.get("Due Date")):
            derived_due = derive_due_date_ui(
                data.get("Invoice Date", ""),
                data.get("Payment Terms", ""),
                ""
            )
            if not is_blank_raw(derived_due):
                data["Due Date"] = derived_due

    # ---- Line items ----
    for idx, raw_item in enumerate(raw_items):
        if idx >= len(display_items):
            display_items.append({})

        item = display_items[idx]
        raw_item = raw_item or {}

        qty = num(item.get("Quantity"))
        unit_price = num(item.get("Unit Price"))
        amount = num(item.get("Amount"))
        disc_rate = num(item.get("Discount Rate Per Item"))
        disc_amt = num(item.get("Discount Amount Per Item"))
        tax_rate = num(item.get("Tax Rate Per Item"))
        tax_amt = num(item.get("Tax Amount Per Item"))

        # Fill only when RAW field is blank
        if is_blank_raw(raw_item.get("amount")) and qty is not None and unit_price is not None:
            item["Amount"] = round(qty * unit_price, 2)
            amount = num(item.get("Amount"))

        if is_blank_raw(raw_item.get("quantity")) and amount is not None and unit_price not in (None, 0):
            item["Quantity"] = round(amount / unit_price, 6)
            qty = num(item.get("Quantity"))

        if is_blank_raw(raw_item.get("unit_price")) and amount is not None and qty not in (None, 0):
            item["Unit Price"] = round(amount / qty, 6)
            unit_price = num(item.get("Unit Price"))

        # Discount
        if is_blank_raw(raw_item.get("discount_amount_per_item")) and disc_rate is not None and amount is not None:
            item["Discount Amount Per Item"] = round(amount * (disc_rate / 100.0), 2)
            disc_amt = num(item.get("Discount Amount Per Item"))

        if is_blank_raw(raw_item.get("discount_rate_per_item")) and disc_amt is not None and amount not in (None, 0):
            item["Discount Rate Per Item"] = round((disc_amt / amount) * 100.0, 4)
            disc_rate = num(item.get("Discount Rate Per Item"))

        # Tax
        taxable_base = None if amount is None else amount - abs(disc_amt if disc_amt is not None else 0.0)

        if is_blank_raw(raw_item.get("tax_amount_per_item")) and taxable_base is not None and tax_rate is not None:
            item["Tax Amount Per Item"] = round(taxable_base * (tax_rate / 100.0), 2)
            tax_amt = num(item.get("Tax Amount Per Item"))

        if is_blank_raw(raw_item.get("tax_rate_per_item")) and taxable_base not in (None, 0) and tax_amt is not None:
            item["Tax Rate Per Item"] = round((tax_amt / taxable_base) * 100.0, 4)

        # Line total
        if is_blank_raw(raw_item.get("line_total")):
            amount = num(item.get("Amount"))
            disc_amt = num(item.get("Discount Amount Per Item")) or 0.0
            tax_amt = num(item.get("Tax Amount Per Item")) or 0.0
            if amount is not None:
                item["Line Total"] = round(amount - abs(disc_amt) + tax_amt, 2)

    data["Itemized Data"] = display_items

    # ---- Summary ----
    amounts = []
    discount_amounts = []
    tax_amounts = []
    taxable_bases = []

    for item in display_items:
        if not isinstance(item, dict):
            continue
        a = num(item.get("Amount"))
        d = num(item.get("Discount Amount Per Item")) or 0.0
        t = num(item.get("Tax Amount Per Item")) or 0.0

        if a is not None:
            amounts.append(a)
            taxable_bases.append(a - d)

        if num(item.get("Discount Amount Per Item")) is not None:
            discount_amounts.append(d)

        if num(item.get("Tax Amount Per Item")) is not None:
            tax_amounts.append(t)

    subtotal_calc = round(sum(amounts), 2) if amounts else None
    total_discount_calc = round(sum(discount_amounts), 2) if discount_amounts else None
    total_tax_calc = round(sum(tax_amounts), 2) if tax_amounts else None
    taxable_base_total = round(sum(taxable_bases), 2) if taxable_bases else None

    if is_blank_raw(summary.get("subtotal")) and subtotal_calc is not None:
        data["Subtotal"] = subtotal_calc

    if is_blank_raw(summary.get("total_discount_amount")) and total_discount_calc is not None:
        data["Total Discount Amount"] = total_discount_calc

    if is_blank_raw(summary.get("tax_amount")) and total_tax_calc is not None:
        data["Total Tax"] = total_tax_calc

    if is_blank_raw(summary.get("discount_rate")):
        subtotal_for_rate = num(data.get("Subtotal"))
        total_discount_for_rate = num(data.get("Total Discount Amount"))
        if subtotal_for_rate not in (None, 0) and total_discount_for_rate is not None:
            data["Discount Rate"] = round((total_discount_for_rate / subtotal_for_rate) * 100.0, 4)

    if is_blank_raw(summary.get("tax_rate")):
        total_tax_for_rate = num(data.get("Total Tax"))
        if taxable_base_total not in (None, 0) and total_tax_for_rate is not None:
            data["Tax Percentage"] = round((total_tax_for_rate / taxable_base_total) * 100.0, 4)

    if is_blank_raw(summary.get("total_amount")):
        subtotal_for_total = num(data.get("Subtotal"))
        total_discount_for_total = num(data.get("Total Discount Amount")) or 0.0
        total_tax_for_total = num(data.get("Total Tax")) or 0.0
        if subtotal_for_total is not None:
            data["Total Amount"] = round(subtotal_for_total - total_discount_for_total + total_tax_for_total, 2)

    return data

from decimal import Decimal
import pandas as pd

def fill_line_item_missing_fields_ui(df: pd.DataFrame, raw_items: list = None, summary_total_discount_amount=0.0) -> pd.DataFrame:
    """
    UI-only line-item post-processing.
    Fills ONLY empty fields. Never overwrites existing values.
    """
    def is_raw_blank(raw_item: dict, field: str) -> bool:
        v = raw_item.get(field, None)
        if v is None:
            return True
        if isinstance(v, str) and v.strip() == "":
            return True
        return False
    out = df.copy()
    summary_discount_present = not (
        summary_total_discount_amount is None
        or str(summary_total_discount_amount).strip() == ""
        or clean_float(summary_total_discount_amount) == 0.0
    )

    def is_blank(v):
        if v is None:
            return True
        try:
            return pd.isna(v) or str(v).strip() == ""
        except Exception:
            return False

    def is_missing_numeric(v):
        if v is None:
            return True
    
        try:
            if pd.isna(v):
                return True
        except Exception:
            pass
    
        if str(v).strip() == "":
            return True
    
        try:
            return clean_float(v) == 0.0
        except Exception:
            return False

    def num(v):
        if is_blank(v):
            return None
        try:
            return float(clean_float(v))
        except Exception:
            return None

    for idx in out.index:
        # Prefer each row's own attached raw data (set by detect_and_reclassify_pseudo_lines)
        # over positional raw_items[idx] lookup, which misaligns once any earlier
        # item has been dropped as a pseudo-line.
        embedded_raw = out.at[idx, "_raw_item"] if "_raw_item" in out.columns else None
        raw_item = embedded_raw if isinstance(embedded_raw, dict) else (raw_items[idx] if raw_items and idx < len(raw_items) else {})
        amount = num(out.at[idx, "Amount"]) if "Amount" in out.columns else None
        quantity = num(out.at[idx, "Quantity"]) if "Quantity" in out.columns else None
        unit_price = num(out.at[idx, "Unit Price"]) if "Unit Price" in out.columns else None
        disc_amt = num(out.at[idx, "Discount Amount Per Item"]) if "Discount Amount Per Item" in out.columns else None
        disc_rate = num(out.at[idx, "Discount Rate Per Item"]) if "Discount Rate Per Item" in out.columns else None
        tax_rate = num(out.at[idx, "Tax Rate Per Item"]) if "Tax Rate Per Item" in out.columns else None
        tax_amt = num(out.at[idx, "Tax Amount Per Item"]) if "Tax Amount Per Item" in out.columns else None
        line_total = num(out.at[idx, "Line Total"]) if "Line Total" in out.columns else None

        # amount = quantity × unit_price
        if "Amount" in out.columns and is_raw_blank(raw_item, "amount"):
            if quantity is not None and unit_price not in (None, 0):
                amount = round(quantity * unit_price, 2)
                out.at[idx, "Amount"] = amount

        # quantity = amount ÷ unit_price
        if "Quantity" in out.columns and is_raw_blank(raw_item, "quantity"):
            if amount is not None and unit_price not in (None, 0):
                quantity = round(amount / unit_price, 6)
                out.at[idx, "Quantity"] = quantity

        # unit_price = amount ÷ quantity
        if "Unit Price" in out.columns and is_raw_blank(raw_item, "unit_price"):
            if amount is not None and quantity not in (None, 0):
                unit_price = round(amount / quantity, 6)
                out.at[idx, "Unit Price"] = unit_price

        # refresh numbers after possible fills
        amount = num(out.at[idx, "Amount"]) if "Amount" in out.columns else amount
        quantity = num(out.at[idx, "Quantity"]) if "Quantity" in out.columns else quantity
        unit_price = num(out.at[idx, "Unit Price"]) if "Unit Price" in out.columns else unit_price

        # discount_amount_per_item = amount × discount_rate_per_item ÷ 100
        if "Discount Amount Per Item" in out.columns and is_blank(out.at[idx, "Discount Amount Per Item"]):
            if disc_rate is not None and amount is not None:
                disc_amt = round(amount * (disc_rate / 100.0), 2)
                out.at[idx, "Discount Amount Per Item"] = disc_amt
            elif not summary_discount_present:
                # no discount signal -> no discount
                disc_amt = 0.0
                out.at[idx, "Discount Amount Per Item"] = 0.0

        # discount_rate_per_item = (discount_amount_per_item ÷ amount) × 100
        if "Discount Rate Per Item" in out.columns and is_raw_blank(raw_item, "discount_rate_per_item"):
            if disc_amt is not None and amount not in (None, 0):
                disc_rate = round((disc_amt / amount) * 100.0, 4)
                out.at[idx, "Discount Rate Per Item"] = disc_rate

        # taxable_base = amount − discount_amount_per_item
        amount = num(out.at[idx, "Amount"]) if "Amount" in out.columns else amount
        disc_amt = num(out.at[idx, "Discount Amount Per Item"]) if "Discount Amount Per Item" in out.columns else disc_amt
        taxable_base = None if amount is None else amount - abs(disc_amt if disc_amt is not None else 0.0)
        # tax_amount_per_item = taxable_base × tax_rate_per_item ÷ 100
        if "Tax Amount Per Item" in out.columns and is_blank(out.at[idx, "Tax Amount Per Item"]):
            if taxable_base is not None and tax_rate is not None:
                tax_amt = round(taxable_base * (tax_rate / 100.0), 2)
                out.at[idx, "Tax Amount Per Item"] = tax_amt

        # tax_rate_per_item = (tax_amount_per_item ÷ taxable_base) × 100
        if "Tax Rate Per Item" in out.columns and is_raw_blank(raw_item, "tax_rate_per_item"):
            tax_amt = num(out.at[idx, "Tax Amount Per Item"])
            if taxable_base not in (None, 0) and tax_amt is not None:
                tax_rate = round((tax_amt / taxable_base) * 100.0, 4)
                out.at[idx, "Tax Rate Per Item"] = tax_rate
        # refresh tax_rate after possible back-calculation
        tax_rate = num(out.at[idx, "Tax Rate Per Item"]) if "Tax Rate Per Item" in out.columns else tax_rate

        # Line Total only if blank
        if "Line Total" in out.columns and is_raw_blank(raw_item, "line_total"):
            amount = num(out.at[idx, "Amount"])
            disc_amt = num(out.at[idx, "Discount Amount Per Item"]) or 0.0
            tax_amt = num(out.at[idx, "Tax Amount Per Item"]) or 0.0
            if amount is not None:
                out.at[idx, "Line Total"] = round(amount - abs(disc_amt) + tax_amt, 2)

    return out

def detect_and_reclassify_pseudo_lines(display_data: dict, raw_parsed_data: dict) -> dict:
    """
    Detects pseudo-lines based on field patterns, not description.
    
    Pseudo-line conditions:
    1. quantity is 0 or 1, AND unit_price is blank, AND amount is blank
    2. discount fields (rate or amount) OR tax fields (rate or amount) is filled, along with line_total
    3. Promotes value to summary, removes pseudo-line from line items
    """
    from copy import deepcopy

    data = deepcopy(display_data)
    items = data.get("Itemized Data", []) or []
    raw_items = (raw_parsed_data or {}).get("items", []) or []

    def is_blank_val(v):
        if v is None:
            return True
        if isinstance(v, str) and v.strip() == "":
            return True
        try:
            return clean_float(v) == 0.0
        except (ValueError, TypeError):
            return True

    def has_value(v):
        return not is_blank_val(v)

    real_items = []
    for idx, item in enumerate(items):
        raw_item = raw_items[idx] if idx < len(raw_items) else {}

        # Condition 1: quantity is 0 or 1, unit_price blank, amount blank
        qty = item.get("Quantity")
        unit_price = item.get("Unit Price")
        amount = item.get("Amount")

        try:
            qty_val = float(qty) if qty is not None else 0.0
        except (ValueError, TypeError):
            qty_val = 0.0

        qty_ok = qty_val in (0.0, 1.0)
        unit_price_blank = raw_item.get("unit_price", "") in ("", None) or str(raw_item.get("unit_price", "")).strip() == ""
        amount_blank = raw_item.get("amount", "") in ("", None) or str(raw_item.get("amount", "")).strip() == ""

        if not (qty_ok and unit_price_blank and amount_blank):
            # Keep a stable reference to this item's own raw data — once any
            # earlier item is dropped as a pseudo-line, positional indexing into
            # raw_items would otherwise misalign for every item after it.
            item["_raw_item"] = raw_item
            real_items.append(item)
            continue

        # Condition 2: discount or tax fields filled along with line_total
        disc_rate = item.get("Discount Rate Per Item")
        disc_amt = item.get("Discount Amount Per Item")
        tax_rate = item.get("Tax Rate Per Item")
        tax_amt = item.get("Tax Amount Per Item")

        raw_disc_amt = raw_item.get("discount_amount_per_item", "")
        raw_disc_rate = raw_item.get("discount_rate_per_item", "")
        raw_tax_amt = raw_item.get("tax_amount_per_item", "")
        raw_tax_rate = raw_item.get("tax_rate_per_item", "")

        def raw_has_value(v):
            if v is None or str(v).strip() == "":
                return False
            try:
                return clean_float(v) != 0.0
            except (ValueError, TypeError):
                return False

        has_discount = raw_has_value(raw_disc_amt) or raw_has_value(raw_disc_rate)
        has_tax = raw_has_value(raw_tax_amt) or raw_has_value(raw_tax_rate)

        if not (has_discount or has_tax):
            item["_raw_item"] = raw_item
            real_items.append(item)
            continue

        # It's a pseudo-line — promote to summary
        # Condition 3: route to correct summary field, don't override existing

        if has_discount:
            existing_disc = clean_float(data.get("Total Discount Amount", 0))
            if existing_disc == 0.0:
                if raw_has_value(raw_disc_amt):
                    pseudo_disc = clean_float(disc_amt)
                elif raw_has_value(raw_disc_rate) and has_value(item.get("Line Total")):
                    pseudo_disc = clean_float(item.get("Line Total"))
                else:
                    pseudo_disc = 0.0
                if pseudo_disc != 0.0:
                    data["Total Discount Amount"] = pseudo_disc

        if has_tax:
            existing_tax = clean_float(data.get("Total Tax", 0))
            if existing_tax == 0.0:
                if raw_has_value(raw_tax_amt):
                    pseudo_tax = clean_float(tax_amt)
                elif raw_has_value(raw_tax_rate) and has_value(item.get("Line Total")):
                    pseudo_tax = clean_float(item.get("Line Total"))
                else:
                    pseudo_tax = 0.0
                if pseudo_tax != 0.0:
                    data["Total Tax"] = pseudo_tax
    if len(real_items) < len(items):
        total_amount = clean_float(data.get("Total Amount", 0))
        sum_line_totals = sum(clean_float(item.get("Line Total", 0)) for item in items)

        if abs(sum_line_totals - total_amount) < 0.01:
            for item in real_items:
                amount = clean_float(item.get("Amount", 0))
                if amount != 0:
                    # Clear line total so distribute_summary_tax_and_discount recalculates it
                    item["Line Total"] = ""
                    item["line_total"] = ""
    data["Itemized Data"] = real_items
    return data

def distribute_summary_tax_and_discount(display_data: dict, raw_parsed_data: dict) -> dict:
    """
    UI-only post-processing: proportionally distributes summary-level tax_amount
    and total_discount_amount down to line items — only when per-line values are missing.
    Rules:
    - NULL (blank) tax_rate_per_item → line is included in taxable pool
    - 0 tax_rate_per_item → line is excluded, tax_amount_per_item set to 0 explicitly
    - Negative-amount lines (credit notes) → included with signed allocation
    - Residue after rounding → added to line with largest abs(taxable_base)
    - Anomaly: summary_tax > 0 but total_taxable_base == 0 → flag for review
    - Discount: same logic, weight = line.amount
    """
    
    from decimal import Decimal, ROUND_HALF_UP
    from copy import deepcopy
    
    data = deepcopy(display_data)
    items = data.get("Itemized Data", []) or []
    if not items:
        return data

    raw_items = (raw_parsed_data or {}).get("items", []) or []

    summary_tax = Decimal(str(clean_float(data.get("Total Tax", 0))))
    summary_discount = Decimal(str(clean_float(data.get("Total Discount Amount", 0))))

    def is_raw_blank(raw_item: dict, field: str) -> bool:
        """True if the model returned blank/null for this field."""
        v = raw_item.get(field, None)
        if v is None:
            return True
        if isinstance(v, str) and v.strip() == "":
            return True
        return False

    def is_raw_explicitly_zero(raw_item: dict, field: str) -> bool:
        """True if the model explicitly returned 0 for this field."""
        v = raw_item.get(field, None)
        if v is None:
            return False
        if isinstance(v, str) and v.strip() == "":
            return False
        try:
            return clean_float(v) == 0.0
        except (ValueError, TypeError):
            return False

    def get_raw_item_for(item, idx):
        """Each item's own raw data — prefers the stable reference detect_and_reclassify_pseudo_lines
        attaches (_raw_item), since positional raw_items[idx] lookup breaks once
        any earlier item has been dropped as a pseudo-line."""
        if isinstance(item, dict) and "_raw_item" in item:
            return item["_raw_item"]
        return raw_items[idx] if idx < len(raw_items) else {}

    # DISCOUNT DISTRIBUTION
    # Only the RESIDUAL beyond what's already explicitly known per-line gets distributed —
    # otherwise a summary total that was itself derived by summing known line discounts
    # (raw summary field blank) would get dumped again onto whichever lines had none.
    known_discount_sum = sum(
        Decimal(str(clean_float(items[idx].get("Discount Amount Per Item", 0))))
        for idx in range(len(items))
        if not is_raw_blank(get_raw_item_for(items[idx], idx), "discount_amount_per_item")
    )
    residual_discount = summary_discount - known_discount_sum

    if residual_discount != Decimal("0"):
        # Identify lines where discount_amount_per_item is blank in raw output
        dist_candidates = []
        for idx, item in enumerate(items):
            raw_item = get_raw_item_for(item, idx)
            if is_raw_blank(raw_item, "discount_amount_per_item"):
                amount = Decimal(str(clean_float(item.get("Amount", 0))))
                dist_candidates.append((idx, amount))

        total_weight = sum(abs(a) for _, a in dist_candidates)

        if total_weight > Decimal("0"):
            allocated = Decimal("0")
            residue_idx = None
            largest_weight = Decimal("0")

            for idx, amount in dist_candidates:
                weight = abs(amount)
                share = (residual_discount * amount / total_weight).quantize(
                    Decimal("0.01"), rounding=ROUND_HALF_UP
                )
                items[idx]["Discount Amount Per Item"] = float(share)
                allocated += share

                if weight > largest_weight:
                    largest_weight = weight
                    residue_idx = idx

            # Residue correction
            residue = residual_discount - allocated
            if residue != Decimal("0") and residue_idx is not None:
                current = Decimal(str(clean_float(items[residue_idx].get("Discount Amount Per Item", 0))))
                items[residue_idx]["Discount Amount Per Item"] = float(current + residue)
    
    # TAX DISTRIBUTION
    # Only the RESIDUAL beyond what's already explicitly known per-line gets distributed —
    # same reasoning as the discount fix above.
    known_tax_sum = sum(
        Decimal(str(clean_float(items[idx].get("Tax Amount Per Item", 0))))
        for idx in range(len(items))
        if not is_raw_blank(get_raw_item_for(items[idx], idx), "tax_amount_per_item")
    )
    residual_tax = summary_tax - known_tax_sum

    if residual_tax == Decimal("0"):
        # Still calculate line totals that are missing — but never override a
        # Line Total the model actually printed. Pseudo-line reclassification
        # above may have deliberately cleared it, which is the one case where
        # recomputing here is allowed even though raw wasn't blank.
        for idx, item in enumerate(items):
            if not isinstance(item, dict):
                continue
            raw_item = get_raw_item_for(item, idx)
            amount = float(item.get("Amount", 0) or 0)
            disc = abs(float(item.get("Discount Amount Per Item", 0) or 0))
            tax = float(item.get("Tax Amount Per Item", 0) or 0)
            if amount != 0 and (is_raw_blank(raw_item, "line_total") or is_blank_numeric_value(item.get("Line Total"))):
                item["Line Total"] = round(amount - disc + tax, 2)
        data["Itemized Data"] = items
        return data

    # Mark confirmed-zero-tax lines first (tax_rate_per_item == 0 in raw)
    for idx, item in enumerate(items):
        raw_item = get_raw_item_for(item, idx)
        if is_raw_explicitly_zero(raw_item, "tax_rate_per_item"):
            items[idx]["Tax Amount Per Item"] = 0.0

    # Build taxable pool: lines where tax_rate_per_item is NULL in raw output
    # and tax_amount_per_item is also blank in raw output
    taxable_pool = []
    for idx, item in enumerate(items):
        raw_item = get_raw_item_for(item, idx)

        rate_is_blank = is_raw_blank(raw_item, "tax_rate_per_item")
        tax_amt_is_blank = is_raw_blank(raw_item, "tax_amount_per_item")

        if not tax_amt_is_blank:
            # Model gave something for this line — don't touch it
            continue

        amount = Decimal(str(clean_float(item.get("Amount", 0))))
        disc = Decimal(str(clean_float(item.get("Discount Amount Per Item", 0))))
        taxable_base = amount - abs(disc)
        taxable_pool.append((idx, taxable_base))

    # Total taxable base using abs() for weight calculation
    total_taxable_base = sum(abs(base) for _, base in taxable_pool)

    # Anomaly guard
    if total_taxable_base == Decimal("0"):
        data["_tax_distribution_anomaly"] = (
            "summary_tax > 0 but total_taxable_base == 0 — review required"
        )
        data["Itemized Data"] = items
        return data

    # Proportional distribution
    allocated = Decimal("0")
    residue_idx = None
    largest_abs_base = Decimal("0")

    for idx, taxable_base in taxable_pool:
        # Weight = abs(base), but share is signed
        share = (residual_tax * taxable_base / total_taxable_base).quantize(
            Decimal("0.01"), rounding=ROUND_HALF_UP
        )
        items[idx]["Tax Amount Per Item"] = float(share)
        allocated += share

        if abs(taxable_base) > largest_abs_base:
            largest_abs_base = abs(taxable_base)
            residue_idx = idx

    # Residue correction — guarantees SUM == residual_tax exactly
    residue = residual_tax - allocated
    if residue != Decimal("0") and residue_idx is not None:
        current = Decimal(str(clean_float(items[residue_idx].get("Tax Amount Per Item", 0))))
        items[residue_idx]["Tax Amount Per Item"] = float(current + residue)

    # Same rule as above: only fill Line Total when raw was blank (or was
    # deliberately cleared by pseudo-line reclassification) — never override
    # a value the model actually printed.
    for idx, item in enumerate(items):
        if not isinstance(item, dict):
            continue
        raw_item = get_raw_item_for(item, idx)
        amount = float(item.get("Amount", 0) or 0)
        disc = abs(float(item.get("Discount Amount Per Item", 0) or 0))
        tax = float(item.get("Tax Amount Per Item", 0) or 0)
        if amount != 0 and (is_raw_blank(raw_item, "line_total") or is_blank_numeric_value(item.get("Line Total"))):
            item["Line Total"] = round(amount - disc + tax, 2)

    data["Itemized Data"] = items
    return data


def prepare_processed_invoice_data(raw_parsed_data: dict, edited_data: dict) -> dict:
    """Apply code #1's UI/download postprocessing sequence to a mapped invoice."""
    processed = build_display_data_from_raw(raw_parsed_data or {}, edited_data or {})
    processed = detect_and_reclassify_pseudo_lines(processed, raw_parsed_data or {})
    processed = distribute_summary_tax_and_discount(processed, raw_parsed_data or {})

    items_df = pd.DataFrame(processed.get("Itemized Data", []) or [])
    if not items_df.empty:
        items_df = fill_line_item_missing_fields_ui(
            items_df,
            raw_items=(raw_parsed_data or {}).get("items", []),
            summary_total_discount_amount=processed.get("Total Discount Amount", 0.0),
        )
        processed["Itemized Data"] = items_df.to_dict("records")

    items_list = processed.get("Itemized Data", []) or []
    total_amount_sum = sum(clean_float(i.get("Amount", 0)) for i in items_list if isinstance(i, dict))
    total_tax_sum = sum(clean_float(i.get("Tax Amount Per Item", 0)) for i in items_list if isinstance(i, dict))
    total_disc_sum = sum(clean_float(i.get("Discount Amount Per Item", 0)) for i in items_list if isinstance(i, dict))
    taxable_base = total_amount_sum - abs(total_disc_sum)

    raw_summary_for_tax = (raw_parsed_data or {}).get("summary", raw_parsed_data or {})
    raw_tax_rate = raw_summary_for_tax.get("tax_rate")
    raw_tax_rate_blank = raw_tax_rate is None or (isinstance(raw_tax_rate, str) and raw_tax_rate.strip() == "")

    # Only derive/recompute Tax Percentage when the model didn't print one — never override a real extracted value.
    if raw_tax_rate_blank:
        if taxable_base > 0 and total_tax_sum != 0:
            processed["Tax Percentage"] = round((total_tax_sum / taxable_base) * 100, 4)
        elif clean_float(processed.get("Tax Percentage", 0)) == 0:
            tax_rates = [
                clean_float(i.get("Tax Rate Per Item", 0))
                for i in items_list
                if isinstance(i, dict) and clean_float(i.get("Tax Rate Per Item", 0)) != 0
            ]
            if tax_rates:
                processed["Tax Percentage"] = round(sum(tax_rates) / len(tax_rates), 4)

    raw_discount_rate = raw_summary_for_tax.get("discount_rate")
    raw_discount_rate_blank = raw_discount_rate is None or (isinstance(raw_discount_rate, str) and raw_discount_rate.strip() == "")

    # Same rule for Discount Rate — only derive it when the model didn't print one.
    if raw_discount_rate_blank and total_amount_sum > 0 and total_disc_sum != 0:
        processed["Discount Rate"] = round((abs(total_disc_sum) / total_amount_sum) * 100, 4)

    return limit_invoice_decimals(processed)

def display_data_for_result(result: dict) -> dict:
    """Use saved user edits as authoritative; otherwise run initial postprocessing."""
    result = result or {}
    edited = deepcopy(result.get("edited_data", {}) or {})
    if result.get("user_saved_edits"):
        return limit_invoice_decimals(edited)
    return prepare_processed_invoice_data(result.get("raw_parsed_data", {}), edited)

def flatten_invoice_to_rows(invoice_data) -> list:
    EXPECTED_BANK_FIELDS = [
        "bank_name", "bank_account_number", "bank_acc_name",
        "bank_iban", "bank_swift", "bank_routing", "bank_branch"
    ]

    def fmt_text(value):
        if value is None or str(value).strip() == "":
            return "NA"
        return str(value).strip()

    def fmt_amount(value):
        if value is None or (isinstance(value, str) and value.strip() == ""):
            return 0
        return limit_decimal_places(value)

    rows = []
    invoice_data = invoice_data or {}
    line_items = invoice_data.get("Itemized Data", []) or []

    bank_details = {}
    nested = invoice_data.get("Bank Details", {}) or {}
    if isinstance(nested, dict):
        for k, v in nested.items():
            key_name = k if str(k).startswith("bank_") else f"bank_{k}"
            bank_details[key_name] = v
    for k, v in invoice_data.items():
        if isinstance(k, str) and k.lower().startswith("bank_"):
            bank_details[k] = v
    for f in EXPECTED_BANK_FIELDS:
        bank_details.setdefault(f, "")

    def base_invoice_info():
        return {
            "Invoice Number":        fmt_text(invoice_data.get("Invoice Number", "")),
            "PO Number":             fmt_text(invoice_data.get("PO Number", "")),
            "Invoice Date":          fmt_text(invoice_data.get("Invoice Date", "")),
            "Payment Terms":         fmt_text(invoice_data.get("Payment Terms", "")),
            "Due Date":              fmt_text(invoice_data.get("Due Date", "")),
            "Currency":              fmt_text(invoice_data.get("Currency", "")),
            "Tax ID":                fmt_text(invoice_data.get("Tax ID", "")),
            "Subtotal":              fmt_amount(invoice_data.get("Subtotal", 0.0)),
            "Tax Percentage":        fmt_amount(invoice_data.get("Tax Percentage", 0.0)),
            "Total Tax":             fmt_amount(invoice_data.get("Total Tax", 0.0)),
            "Discount Rate":         fmt_amount(invoice_data.get("Discount Rate", 0.0)),
            "Total Discount Amount": fmt_amount(invoice_data.get("Total Discount Amount", 0.0)),
            "Total Amount":          fmt_amount(invoice_data.get("Total Amount", 0.0)),
            "Sender Name":           fmt_text(invoice_data.get("Sender Name", "") or (invoice_data.get("Sender", {}) or {}).get("Name", "")),
            "Sender Address":        fmt_text(invoice_data.get("Sender Address", "") or (invoice_data.get("Sender", {}) or {}).get("Address", "")),
            "Recipient Name":        fmt_text(invoice_data.get("Recipient Name", "") or (invoice_data.get("Recipient", {}) or {}).get("Name", "")),
            "Recipient Address":     fmt_text(invoice_data.get("Recipient Address", "") or (invoice_data.get("Recipient", {}) or {}).get("Address", "")),
        }

    def add_bank(row):
        for k in EXPECTED_BANK_FIELDS:
            row[k] = fmt_text(bank_details.get(k, ""))
        return row

    if not line_items:
        row = add_bank(base_invoice_info())
        row.update({
            "Item Description":       "NA",
            "Service Name":           "NA",
            "Service Start Date":     "NA",
            "Service End Date":       "NA",
            "SKU":                    "NA",
            "Item Quantity":          0,
            "Item Unit Price":        0,
            "Item Amount":            0,
            "Discount Rate Per Item": 0,
            "Discount Amount Per Item": 0,
            "Tax Rate Per Item":      0,
            "Tax Amount Per Item":    0,
            "Item Line Total":        0,
            "IO Number/Cost Centre":  "NA",
        })
        rows.append(row)
        return rows

    for item in line_items:
        if not isinstance(item, dict):
            item = {}
        row = add_bank(base_invoice_info())
        row.update({
            "Item Description":       fmt_text(item.get("Description", item.get("Item Description", ""))),
            "Service Name":           fmt_text(item.get("Service Name", "")),
            "Service Start Date":     fmt_text(item.get("Service Start Date", "")),
            "Service End Date":       fmt_text(item.get("Service End Date", "")),
            "SKU":                    fmt_text(item.get("SKU", "")),
            "Item Quantity":          fmt_amount(item.get("Quantity", item.get("Item Quantity", 0))),
            "Item Unit Price":        fmt_amount(item.get("Unit Price", item.get("Item Unit Price", 0))),
            "Item Amount":            fmt_amount(item.get("Amount", item.get("Item Amount", 0))),
            "Discount Rate Per Item": fmt_amount(item.get("Discount Rate Per Item", 0)),
            "Discount Amount Per Item": fmt_amount(item.get("Discount Amount Per Item", 0)),
            "Tax Rate Per Item":      fmt_amount(item.get("Tax Rate Per Item", 0)),
            "Tax Amount Per Item":    fmt_amount(item.get("Tax Amount Per Item", 0)),
            "Item Line Total":        fmt_amount(item.get("Line Total", item.get("Item Line Total", item.get("Amount", 0)))),
            "IO Number/Cost Centre":  fmt_text(item.get("IO Number/Cost Centre", "")),
        })
        rows.append(row)
    return rows


# -----------------------------
# JSONL conversion
# -----------------------------
def build_page_filenames(file_name: str, num_pages: int) -> list[str]:
    """
    Single page → ["PI-2025-Dec-040.png"]
    Multi page  → ["PI-2025-Dec-051_page1.png", "PI-2025-Dec-051_page2.png"]
    Always a list, always .png, single page keeps original name with no suffix.
    """
    stem = Path(file_name).stem
    if num_pages == 1:
        return [f"{stem}.png"]
    return [f"{stem}_page{i+1}.png" for i in range(num_pages)]


def convert_to_jsonl_record(edited_data, file_name="", num_pages: int = 1):
    ed       = edited_data or {}
    currency = ed.get("Currency", "")
    items    = ed.get("Itemized Data", []) or []
    bank     = ed.get("Bank Details", {}) or {}

    header = {
        "invoice_no":    ed.get("Invoice Number", ""),
        "po_no":         ed.get("PO Number", ""),
        "invoice_date":  ed.get("Invoice Date", ""),
        "payment_terms": ed.get("Payment Terms", ""),
        "due_date":      ed.get("Due Date", ""),
        "sender_name":   ed.get("Sender Name", "") or (ed.get("Sender", {}) or {}).get("Name", ""),
        "sender_addr":   ed.get("Sender Address", "") or (ed.get("Sender", {}) or {}).get("Address", ""),
        "tax_id":        ed.get("Tax ID", ""),
        "rcpt_name":     ed.get("Recipient Name", "") or (ed.get("Recipient", {}) or {}).get("Name", ""),
        "rcpt_addr":     ed.get("Recipient Address", "") or (ed.get("Recipient", {}) or {}).get("Address", ""),
        "bank_iban":     bank.get("bank_iban", ""),
        "bank_name":     bank.get("bank_name", ""),
        "bank_acc_no":   bank.get("bank_account_number", ""),
        "bank_routing":  bank.get("bank_routing", ""),
        "bank_swift":    bank.get("bank_swift", ""),
        "bank_acc_name": bank.get("bank_acc_name", ""),
        "bank_branch":   bank.get("bank_branch", ""),
    }

    jsonl_items = []
    for it in items:
        if not isinstance(it, dict):
            continue
        jsonl_items.append({
            "service_name":             it.get("Service Name", ""),
            "service_start_date":       it.get("Service Start Date", ""),
            "service_end_date":         it.get("Service End Date", ""),
            "descriptions":             it.get("Description", ""),
            "SKU":                      it.get("SKU", ""),
            "quantity":                 it.get("Quantity", ""),
            "unit_price":               it.get("Unit Price", ""),
            "amount":                   it.get("Amount", ""),
            "discount_rate_per_item":   it.get("Discount Rate Per Item", ""),
            "discount_amount_per_item": it.get("Discount Amount Per Item", ""),
            "tax_rate_per_item":        it.get("Tax Rate Per Item", ""),
            "tax_amount_per_item":      it.get("Tax Amount Per Item", ""),
            "line_total":               it.get("Line Total", ""),
        })

    summary = {
        "subtotal":              ed.get("Subtotal", ""),
        "tax_rate":              ed.get("Tax Percentage", ""),
        "tax_amount":            ed.get("Total Tax", ""),
        "discount_rate":         ed.get("Discount Rate", ""),
        "total_discount_amount": ed.get("Total Discount Amount", ""),
        "total_amount":          ed.get("Total Amount", ""),
        "currency":              currency,
    }

    return {
        "file_name": build_page_filenames(file_name, num_pages),
        "gt_parse": {
            "header":  header,
            "items":   jsonl_items,
            "summary": summary,
        }
    }


def build_excel_bytes(df):
    buf = BytesIO()
    with pd.ExcelWriter(buf, engine="openpyxl") as writer:
        df.to_excel(writer, index=False, sheet_name="Invoices")
        ws = writer.sheets["Invoices"]
        for col_cells in ws.columns:
            max_len    = 0
            col_letter = col_cells[0].column_letter
            for cell in col_cells:
                try:
                    if cell.value:
                        max_len = max(max_len, len(str(cell.value)))
                except Exception:
                    pass
            ws.column_dimensions[col_letter].width = min(max_len + 3, 50)
    return buf.getvalue()


# =============================================================================
# Extraction helper — inference + code #1 postprocessing
# =============================================================================
def _run_extraction(pages: List[Image.Image]) -> tuple[str | None, dict, dict]:
    raw_json = run_inference_vllm(pages)
    if not raw_json:
        return None, {}, {}

    raw_parsed_data, parsed_data = parse_vllm_json(raw_json)
    if not parsed_data:
        st.warning("Failed to parse JSON response from model.")
        return raw_json, raw_parsed_data or {}, {}

    mapped = validate_and_calculate_taxes(parsed_data)
    safe_mapped = mapped if isinstance(mapped, dict) else {}
    display_ready = prepare_processed_invoice_data(raw_parsed_data or {}, safe_mapped)
    return raw_json, raw_parsed_data or {}, display_ready if isinstance(display_ready, dict) else safe_mapped


# -----------------------------
# Session scaffolding
# -----------------------------
if "batch_results" not in st.session_state:
    st.session_state.batch_results = {}
if "current_file_hash" not in st.session_state:
    st.session_state.current_file_hash = None
if "is_processing_batch" not in st.session_state:
    st.session_state.is_processing_batch = False
if "confirm_back" not in st.session_state:
    st.session_state.confirm_back = False

frame_left, frame_right = st.columns([1, 1], vertical_alignment="top")

# =============================================================================
# UPLOAD & BATCH PROCESSING
# =============================================================================
if not st.session_state.is_processing_batch and len(st.session_state.batch_results) == 0:
    with frame_left:
        st.header("📤 Upload Invoices")
        uploaded_files = st.file_uploader(
            "Upload invoice images or PDFs — all pages of a PDF are sent in one request",
            type=["png", "jpg", "jpeg", "pdf"],
            accept_multiple_files=True
        )

        if uploaded_files:
            st.session_state.is_processing_batch = True
            progress_bar = st.progress(0)
            status_text  = st.empty()

            for idx, uploaded_file in enumerate(uploaded_files):
                status_text.text(f"Processing {idx+1}/{len(uploaded_files)}: {uploaded_file.name}")
                uploaded_bytes = uploaded_file.read()
                file_hash      = hashlib.sha256(uploaded_bytes).hexdigest()

                if file_hash in st.session_state.batch_results:
                    progress_bar.progress((idx + 1) / len(uploaded_files))
                    continue

                pages  = []
                is_pdf = (
                    uploaded_file.name.lower().endswith(".pdf")
                    or (hasattr(uploaded_file, "type") and uploaded_file.type == "application/pdf")
                )
                if is_pdf:
                    if convert_from_bytes is None:
                        st.warning(f"PDF {uploaded_file.name} could not be rendered (pdf2image/poppler missing).")
                        progress_bar.progress((idx + 1) / len(uploaded_files))
                        continue
                    try:
                        pdf_pages = convert_from_bytes(uploaded_bytes, dpi=300)
                        pages     = [p.convert("RGB") for p in pdf_pages]
                        if not pages:
                            st.warning(f"PDF {uploaded_file.name} has no pages.")
                            progress_bar.progress((idx + 1) / len(uploaded_files))
                            continue
                        st.info(f"{uploaded_file.name}: {len(pages)} page(s) detected — all sent in one request.")
                    except Exception as exc:
                        st.warning(f"Could not render PDF {uploaded_file.name}: {exc}")
                        progress_bar.progress((idx + 1) / len(uploaded_files))
                        continue
                else:
                    try:
                        img = Image.open(BytesIO(uploaded_bytes)).convert("RGB")
                        pages.append(img)
                    except Exception:
                        st.warning(f"Failed to open {uploaded_file.name}.")
                        progress_bar.progress((idx + 1) / len(uploaded_files))
                        continue

                if not pages:
                    progress_bar.progress((idx + 1) / len(uploaded_files))
                    continue

                raw_json, raw_parsed_data, safe_mapped = _run_extraction(pages)
                if raw_json is None:
                    st.warning(f"No response from vLLM for {uploaded_file.name}")

                st.session_state.batch_results[file_hash] = {
                    "file_name":    uploaded_file.name,
                    "pages":        pages,
                    "current_page": 0,
                    "raw_pred":     raw_json,
                    "raw_parsed_data": raw_parsed_data,
                    "mapped_data":  safe_mapped,
                    "edited_data":  safe_mapped.copy(),
                    "user_saved_edits": False
                }

                progress_bar.progress((idx + 1) / len(uploaded_files))

            status_text.text("✅ All files processed!")
            st.session_state.is_processing_batch = False
            st.rerun()

    with frame_right:
        st.caption("Preview & editor will appear here after extraction.")

# =============================================================================
# REVIEW & EDIT
# =============================================================================
elif len(st.session_state.batch_results) > 0:

    with frame_left:
        all_rows = []
        for file_hash, result in st.session_state.batch_results.items():
            processed_result = display_data_for_result(result)
            rows = flatten_invoice_to_rows(processed_result)
            for r in rows:
                r["Source File"] = result.get("file_name", file_hash)
            all_rows.extend(rows)

        if all_rows:
            full_df = pd.DataFrame(all_rows)
            cols    = list(full_df.columns)
            if "Source File" in cols:
                cols    = ["Source File"] + [c for c in cols if c != "Source File"]
                full_df = full_df[cols]

            dl_cols = st.columns(3)
            with dl_cols[0]:
                st.download_button("📦 All CSV",
                                   full_df.to_csv(index=False).encode("utf-8"),
                                   file_name="all_extracted_invoices.csv",
                                   mime="text/csv",
                                   key="download_all_csv")
            with dl_cols[1]:
                st.download_button("📦 All Excel",
                                   build_excel_bytes(full_df),
                                   file_name="all_extracted_invoices.xlsx",
                                   mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
                                   key="download_all_xlsx")
            with dl_cols[2]:
                jsonl_lines = [
                    json.dumps(
                        convert_to_jsonl_record(
                            display_data_for_result(res),
                            res.get("file_name", ""),
                            len(res.get("pages", []))
                        ),
                        ensure_ascii=False
                    )
                    for res in st.session_state.batch_results.values()
                ]
                st.download_button("📦 All JSONL",
                                   "\n".join(jsonl_lines).encode("utf-8"),
                                   file_name="all_extracted_invoices.jsonl",
                                   mime="application/jsonl",
                                   key="download_all_jsonl")

    with frame_right:
        if st.button("⬅️ Back to Upload"):
            st.session_state[f"confirm_back"] = True

        if st.session_state.get("confirm_back"):
            st.warning("⚠️ All extracted data will be lost. Are you sure?")
            col_yes, col_no = st.columns(2)
            with col_yes:
                if st.button("✅ Yes, go back"):
                    st.session_state.batch_results.clear()
                    st.session_state.current_file_hash = None
                    st.session_state.is_processing_batch = False
                    st.session_state.confirm_back = False
                    st.rerun()
            with col_no:
                if st.button("❌ Cancel"):
                    st.session_state.confirm_back = False
                    st.rerun()

    with frame_left:
        file_options     = {f"{v['file_name']} ({k[:6]})": k for k, v in st.session_state.batch_results.items()}
        selected_display = st.selectbox("Select invoice to view/edit:",
                                        options=list(file_options.keys()),
                                        index=0,
                                        key="file_selector")
        selected_hash    = file_options[selected_display]

        if st.session_state.current_file_hash != selected_hash:
            if st.session_state.current_file_hash is not None:
                old_hash = st.session_state.current_file_hash
                for k in [k for k in st.session_state.keys() if k.endswith(f"_{old_hash}")]:
                    del st.session_state[k]
            st.session_state.current_file_hash = selected_hash

    current      = st.session_state.batch_results[selected_hash]
    pages        = current["pages"]
    num_pages    = len(pages)
    cur_page_idx = current.get("current_page", 0)
    form_data    = display_data_for_result(current)

    bank          = form_data.get("Bank Details", {}) if isinstance(form_data.get("Bank Details", {}), dict) else {}
    form_currency = form_data.get("Currency", "")

    # State defaults — editable numeric values stay numeric so saves round-trip cleanly.
    state_defaults = {
        f"Invoice Number_{selected_hash}":           form_data.get("Invoice Number", ""),
        f"PO Number_{selected_hash}":                form_data.get("PO Number", ""),
        f"Payment Terms_{selected_hash}":            form_data.get("Payment Terms", ""),
        f"Tax ID_{selected_hash}":                   form_data.get("Tax ID", ""),
        f"Currency_{selected_hash}":                 form_data.get("Currency", "USD") or "USD",
        f"Currency_Custom_{selected_hash}":          form_data.get("Currency", "") if form_data.get("Currency") not in ["USD","EUR","GBP","INR"] else "",
        f"Subtotal_{selected_hash}":                 editable_numeric_value("Subtotal", form_data.get("Subtotal", 0.0)),
        f"Tax Percentage_{selected_hash}":           editable_numeric_value("Tax Percentage", form_data.get("Tax Percentage", 0.0)),
        f"Total Tax_{selected_hash}":                editable_numeric_value("Total Tax", form_data.get("Total Tax", 0.0)),
        f"Discount Rate_{selected_hash}":            editable_numeric_value("Discount Rate", form_data.get("Discount Rate", 0.0)),
        f"Total Discount Amount_{selected_hash}":    editable_numeric_value("Total Discount Amount", form_data.get("Total Discount Amount", 0.0)),
        f"Total Amount_{selected_hash}":             editable_numeric_value("Total Amount", form_data.get("Total Amount", 0.0)),
        f"Sender Name_{selected_hash}":              form_data.get("Sender Name", ""),
        f"Sender Address_{selected_hash}":           form_data.get("Sender Address", ""),
        f"Recipient Name_{selected_hash}":           form_data.get("Recipient Name", ""),
        f"Recipient Address_{selected_hash}":        form_data.get("Recipient Address", ""),
        f"Bank_bank_name_{selected_hash}":           bank.get("bank_name", ""),
        f"Bank_bank_account_number_{selected_hash}": bank.get("bank_account_number", "") or bank.get("bank_acc_no", ""),
        f"Bank_bank_acc_name_{selected_hash}":       bank.get("bank_acc_name", "") or bank.get("bank_account_holder", ""),
        f"Bank_bank_iban_{selected_hash}":           bank.get("bank_iban", ""),
        f"Bank_bank_swift_{selected_hash}":          bank.get("bank_swift", ""),
        f"Bank_bank_routing_{selected_hash}":        bank.get("bank_routing", ""),
        f"Bank_bank_branch_{selected_hash}":         bank.get("bank_branch", ""),
    }
    for key, default in state_defaults.items():
        if key not in st.session_state:
            st.session_state[key] = default

    numeric_state_fields = {
        "Subtotal": f"Subtotal_{selected_hash}",
        "Tax Percentage": f"Tax Percentage_{selected_hash}",
        "Total Tax": f"Total Tax_{selected_hash}",
        "Discount Rate": f"Discount Rate_{selected_hash}",
        "Total Discount Amount": f"Total Discount Amount_{selected_hash}",
        "Total Amount": f"Total Amount_{selected_hash}",
    }
    for field, key in numeric_state_fields.items():
        st.session_state[key] = editable_numeric_value(field, st.session_state.get(key, 0.0))

    if f"Invoice Date_{selected_hash}" not in st.session_state:
        st.session_state[f"Invoice Date_{selected_hash}"] = parse_date_to_object(form_data.get("Invoice Date", ""), form_currency)
    if f"Due Date_{selected_hash}" not in st.session_state:
        st.session_state[f"Due Date_{selected_hash}"] = parse_date_to_object(form_data.get("Due Date", ""), form_currency)

    # Raw strings from model — used for ground truth output, never overwritten by date picker
    if f"Invoice Date_raw_{selected_hash}" not in st.session_state:
        st.session_state[f"Invoice Date_raw_{selected_hash}"] = form_data.get("Invoice Date", "")
    if f"Due Date_raw_{selected_hash}" not in st.session_state:
        st.session_state[f"Due Date_raw_{selected_hash}"] = form_data.get("Due Date", "")

    # =========================================================================
    # LEFT COLUMN — image preview
    # =========================================================================
    with frame_left:
        st.image(pages[cur_page_idx],
                 caption=f"{current['file_name']} — Page {cur_page_idx + 1} of {num_pages}",
                 use_container_width=True)

        if num_pages > 1:
            nav_cols = st.columns([1, 2, 1])
            with nav_cols[0]:
                if st.button("⬅️ Prev", disabled=(cur_page_idx == 0), key=f"prev_page_{selected_hash}"):
                    st.session_state.batch_results[selected_hash]["current_page"] = cur_page_idx - 1
                    st.rerun()
            with nav_cols[1]:
                st.markdown(f"<p style='text-align:center;'>Page {cur_page_idx + 1} / {num_pages}</p>",
                            unsafe_allow_html=True)
            with nav_cols[2]:
                if st.button("Next ➡️", disabled=(cur_page_idx == num_pages - 1), key=f"next_page_{selected_hash}"):
                    st.session_state.batch_results[selected_hash]["current_page"] = cur_page_idx + 1
                    st.rerun()

        st.write(f"**File Hash:** `{selected_hash[:8]}…`  |  **Pages:** {num_pages}")

        col_a, col_b = st.columns(2)

        with col_a:
            if st.button("🔁 Re-Run (All Pages)", key=f"rerun_all_{selected_hash}",
                         help=f"Re-send all {num_pages} page(s) in one request"):
                with st.spinner(f"Re-running inference on all {num_pages} page(s)…"):
                    try:
                        raw_json, raw_parsed_data, safe_mapped = _run_extraction(pages)
                        if raw_json is None:
                            st.error("No response from vLLM.")
                        else:
                            st.session_state.batch_results[selected_hash]["raw_pred"]    = raw_json
                            st.session_state.batch_results[selected_hash]["raw_parsed_data"] = raw_parsed_data
                            st.session_state.batch_results[selected_hash]["mapped_data"] = safe_mapped
                            st.session_state.batch_results[selected_hash]["edited_data"] = safe_mapped.copy()
                            st.session_state.batch_results[selected_hash]["user_saved_edits"] = False
                            for k in [k for k in st.session_state.keys() if k.endswith(f"_{selected_hash}")]:
                                del st.session_state[k]
                            st.success("✅ Re-run complete (all pages)")
                            st.rerun()
                    except Exception as e:
                        st.error(f"Re-run failed: {e}")

        with col_b:
            if st.button(f"🔍 Extract Page {cur_page_idx + 1} Only", key=f"extract_page_{selected_hash}",
                         help="Send only the currently displayed page"):
                with st.spinner(f"Running inference on page {cur_page_idx + 1} only…"):
                    try:
                        raw_json, raw_parsed_data, safe_mapped = _run_extraction([pages[cur_page_idx]])
                        if raw_json is None:
                            st.error("No response from vLLM.")
                        else:
                            st.session_state.batch_results[selected_hash]["raw_pred"]    = raw_json
                            st.session_state.batch_results[selected_hash]["raw_parsed_data"] = raw_parsed_data
                            st.session_state.batch_results[selected_hash]["mapped_data"] = safe_mapped
                            st.session_state.batch_results[selected_hash]["edited_data"] = safe_mapped.copy()
                            st.session_state.batch_results[selected_hash]["user_saved_edits"] = False
                            for k in [k for k in st.session_state.keys() if k.endswith(f"_{selected_hash}")]:
                                del st.session_state[k]
                            st.success(f"✅ Page {cur_page_idx + 1} extracted!")
                            st.rerun()
                    except Exception as e:
                        st.error(f"Extraction failed: {e}")

        with st.expander("🔍 Show raw model output"):
            raw_pred = current.get("raw_pred")
            if raw_pred is None:
                st.warning("No raw output available.")
            else:
                st.code(str(raw_pred), language="json")

    # =========================================================================
    # RIGHT COLUMN — editable form
    # =========================================================================
    with frame_right:
        st.subheader(f"Editable Invoice: {current['file_name']}")

        # Live widgets: keep the data editor outside a form so dynamic row adds/deletes are saved before downloads.
        tabs = st.tabs(["Invoice Details", "Sender/Recipient", "Bank Details", "Line Items"])

        with tabs[0]:
            st.text_input("Invoice Number", key=f"Invoice Number_{selected_hash}")
            st.text_input("PO Number",      key=f"PO Number_{selected_hash}")

            st.write("**Invoice Date:**")
            raw_invoice_date = form_data.get("Invoice Date", "")
            if raw_invoice_date:
                st.info(f"📅 Model extracted: {raw_invoice_date}")
            st.date_input("Select date:", key=f"Invoice Date_{selected_hash}",
                          format="DD/MM/YYYY", label_visibility="collapsed")

            st.text_input("Payment Terms", key=f"Payment Terms_{selected_hash}")

            st.write("**Due Date:**")
            raw_due_date = form_data.get("Due Date", "")
            if raw_due_date:
                st.info(f"📅 Model extracted: {raw_due_date}")
            st.date_input("Select date:", key=f"Due Date_{selected_hash}",
                          format="DD/MM/YYYY", label_visibility="collapsed")

            st.text_input("Tax ID / VAT Number", key=f"Tax ID_{selected_hash}")

            curr_options = ["USD", "EUR", "GBP", "INR", "Other"]
            if st.session_state[f"Currency_{selected_hash}"] not in curr_options:
                st.session_state[f"Currency_{selected_hash}"] = "Other"
            st.selectbox("Currency", options=curr_options, key=f"Currency_{selected_hash}")
            if st.session_state.get(f"Currency_{selected_hash}") == "Other":
                st.text_input("Specify Currency", key=f"Currency_Custom_{selected_hash}")

            st.number_input("Subtotal",              key=f"Subtotal_{selected_hash}",              format="%.3f", step=0.001)
            st.number_input("Tax %",                 key=f"Tax Percentage_{selected_hash}",        format="%.3f", step=0.001)
            st.number_input("Total Tax",             key=f"Total Tax_{selected_hash}",             format="%.3f", step=0.001)
            st.number_input("Discount Rate %",       key=f"Discount Rate_{selected_hash}",         format="%.3f", step=0.001)
            st.number_input("Total Discount Amount", key=f"Total Discount Amount_{selected_hash}", format="%.3f", step=0.001)
            st.number_input("Total Amount",          key=f"Total Amount_{selected_hash}",          format="%.3f", step=0.001)

        with tabs[1]:
            st.text_input("Sender Name",       key=f"Sender Name_{selected_hash}")
            st.text_area ("Sender Address",    key=f"Sender Address_{selected_hash}",    height=80)
            st.text_input("Recipient Name",    key=f"Recipient Name_{selected_hash}")
            st.text_area ("Recipient Address", key=f"Recipient Address_{selected_hash}", height=80)

        with tabs[2]:
            st.text_input("Bank Name",       key=f"Bank_bank_name_{selected_hash}")
            st.text_input("Account Number",  key=f"Bank_bank_account_number_{selected_hash}")
            st.text_input("Account Name",    key=f"Bank_bank_acc_name_{selected_hash}")
            st.text_input("IBAN",            key=f"Bank_bank_iban_{selected_hash}")
            st.text_input("SWIFT",           key=f"Bank_bank_swift_{selected_hash}")
            st.text_input("Routing",         key=f"Bank_bank_routing_{selected_hash}")
            st.text_input("Branch",          key=f"Bank_bank_branch_{selected_hash}")

        with tabs[3]:
            items_state_key = f"items_df_{selected_hash}"

            if items_state_key not in st.session_state:
                item_rows  = form_data.get("Itemized Data", []) or []
                normalized = []
                for it in item_rows:
                    if not isinstance(it, dict):
                        it = {}
                    normalized.append({
                        "Description":              it.get("Description", it.get("Item Description", "")),
                        "Service Name":             it.get("Service Name", ""),
                        "Service Start Date":       it.get("Service Start Date", ""),
                        "Service End Date":         it.get("Service End Date", ""),
                        "SKU":                      it.get("SKU", ""),
                        "Quantity":                 editable_numeric_value("Quantity", it.get("Quantity", it.get("Item Quantity", 0.0))),
                        "Unit Price":               editable_numeric_value("Unit Price", it.get("Unit Price", it.get("Item Unit Price", 0.0))),
                        "Amount":                   editable_numeric_value("Amount", it.get("Amount", it.get("Item Amount", 0.0))),
                        "Discount Rate Per Item":   editable_numeric_value("Discount Rate Per Item", it.get("Discount Rate Per Item", 0.0)),
                        "Discount Amount Per Item": editable_numeric_value("Discount Amount Per Item", it.get("Discount Amount Per Item", 0.0)),
                        "Tax Rate Per Item":        editable_numeric_value("Tax Rate Per Item", it.get("Tax Rate Per Item", 0.0)),
                        "Tax Amount Per Item":      editable_numeric_value("Tax Amount Per Item", it.get("Tax Amount Per Item", 0.0)),
                        "IO Number/Cost Centre":    it.get("IO Number/Cost Centre", ""),
                        "Line Total":               editable_numeric_value("Line Total", it.get("Line Total", it.get("Item Line Total", 0.0))),
                    })
                st.session_state[items_state_key] = coerce_numeric_items_df(
                    pd.DataFrame(normalized) if normalized
                    else pd.DataFrame(columns=[
                        "Description", "Service Name", "Service Start Date", "Service End Date", "SKU",
                        "Quantity", "Unit Price", "Amount",
                        "Discount Rate Per Item", "Discount Amount Per Item",
                        "Tax Rate Per Item", "Tax Amount Per Item",
                        "IO Number/Cost Centre", "Line Total"
                    ])
                )

            items_df = coerce_numeric_items_df(st.session_state[items_state_key])

            column_config = {
                "Description":              st.column_config.TextColumn("Description",    width="large"),
                "Service Name":             st.column_config.TextColumn("Service Name",   width="medium"),
                "Service Start Date":       st.column_config.TextColumn("Svc Start",      width="small"),
                "Service End Date":         st.column_config.TextColumn("Svc End",        width="small"),
                "SKU":                      st.column_config.TextColumn("SKU",            width="small"),
                "Quantity":                 st.column_config.NumberColumn("Qty",          width="small", format="%.3f", step=0.001),
                "Unit Price":               st.column_config.NumberColumn("Unit Price",   width="small", format="%.3f", step=0.001),
                "Amount":                   st.column_config.NumberColumn("Amount",       width="small", format="%.3f", step=0.001),
                "Discount Rate Per Item":   st.column_config.NumberColumn("Disc %",       width="small", format="%.3f", step=0.001),
                "Discount Amount Per Item": st.column_config.NumberColumn("Disc Amt",     width="small", format="%.3f", step=0.001),
                "Tax Rate Per Item":        st.column_config.NumberColumn("Tax %",        width="small", format="%.3f", step=0.001),
                "Tax Amount Per Item":      st.column_config.NumberColumn("Tax Amt",      width="small", format="%.3f", step=0.001),
                "IO Number/Cost Centre":    st.column_config.TextColumn("IO/Cost Centre", width="medium"),
                "Line Total":               st.column_config.NumberColumn("Line Total",   width="small", format="%.3f", step=0.001),
            }

            editor_key = f"items_editor_{selected_hash}"
            edited_df = st.data_editor(
                items_df,
                num_rows="dynamic",
                key=editor_key,
                use_container_width=True,
                height=DATA_EDITOR_HEIGHT,
                column_config=column_config,
            )

        saved = st.button("💾 Save All Edits", key=f"save_all_edits_{selected_hash}")

        currency = st.session_state.get(f"Currency_{selected_hash}", "USD")
        if currency == "Other":
            currency = st.session_state.get(f"Currency_Custom_{selected_hash}", "")

        items_state_key  = f"items_df_{selected_hash}"
        current_items_df = coerce_numeric_items_df(st.session_state.get(items_state_key, pd.DataFrame()))
        st.session_state[items_state_key] = current_items_df
        line_items_list  = current_items_df.to_dict("records")

        def _build_updated_dict():
            current_df = apply_data_editor_state(
                st.session_state.get(items_state_key, pd.DataFrame()),
                st.session_state.get(f"items_editor_{selected_hash}"),
            )
            current_line_items = current_df.to_dict("records")
            updated = {
                "Invoice Number":        st.session_state.get(f"Invoice Number_{selected_hash}", ""),
                "PO Number":             st.session_state.get(f"PO Number_{selected_hash}", ""),
                # Raw model string — not the date picker's reformatted value
                "Invoice Date":          st.session_state.get(f"Invoice Date_raw_{selected_hash}", ""),
                "Payment Terms":         st.session_state.get(f"Payment Terms_{selected_hash}", ""),
                "Due Date":              st.session_state.get(f"Due Date_raw_{selected_hash}", ""),
                "Tax ID":                st.session_state.get(f"Tax ID_{selected_hash}", ""),
                "Currency":              currency,
                "Subtotal":              editable_numeric_value("Subtotal", st.session_state.get(f"Subtotal_{selected_hash}", 0.0)),
                "Tax Percentage":        editable_numeric_value("Tax Percentage", st.session_state.get(f"Tax Percentage_{selected_hash}", 0.0)),
                "Total Tax":             editable_numeric_value("Total Tax", st.session_state.get(f"Total Tax_{selected_hash}", 0.0)),
                "Discount Rate":         editable_numeric_value("Discount Rate", st.session_state.get(f"Discount Rate_{selected_hash}", 0.0)),
                "Total Discount Amount": editable_numeric_value("Total Discount Amount", st.session_state.get(f"Total Discount Amount_{selected_hash}", 0.0)),
                "Total Amount":          editable_numeric_value("Total Amount", st.session_state.get(f"Total Amount_{selected_hash}", 0.0)),
                "Sender Name":           st.session_state.get(f"Sender Name_{selected_hash}", ""),
                "Sender Address":        st.session_state.get(f"Sender Address_{selected_hash}", ""),
                "Recipient Name":        st.session_state.get(f"Recipient Name_{selected_hash}", ""),
                "Recipient Address":     st.session_state.get(f"Recipient Address_{selected_hash}", ""),
                "Bank Details": {
                    "bank_name":           st.session_state.get(f"Bank_bank_name_{selected_hash}", ""),
                    "bank_account_number": st.session_state.get(f"Bank_bank_account_number_{selected_hash}", ""),
                    "bank_acc_name":       st.session_state.get(f"Bank_bank_acc_name_{selected_hash}", ""),
                    "bank_iban":           st.session_state.get(f"Bank_bank_iban_{selected_hash}", ""),
                    "bank_swift":          st.session_state.get(f"Bank_bank_swift_{selected_hash}", ""),
                    "bank_routing":        st.session_state.get(f"Bank_bank_routing_{selected_hash}", ""),
                    "bank_branch":         st.session_state.get(f"Bank_bank_branch_{selected_hash}", ""),
                },
                "Itemized Data": current_line_items,
                "Sender":    {"Name": st.session_state.get(f"Sender Name_{selected_hash}", ""),
                              "Address": st.session_state.get(f"Sender Address_{selected_hash}", "")},
                "Recipient": {"Name": st.session_state.get(f"Recipient Name_{selected_hash}", ""),
                              "Address": st.session_state.get(f"Recipient Address_{selected_hash}", "")},
            }
            return limit_invoice_decimals(updated)

        if saved:
            updated = _build_updated_dict()
            # Persist user-edited date picker values back as the date strings
            invoice_date = st.session_state.get(f"Invoice Date_{selected_hash}")
            due_date     = st.session_state.get(f"Due Date_{selected_hash}")
            if invoice_date is not None:
                try:
                    updated["Invoice Date"] = invoice_date.strftime("%d-%b-%Y")
                    st.session_state[f"Invoice Date_raw_{selected_hash}"] = invoice_date.strftime("%d-%b-%Y")
                except (AttributeError, ValueError):
                    pass
            if due_date is not None:
                try:
                    updated["Due Date"] = due_date.strftime("%d-%b-%Y")
                    st.session_state[f"Due Date_raw_{selected_hash}"] = due_date.strftime("%d-%b-%Y")
                except (AttributeError, ValueError):
                    pass
            updated = limit_invoice_decimals(updated)
            st.session_state.batch_results[selected_hash]["edited_data"] = updated
            st.session_state.batch_results[selected_hash]["user_saved_edits"] = True
            for k in [k for k in list(st.session_state.keys()) if k.endswith(f"_{selected_hash}")]:
                del st.session_state[k]
            st.success("✅ Saved")
            st.rerun()

        download_data = limit_invoice_decimals(_build_updated_dict())
        rows          = flatten_invoice_to_rows(download_data)
        file_df       = pd.DataFrame(rows)
        file_stem     = Path(current["file_name"]).stem

        per_dl_cols = st.columns(3)
        with per_dl_cols[0]:
            st.download_button("📥 CSV",
                               file_df.to_csv(index=False).encode("utf-8"),
                               file_name=f"{file_stem}_full.csv",
                               mime="text/csv",
                               key=f"dl_csv_{selected_hash}")
        with per_dl_cols[1]:
            st.download_button("📥 Excel",
                               build_excel_bytes(file_df),
                               file_name=f"{file_stem}_full.xlsx",
                               mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
                               key=f"dl_xlsx_{selected_hash}")
        with per_dl_cols[2]:
            record    = convert_to_jsonl_record(download_data, current["file_name"], num_pages)
            jsonl_one = json.dumps(record, ensure_ascii=False).encode("utf-8")
            st.download_button("📥 JSONL",
                               jsonl_one,
                               file_name=f"{file_stem}.jsonl",
                               mime="application/jsonl",
                               key=f"dl_jsonl_{selected_hash}")

# =============================================================================
# Processing placeholder
# =============================================================================
elif st.session_state.is_processing_batch:
    with frame_left:
        st.info("⏳ Processing batch… Please wait.")
        st.progress(0)
    with frame_right:
        st.caption("Preview & editor will appear here after extraction.")

else:
    with frame_left:
        st.caption("Ready when you are.")
    with frame_right:
        st.caption("Preview & editor will appear here after extraction.")