Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,6 +1,12 @@
 import os
 # --- Fix: ensure HOME is writable before Streamlit initializes ---
 from pathlib import Path
+def safe_number_input(label, value, key):
+    try:
+        v = float(value)
+    except Exception:
+        v = 0.0
+    return st.number_input(label, value=v, key=key)
 
 _home = os.environ.get("HOME", "")
 if _home in ("", "/", None):
@@ -32,45 +38,10 @@ import pandas as pd
 from PIL import Image
 from huggingface_hub import login
 
-
-# ---------------------------
-# Helper: Update session state on widget change
-# ---------------------------
-def _update_flat_field(field_name: str, widget_key: str):
-    """Update a top-level field in edited_data."""
-    current_hash = st.session_state.get("current_file_hash")
-    if not current_hash or current_hash not in st.session_state.batch_results:
-        return
-    new_value = st.session_state[widget_key]
-    st.session_state.batch_results[current_hash]["edited_data"][field_name] = new_value
-
-
-def _update_nested_field(field_path: str, widget_key: str):
-    """Update a nested field like 'Sender.Name' or 'Bank Details.bank_name'."""
-    current_hash = st.session_state.get("current_file_hash")
-    if not current_hash or current_hash not in st.session_state.batch_results:
-        return
-    new_value = st.session_state[widget_key]
-    data_ref = st.session_state.batch_results[current_hash]["edited_data"]
-    
-    keys = field_path.split(".")
-    d = data_ref
-    for k in keys[:-1]:
-        if k not in d or not isinstance(d[k], dict):
-            d[k] = {}
-        d = d[k]
-    d[keys[-1]] = new_value
-
-
-def _update_line_items(widget_key: str):
-    """Auto-update line items when data_editor changes."""
-    current_hash = st.session_state.get("current_file_hash")
-    if not current_hash or current_hash not in st.session_state.batch_results:
-        return
-    df = st.session_state[widget_key]
-    records = df.to_dict('records')
-    st.session_state.batch_results[current_hash]["edited_data"]["Itemized Data"] = records
-
+# 👇 ADD CALLBACK FUNCTION AT TOP LEVEL
+def trigger_rerun():
+    # Forces immediate rerun after any edit — ensures edited_df is fresh
+    pass
 
 # ---------------------------
 # UI: main
@@ -236,44 +207,25 @@ def run_inference_on_image(image: Image.Image, processor, model, device, decoder
 # ---------------------------
 def map_prediction_to_ui(pred):
     import json, re
-    from collections import defaultdict
+    from datetime import datetime
 
-    # --- parse raw string payloads that embed JSON ---
     def safe_json_load(s):
         if s is None:
             return None
         if isinstance(s, (dict, list)):
             return s
         if isinstance(s, str):
-            s = s.strip()
-            if s == "":
-                return None
             try:
                 return json.loads(s)
             except Exception:
-                # try to extract balanced-brace substrings (simple approach)
-                subs = []
-                stack = []
-                start = None
-                for i, ch in enumerate(s):
-                    if ch == "{":
-                        if not stack:
-                            start = i
-                        stack.append("{")
-                    elif ch == "}":
-                        if stack:
-                            stack.pop()
-                            if not stack and start is not None:
-                                subs.append(s[start:i+1])
-                                start = None
-                for sub in subs:
-                    try:
-                        return json.loads(sub)
-                    except Exception:
-                        continue
+                try:
+                    t = s.strip()
+                    t = t.replace("\\'", "'").replace('\"{', '{').replace('}\"', '}')
+                    return json.loads(t)
+                except Exception:
+                    return None
         return None
 
-    # --- normalize numeric strings like "1,800.00" -> float ---
     def clean_number(x):
         if x is None:
             return 0.0
@@ -282,7 +234,6 @@ def map_prediction_to_ui(pred):
         s = str(x).strip()
         if s == "":
             return 0.0
-        # remove commas and non-number chars except dot and minus
         s = re.sub(r"[,\s]", "", s)
         s = re.sub(r"[^\d\.\-]", "", s)
         if s in ("", ".", "-", "-."):
@@ -292,69 +243,28 @@ def map_prediction_to_ui(pred):
         except Exception:
             return 0.0
 
-    # --- collect all keys -> list of values from pred, recursively ---
-    def collect_keys(obj, out):
-        if isinstance(obj, dict):
-            for k, v in obj.items():
-                lk = str(k).strip().lower()
-                out[lk].append(v)
-                collect_keys(v, out)
-        elif isinstance(obj, list):
-            for it in obj:
-                collect_keys(it, out)
-        else:
-            # primitive: handled via parent key
-            pass
-
-    # --- find list-of-dicts candidates for items (recursively) ---
-    def collect_lists_of_dicts(obj, out_lists):
-        if isinstance(obj, dict):
-            for v in obj.values():
-                if isinstance(v, list) and v and isinstance(v[0], dict):
-                    out_lists.append(v)
-                else:
-                    collect_lists_of_dicts(v, out_lists)
-        elif isinstance(obj, list):
-            for it in obj:
-                if isinstance(it, list) and it and isinstance(it[0], dict):
-                    out_lists.append(it)
-                else:
-                    collect_lists_of_dicts(it, out_lists)
-
-    # --- map item dict -> UI item row using the keys you specified in example ---
-    def map_item_dict(it):
-        if not isinstance(it, dict):
-            return None
-        # lowered keys mapping
-        lower = {str(k).strip().lower(): v for k, v in it.items()}
-        desc = (lower.get("descriptions") or lower.get("description") or lower.get("desc") or lower.get("item") or "")
-        qty = lower.get("quantity") or lower.get("qty") or lower.get("count") or ""
-        unit_price = lower.get("unit_price") or lower.get("price") or ""
-        amount = lower.get("amount") or lower.get("line_total") or lower.get("line total") or lower.get("total") or ""
-        tax = lower.get("tax") or lower.get("tax_amount") or ""
-        line_total = lower.get("line_total") or lower.get("line_total".lower()) or lower.get("line total") or amount
-
-        return {
-            "Description": str(desc).strip(),
-            "Quantity": float(clean_number(qty)),
-            "Unit Price": float(clean_number(unit_price)),
-            "Amount": float(clean_number(amount)),
-            "Tax": float(clean_number(tax)),
-            "Line Total": float(clean_number(line_total))
-        }
-
-    # ----------------- Start mapping -----------------
-    # Try parse if pred is a JSON-like string
-    parsed = safe_json_load(pred) if isinstance(pred, str) else pred
-    if parsed is None and isinstance(pred, str):
-        # not parseable -> fallback to empty UI
-        parsed = None
-
-    if parsed is None and not isinstance(pred, dict):
-        # nothing we can map
-        parsed = pred  # will still allow collect_keys if it's dict; else produce empty ui
+    def parse_date(s):
+        if not s:
+            return ""
+        s = str(s).strip()
+        for fmt in ("%Y-%m-%d", "%d-%m-%Y", "%d/%m/%Y", "%m/%d/%Y", "%d.%m.%Y"):
+            try:
+                return datetime.strptime(s, fmt).strftime("%Y-%m-%d")
+            except Exception:
+                pass
+        m = re.match(r"^(\d{1,2})/(\d{1,2})/(\d{4})$", s)
+        if m:
+            a, b, y = int(m.group(1)), int(m.group(2)), int(m.group(3))
+            if a > 12:
+                d, mo = a, b
+            else:
+                mo, d = a, b
+            try:
+                return datetime(year=y, month=mo, day=d).strftime("%Y-%m-%d")
+            except Exception:
+                return s
+        return s
 
-    # create empty UI template
     ui = {
         "Invoice Number": "",
         "Invoice Date": "",
@@ -374,122 +284,138 @@ def map_prediction_to_ui(pred):
         "Itemized Data": []
     }
 
-    # If parsed is a dict, collect all keys and list-of-dict candidates
-    key_map = defaultdict(list)  # lowercase-key -> list of values
-    list_candidates = []         # list of list-of-dicts found
-    if isinstance(parsed, dict):
-        collect_keys(parsed, key_map)
-        collect_lists_of_dicts(parsed, list_candidates)
-    elif isinstance(pred, dict):
-        # if parsing failed but original pred is dict, use that
-        collect_keys(pred, key_map)
-        collect_lists_of_dicts(pred, list_candidates)
-
-    # Helper to pick first non-empty value from candidate keys
-    def pick_first(*candidate_keys):
-        for k in candidate_keys:
-            lk = k.strip().lower()
-            if lk in key_map:
-                # pick first non-empty
-                for v in key_map[lk]:
-                    if v is None:
-                        continue
-                    # return primitive or string immediately; if dict/list, return as-is
-                    if isinstance(v, (dict, list)):
-                        return v
-                    s = str(v).strip()
-                    if s != "":
-                        return s
-        return None
+    if pred is None:
+        return ui
 
-    # Map simple scalar fields using the exact keys you provided (plus common close variants)
-    ui["Invoice Number"] = pick_first("invoice_no", "invoice_number", "invoiceid", "invoice id") or ""
-    ui["Invoice Date"] = pick_first("invoice_date", "date", "invoice date") or ""
-    ui["Due Date"] = pick_first("due_date", "due_date", "due") or ""
-    ui["Sender Name"] = pick_first("sender_name", "sender") or ""
-    ui["Sender Address"] = pick_first("sender_addr", "sender_address", "sender addr") or ""
-    ui["Recipient Name"] = pick_first("rcpt_name", "recipient_name", "recipient", "rcpt") or ""
-    ui["Recipient Address"] = pick_first("rcpt_addr", "recipient_address", "recipient addr") or ""
+    if isinstance(pred, str):
+        parsed = safe_json_load(pred)
+        if parsed is not None:
+            pred = parsed
 
-    # bank details: gather keys that start with 'bank_' or exact matches
-    bank = {}
-    for bk in ("bank_name", "bank_acc_no", "bank_account_number", "bank_acc_name", "bank_iban", "bank_swift", "bank_routing", "bank_branch", "iban"):
-        val = pick_first(bk, bk.replace("bank_", ""))  # allow both 'iban' and 'bank_iban'
-        if val:
-            # normalize key name to bank_* form
-            if bk == "iban":
-                bank["bank_iban"] = str(val)
-            else:
-                bank[bk] = str(val)
-    ui["Bank Details"] = bank
+    gt = None
+    if isinstance(pred, dict):
+        if "gt_parse" in pred:
+            gp = pred["gt_parse"]
+            gp_parsed = safe_json_load(gp)
+            gt = gp_parsed if gp_parsed is not None else (gp if isinstance(gp, dict) else {})
+        else:
+            gt = pred
+    else:
+        return ui
+
+    header = gt.get("header") or {}
+    items = gt.get("items") or []
+    summary = gt.get("summary") or {}
+
+    ui["Invoice Number"] = header.get("invoice_no") or header.get("invoice_number") or ui["Invoice Number"]
+    ui["Invoice Date"] = str(header.get("invoice_date") or header.get("inv_date") or "")
+    ui["Due Date"] = str(header.get("due_date") or header.get("due") or "")
+
+    ui["Sender Name"] = header.get("sender_name") or header.get("seller_name") or header.get("from_name") or ui["Sender Name"]
+    ui["Sender Address"] = header.get("sender_addr") or header.get("sender_address") or header.get("seller_addr") or ui["Sender Address"]
+    ui["Recipient Name"] = header.get("rcpt_name") or header.get("recipient_name") or header.get("to_name") or ui["Recipient Name"]
+    ui["Recipient Address"] = header.get("rcpt_addr") or header.get("rcpt_address") or header.get("recipient_address") or ui["Recipient Address"]
 
-    # summary / totals
-    ui["Subtotal"] = clean_number(pick_first("subtotal", "sub_total", "sub total") or 0.0)
-    ui["Tax Percentage"] = clean_number(pick_first("tax_rate", "tax_percentage", "tax pct", "tax percentage") or 0.0)
-    ui["Total Tax"] = clean_number(pick_first("tax_amount", "tax", "total_tax") or 0.0)
-    ui["Total Amount"] = clean_number(pick_first("total_amount", "grand_total", "total", "amount") or 0.0)
-    ui["Currency"] = (pick_first("currency") or "").strip()
-
-    # Item extraction:
-    items_rows = []
-
-    # 1) If we found an explicit list-of-dicts candidate, use the first one that looks like invoice items
-    def list_looks_like_items(lst):
-        if not isinstance(lst, list) or not lst:
-            return False
-        if not isinstance(lst[0], dict):
-            return False
-        # check if any expected item key present in first element
-        expected = {"descriptions", "description", "desc", "item", "quantity", "qty", "amount", "unit_price", "line_total", "line_total".lower(), "line_total"}
-        keys0 = {str(k).strip().lower() for k in lst[0].keys()}
-        return bool(expected.intersection(keys0))
-
-    for cand in list_candidates:
-        if list_looks_like_items(cand):
-            for it in cand:
-                row = map_item_dict(it)
-                if row is not None:
-                    items_rows.append(row)
-            # prefer first plausible list
-            if items_rows:
-                break
-
-    # 2) If no list-of-dicts found, try to find a single dict anywhere that looks like an item (e.g., 'items': {...} as dict)
-    if not items_rows:
-        # search key_map values for dicts that have item-like keys
-        for k, vals in key_map.items():
-            for v in vals:
-                if isinstance(v, dict):
-                    # does this dict have an item-like key?
-                    lower_keys = {str(x).strip().lower() for x in v.keys()}
-                    if lower_keys.intersection({"descriptions", "description", "desc", "amount", "line_total", "quantity", "qty", "unit_price"}):
-                        row = map_item_dict(v)
-                        if row is not None:
-                            items_rows.append(row)
-                            # we don't break because there might be multiple item-like dicts at different keys,
-                            # but continue scanning to collect all.
-    # 3) Last resort: if key_map contains 'descriptions' or 'amount' as scalar but no dict, build a single-item row
-    if not items_rows:
-        desc = pick_first("descriptions", "description")
-        amt = pick_first("amount", "line_total")
-        qty = pick_first("quantity", "qty")
-        unit_price = pick_first("unit_price", "price")
-        if desc or amt or qty or unit_price:
-            items_rows.append({
-                "Description": str(desc or ""),
-                "Quantity": float(clean_number(qty)),
-                "Unit Price": float(clean_number(unit_price)),
-                "Amount": float(clean_number(amt)),
-                "Tax": float(clean_number(pick_first("tax", "tax_amount") or 0.0)),
-                "Line Total": float(clean_number(amt or 0.0))
-            })
-
-    ui["Itemized Data"] = items_rows
-
-    # Also set Sender/Recipient convenience fields
     ui["Sender"] = {"Name": ui["Sender Name"], "Address": ui["Sender Address"]}
     ui["Recipient"] = {"Name": ui["Recipient Name"], "Address": ui["Recipient Address"]}
 
+    bank = {}
+    if header.get("bank_name"):
+        bank["bank_name"] = str(header.get("bank_name")).strip()
+    if header.get("bank_acc_no"):
+        bank["bank_account_number"] = str(header.get("bank_acc_no")).strip()
+    if header.get("bank_account_number"):
+        bank["bank_account_number"] = bank.get("bank_account_number") or str(header.get("bank_account_number")).strip()
+    if header.get("bank_iban"):
+        bank["bank_iban"] = str(header.get("bank_iban")).strip()
+    if header.get("bank_routing"):
+        bank["bank_routing"] = str(header.get("bank_routing")).strip()
+    if header.get("bank_swift"):
+        bank["bank_swift"] = str(header.get("bank_swift")).strip()
+    if header.get("bank_branch"):
+        bank["bank_branch"] = str(header.get("bank_branch")).strip()
+    if header.get("bank_acc_name"):
+        bank["bank_acc_name"] = str(header.get("bank_acc_name")).strip()
+    hb = header.get("bank")
+    if isinstance(hb, dict):
+        for k, v in hb.items():
+            if not v:
+                continue
+            lk = k.lower()
+            if "iban" in lk:
+                bank["bank_iban"] = bank.get("bank_iban") or str(v).strip()
+            elif "swift" in lk:
+                bank["bank_swift"] = bank.get("bank_swift") or str(v).strip()
+            elif "acc" in lk or "account" in lk:
+                bank["bank_account_number"] = bank.get("bank_account_number") or str(v).strip()
+            elif "name" in lk and "bank" in lk:
+                bank["bank_name"] = bank.get("bank_name") or str(v).strip()
+            elif "branch" in lk:
+                bank["bank_branch"] = bank.get("bank_branch") or str(v).strip()
+            elif "acc_name" in lk or "account_name" in lk:
+                bank["bank_acc_name"] = bank.get("bank_acc_name") or str(v).strip()
+
+    ui["Bank Details"] = bank
+
+    ui["Subtotal"] = clean_number(summary.get("subtotal") or summary.get("sub_total") or summary.get("subTotal"))
+    ui["Tax Percentage"] = clean_number(summary.get("tax_rate") or summary.get("taxRate") or summary.get("tax_percentage"))
+    ui["Total Tax"] = clean_number(summary.get("tax_amount") or summary.get("tax") or summary.get("taxAmount"))
+    ui["Total Amount"] = clean_number(summary.get("total_amount") or summary.get("grand_total") or summary.get("total") or summary.get("amount_total"))
+    ui["Currency"] = summary.get("currency") or header.get("currency") or ui["Currency"] or ""
+
+    normalized_items = []
+
+    if isinstance(items, str):
+        parsed_items = safe_json_load(items)
+        if parsed_items is not None:
+            items = parsed_items
+
+    if isinstance(items, dict):
+        if any(isinstance(v, list) for v in items.values()):
+            list_cols = {k: v for k, v in items.items() if isinstance(v, list)}
+            max_len = max((len(v) for v in list_cols.values()), default=0)
+            for i in range(max_len):
+                row = {}
+                for k, v in items.items():
+                    if isinstance(v, list):
+                        row[k] = v[i] if i < len(v) else ""
+                    else:
+                        row[k] = v
+                normalized_items.append(row)
+        else:
+            normalized_items.append(items)
+    elif isinstance(items, list):
+        normalized_items = items
+    else:
+        normalized_items = []
+
+    item_rows = []
+    for it in normalized_items:
+        if not isinstance(it, dict):
+            item_rows.append({"Description": str(it), "Quantity": 1, "Unit Price": 0.0, "Amount": 0.0, "Tax": 0.0, "Line Total": 0.0})
+            continue
+        desc = it.get("descriptions") or it.get("description") or it.get("desc") or it.get("item") or it.get("name") or ""
+        qty = it.get("quantity") or it.get("qty") or it.get("Quantity") or ""
+        unit = it.get("unit_price") or it.get("unitPrice") or it.get("price") or ""
+        amt = it.get("amount") or it.get("Line_total") or it.get("line_total") or it.get("total") or ""
+
+        # Extract item-level tax if available under common keys
+        tax_val = it.get("tax") or it.get("tax_amount") or it.get("line_tax") or it.get("item_tax") or it.get("taxAmount") or ""
+
+        # Extract explicit line total if present; otherwise fall back to amount
+        line_total_val = it.get("Line_total") or it.get("line_total") or it.get("lineTotal") or amt
+
+        item_rows.append({
+            "Description": str(desc).strip(),
+            "Quantity": float(clean_number(qty)),
+            "Unit Price": float(clean_number(unit)),
+            "Amount": float(clean_number(amt)),
+            "Tax": float(clean_number(tax_val)),
+            "Line Total": float(clean_number(line_total_val))
+        })
+
+    ui["Itemized Data"] = item_rows
+
     return ui
 
 # ---------------------------
@@ -499,53 +425,35 @@ def flatten_invoice_to_rows(invoice_data) -> list:
     """
     Converts nested invoice data into a flat list of rows (one per line item),
     with invoice-level and sender/recipient/bank fields repeated in each row.
-    Ensures all expected bank fields are present, even if empty.
     """
-    # Define all expected bank fields (must match UI keys)
-    EXPECTED_BANK_FIELDS = [
-        "bank_name",
-        "bank_account_number",
-        "bank_acc_name",
-        "bank_iban",
-        "bank_swift",
-        "bank_routing",
-        "bank_branch"
-    ]
-
     rows = []
     line_items = invoice_data.get("Itemized Data", [])
 
-    # Prepare base invoice info
-    base_row = {
-        "Invoice Number": invoice_data.get("Invoice Number", ""),
-        "Invoice Date": invoice_data.get("Invoice Date", ""),
-        "Due Date": invoice_data.get("Due Date", ""),
-        "Currency": invoice_data.get("Currency", ""),
-        "Subtotal": invoice_data.get("Subtotal", 0.0),
-        "Tax Percentage": invoice_data.get("Tax Percentage", 0.0),
-        "Total Tax": invoice_data.get("Total Tax", 0.0),
-        "Total Amount": invoice_data.get("Total Amount", 0.0),
-        "Sender Name": invoice_data.get("Sender", {}).get("Name", ""),
-        "Sender Address": invoice_data.get("Sender", {}).get("Address", ""),
-        "Recipient Name": invoice_data.get("Recipient", {}).get("Name", ""),
-        "Recipient Address": invoice_data.get("Recipient", {}).get("Address", ""),
-    }
-
-    # Initialize all expected bank fields to empty string
-    bank_row = {field: "" for field in EXPECTED_BANK_FIELDS}
-
-    # Override with actual values if present
-    bank_details = invoice_data.get("Bank Details", {}) or {}
-    for k, v in bank_details.items():
-        key_name = k if str(k).startswith("bank_") else f"bank_{k}"
-        if key_name in bank_row:
-            bank_row[key_name] = str(v).strip() if v is not None else ""
-
-    base_row.update(bank_row)
-
     if not line_items:
         # If no line items, create one row with invoice info only
-        row = base_row.copy()
+        row = {
+            "Invoice Number": invoice_data.get("Invoice Number", ""),
+            "Invoice Date": invoice_data.get("Invoice Date", ""),
+            "Due Date": invoice_data.get("Due Date", ""),
+            "Currency": invoice_data.get("Currency", ""),
+            "Subtotal": invoice_data.get("Subtotal", 0.0),
+            "Tax Percentage": invoice_data.get("Tax Percentage", 0.0),
+            "Total Tax": invoice_data.get("Total Tax", 0.0),
+            "Total Amount": invoice_data.get("Total Amount", 0.0),
+            "Sender Name": invoice_data.get("Sender", {}).get("Name", ""),
+            "Sender Address": invoice_data.get("Sender", {}).get("Address", ""),
+            "Recipient Name": invoice_data.get("Recipient", {}).get("Name", ""),
+            "Recipient Address": invoice_data.get("Recipient", {}).get("Address", ""),
+        }
+
+        # Flatten bank details
+        bank = invoice_data.get("Bank Details", {})
+        for k, v in bank.items():
+            # Avoid double-prefixing if key already contains 'bank_'
+            key_name = k if str(k).startswith("bank_") else f"bank_{k}"
+            row[key_name] = v
+
+        # Add empty line item fields
         row.update({
             "Item Description": "",
             "Item Quantity": 0,
@@ -559,7 +467,29 @@ def flatten_invoice_to_rows(invoice_data) -> list:
 
     # For each line item, create a row with all invoice context
     for item in line_items:
-        row = base_row.copy()
+        row = {
+            "Invoice Number": invoice_data.get("Invoice Number", ""),
+            "Invoice Date": invoice_data.get("Invoice Date", ""),
+            "Due Date": invoice_data.get("Due Date", ""),
+            "Currency": invoice_data.get("Currency", ""),
+            "Subtotal": invoice_data.get("Subtotal", 0.0),
+            "Tax Percentage": invoice_data.get("Tax Percentage", 0.0),
+            "Total Tax": invoice_data.get("Total Tax", 0.0),
+            "Total Amount": invoice_data.get("Total Amount", 0.0),
+            "Sender Name": invoice_data.get("Sender", {}).get("Name", ""),
+            "Sender Address": invoice_data.get("Sender", {}).get("Address", ""),
+            "Recipient Name": invoice_data.get("Recipient", {}).get("Name", ""),
+            "Recipient Address": invoice_data.get("Recipient", {}).get("Address", ""),
+        }
+
+        # Flatten bank details
+        bank = invoice_data.get("Bank Details", {})
+        for k, v in bank.items():
+            # Avoid double-prefixing if key already contains 'bank_'
+            key_name = k if str(k).startswith("bank_") else f"bank_{k}"
+            row[key_name] = v
+
+        # Add line item fields
         row.update({
             "Item Description": item.get("Description", ""),
             "Item Quantity": item.get("Quantity", 0),
@@ -568,6 +498,7 @@ def flatten_invoice_to_rows(invoice_data) -> list:
             "Item Tax": item.get("Tax", 0.0),
             "Item Line Total": item.get("Line Total", item.get("Amount", 0.0)),
         })
+
         rows.append(row)
 
     return rows
@@ -678,15 +609,15 @@ if not st.session_state.is_processing_batch and len(st.session_state.batch_resul
 # RESULTS VIEW — Show selector + editable form
 # ---------------------------
 elif len(st.session_state.batch_results) > 0:
-    # ---------------------------
-    # Global Download All — produce a single Excel file (concatenated rows) and trigger direct download
-    # ---------------------------
+            # ---------------------------
+        # Global Download All — produce a single Excel file (concatenated rows) and trigger direct download
+        # ---------------------------
     if st.button("📦 Download All Results (Excel)", key="download_all"):
-        # Collect rows from all invoices and concatenate into one DataFrame
+            # Collect rows from all invoices and concatenate into one DataFrame
         all_rows = []
         for file_hash, result in st.session_state.batch_results.items():
             rows = flatten_invoice_to_rows(result["edited_data"])
-            # Annotate rows with source file name so user can identify which invoice each row came from
+                # Annotate rows with source file name so user can identify which invoice each row came from
             for r in rows:
                 r["Source File"] = result.get("file_name", file_hash)
             all_rows.extend(rows)
@@ -696,13 +627,13 @@ elif len(st.session_state.batch_results) > 0:
         else:
             full_df = pd.DataFrame(all_rows)
 
-            # Reorder columns to put Source File first
+                # Reorder columns to put Source File first
             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]
 
-            # Try to write XLSX (preferred). If engine not available, fall back to CSV.
+                # Try to write XLSX (preferred). If engine not available, fall back to CSV.
             buffer = BytesIO()
             dl_filename = "all_extracted_invoices.xlsx"
             tried_xlsx = False
@@ -723,7 +654,7 @@ elif len(st.session_state.batch_results) > 0:
                 dl_filename = "all_extracted_invoices.csv"
                 mime = "text/csv"
 
-            # Trigger immediate download via a data URI and small HTML snippet
+                # Trigger immediate download via a data URI and small HTML snippet
             import base64
             import streamlit.components.v1 as components
             b64 = base64.b64encode(file_bytes).decode()
@@ -765,7 +696,8 @@ elif len(st.session_state.batch_results) > 0:
     # Get current file data
     current = st.session_state.batch_results[selected_hash]
     image = current["image"]
-    data = current["edited_data"]  # user edits go here
+    import copy
+    form_data = copy.deepcopy(current["edited_data"]) 
 
     # Layout
     left_col, right_col = st.columns([1, 1])
@@ -802,7 +734,6 @@ elif len(st.session_state.batch_results) > 0:
                     st.rerun()
                 except Exception as e:
                     st.error(f"Re-run failed: {e}")
-
         tabs = st.tabs(["Invoice Details", "Sender/Recipient info", "Bank Details", "Line Items"])
 
         st.markdown(
@@ -831,225 +762,88 @@ elif len(st.session_state.batch_results) > 0:
         # ---------- Invoice Details ----------
         with tabs[0]:
             with st.container():
-                widget_key = f"invoice_number_{selected_hash}"
-                st.text_input(
-                    "Invoice Number",
-                    value=data.get('Invoice Number', ''),
-                    key=widget_key,
-                    on_change=_update_flat_field,
-                    args=("Invoice Number", widget_key)
-                )
-
-                widget_key = f"invoice_date_text_{selected_hash}"
-                st.text_input(
-                    "Invoice Date",
-                    value=str(data.get('Invoice Date', '')).strip(),
-                    key=widget_key,
-                    on_change=_update_flat_field,
-                    args=("Invoice Date", widget_key)
-                )
-
-                widget_key = f"due_date_text_{selected_hash}"
-                st.text_input(
-                    "Due Date",
-                    value=str(data.get('Due Date', '')).strip(),
-                    key=widget_key,
-                    on_change=_update_flat_field,
-                    args=("Due Date", widget_key)
-                )
-
+                form_data['Invoice Number'] = st.text_input("Invoice Number", value=form_data.get('Invoice Number', ''), key=f"invoice_number_{selected_hash}")
+                form_data['Invoice Date'] = st.text_input("Invoice Date", value=str(form_data.get('Invoice Date', '')).strip(), key=f"invoice_date_text_{selected_hash}")
+                form_data['Due Date'] = st.text_input("Due Date", value=str(form_data.get('Due Date', '')).strip(), key=f"due_date_text_{selected_hash}")
                 curr_options = ['USD', 'EUR', 'GBP', 'INR', 'Other']
-                curr_value = data.get('Currency', 'USD')
+                curr_value = form_data.get('Currency', 'USD')
                 curr_index = curr_options.index(curr_value) if curr_value in curr_options else (len(curr_options) - 1)
-                new_curr = st.selectbox(
-                    "Currency",
-                    options=curr_options,
-                    index=curr_index,
-                    key=f"currency_select_{selected_hash}",
-                    on_change=_update_flat_field,
-                    args=("Currency", f"currency_select_{selected_hash}")
-                )
+                new_curr = st.selectbox("Currency", options=curr_options, index=curr_index, key=f"currency_select_{selected_hash}")
                 if new_curr == 'Other':
-                    widget_key = f"custom_currency_{selected_hash}"
-                    st.text_input(
-                        "Specify Currency",
-                        value=data.get('Currency', ''),
-                        key=widget_key,
-                        on_change=_update_flat_field,
-                        args=("Currency", widget_key)
-                    )
-
-                def safe_number_input(label, value, key, field_name):
-                    try:
-                        v = float(value)
-                    except Exception:
-                        v = 0.0
-                    st.number_input(
-                        label,
-                        value=v,
-                        key=key,
-                        on_change=_update_flat_field,
-                        args=(field_name, key)
-                    )
-
-                safe_number_input("Subtotal", data.get('Subtotal', 0.0), f"subtotal_{selected_hash}", "Subtotal")
-                safe_number_input("Tax Percentage", data.get('Tax Percentage', 0.0), f"tax_pct_{selected_hash}", "Tax Percentage")
-                safe_number_input("Total Tax", data.get('Total Tax', 0.0), f"total_tax_{selected_hash}", "Total Tax")
-                safe_number_input("Total Amount", data.get('Total Amount', 0.0), f"total_amount_{selected_hash}", "Total Amount")
+                    new_curr = st.text_input("Specify Currency", value=form_data.get('Currency', ''), key=f"custom_currency_{selected_hash}")
+                form_data['Currency'] = new_curr
+                form_data['Subtotal'] = safe_number_input("Subtotal", form_data.get('Subtotal', 0.0), f"subtotal_{selected_hash}")
+                form_data['Tax Percentage'] = safe_number_input("Tax Percentage", form_data.get('Tax Percentage', 0.0), f"tax_pct_{selected_hash}")
+                form_data['Total Tax'] = safe_number_input("Total Tax", form_data.get('Total Tax', 0.0), f"total_tax_{selected_hash}")
+                form_data['Total Amount'] = safe_number_input("Total Amount", form_data.get('Total Amount', 0.0), f"total_amount_{selected_hash}")
 
         # ---------- Sender / Recipient ----------
         with tabs[1]:
-            if 'Sender' not in data:
-                data['Sender'] = {'Name': '', 'Address': ''}
-            if 'Recipient' not in data:
-                data['Recipient'] = {'Name': '', 'Address': ''}
-
-            sender_info = data['Sender']
-            recipient_info = data['Recipient']
-
+            sender_info = form_data.get('Sender', {'Name': '', 'Address': ''})
+            recipient_info = form_data.get('Recipient', {'Name': '', 'Address': ''})
             with st.container():
-                widget_key = f"sender_name_{selected_hash}"
-                st.text_input(
-                    "Sender Name*",
-                    value=sender_info.get('Name', ''),
-                    key=widget_key,
-                    on_change=_update_nested_field,
-                    args=("Sender.Name", widget_key)
-                )
-
-                widget_key = f"sender_address_{selected_hash}"
-                st.text_area(
-                    "Sender Address*",
-                    value=sender_info.get('Address', ''),
-                    key=widget_key,
-                    on_change=_update_nested_field,
-                    args=("Sender.Address", widget_key)
-                )
-
-                widget_key = f"recipient_name_{selected_hash}"
-                st.text_input(
-                    "Recipient Name*",
-                    value=recipient_info.get('Name', ''),
-                    key=widget_key,
-                    on_change=_update_nested_field,
-                    args=("Recipient.Name", widget_key)
-                )
-
-                widget_key = f"recipient_address_{selected_hash}"
-                st.text_area(
-                    "Recipient Address*",
-                    value=recipient_info.get('Address', ''),
-                    key=widget_key,
-                    on_change=_update_nested_field,
-                    args=("Recipient.Address", widget_key)
-                )
-
+                sender_info['Name'] = st.text_input("Sender Name*", value=sender_info.get('Name', ''), key=f"sender_name_{selected_hash}")
+                sender_info['Address'] = st.text_area("Sender Address*", value=sender_info.get('Address', ''), key=f"sender_address_{selected_hash}")
+                recipient_info['Name'] = st.text_input("Recipient Name*", value=recipient_info.get('Name', ''), key=f"recipient_name_{selected_hash}")
+                recipient_info['Address'] = st.text_area("Recipient Address*", value=recipient_info.get('Address', ''), key=f"recipient_address_{selected_hash}")
                 if st.button("⇄ Swap", help="Swap sender and recipient information", key=f"swap_{selected_hash}"):
-                    data['Sender'], data['Recipient'] = data['Recipient'], data['Sender']
-                    st.session_state.batch_results[selected_hash]["edited_data"] = data
+                    form_data['Sender'], form_data['Recipient'] = form_data['Recipient'], form_data['Sender']
                     st.rerun()
 
         # ---------- Bank Details ----------
         with tabs[2]:
-            bank_info = data.get('Bank Details', {}) or {}
+            bank_info = form_data.get('Bank Details', {}) or {}
             with st.container():
-                widget_key = f"bank_name_{selected_hash}"
-                st.text_input(
-                    "Bank Name",
-                    value=bank_info.get('bank_name', ''),
-                    key=widget_key,
-                    on_change=_update_nested_field,
-                    args=("Bank Details.bank_name", widget_key)
-                )
-
-                widget_key = f"bank_account_{selected_hash}"
-                st.text_input(
-                    "Account Number",
-                    value=bank_info.get('bank_account_number', '') or bank_info.get('bank_acc_no',''),
-                    key=widget_key,
-                    on_change=_update_nested_field,
-                    args=("Bank Details.bank_account_number", widget_key)
-                )
-
-                widget_key = f"bank_acc_name_{selected_hash}"
-                st.text_input(
-                    "Bank Account Name",
-                    value=bank_info.get('bank_acc_name', '') or bank_info.get('bank_acc_name', ''),
-                    key=widget_key,
-                    on_change=_update_nested_field,
-                    args=("Bank Details.bank_acc_name", widget_key)
-                )
-
-                widget_key = f"iban_{selected_hash}"
-                st.text_input(
-                    "IBAN",
-                    value=bank_info.get('bank_iban', ''),
-                    key=widget_key,
-                    on_change=_update_nested_field,
-                    args=("Bank Details.bank_iban", widget_key)
-                )
-
-                widget_key = f"swift_code_{selected_hash}"
-                st.text_input(
-                    "SWIFT Code",
-                    value=bank_info.get('bank_swift', ''),
-                    key=widget_key,
-                    on_change=_update_nested_field,
-                    args=("Bank Details.bank_swift", widget_key)
-                )
-
-                widget_key = f"routing_{selected_hash}"
-                st.text_input(
-                    "Routing Number",
-                    value=bank_info.get('bank_routing', ''),
-                    key=widget_key,
-                    on_change=_update_nested_field,
-                    args=("Bank Details.bank_routing", widget_key)
-                )
-
-                widget_key = f"branch_{selected_hash}"
-                st.text_input(
-                    "Branch",
-                    value=bank_info.get('bank_branch', ''),
-                    key=widget_key,
-                    on_change=_update_nested_field,
-                    args=("Bank Details.bank_branch", widget_key)
-                )
+                bank_info['bank_name'] = st.text_input("Bank Name", value=bank_info.get('bank_name', ''), key=f"bank_name_{selected_hash}")
+                bank_info['bank_account_number'] = st.text_input("Account Number", value=bank_info.get('bank_account_number', '') or bank_info.get('bank_acc_no',''), key=f"bank_account_{selected_hash}")
+                bank_info['bank_acc_name'] = st.text_input("Bank Account Name", value=bank_info.get('bank_acc_name', '') or bank_info.get('bank_acc_name', ''), key=f"bank_acc_name_{selected_hash}")
+                bank_info['bank_iban'] = st.text_input("IBAN", value=bank_info.get('bank_iban', ''), key=f"iban_{selected_hash}")
+                bank_info['bank_swift'] = st.text_input("SWIFT Code", value=bank_info.get('bank_swift', ''), key=f"swift_code_{selected_hash}")
+                bank_info['bank_routing'] = st.text_input("Routing Number", value=bank_info.get('bank_routing', ''), key=f"routing_{selected_hash}")
+                bank_info['bank_branch'] = st.text_input("Branch", value=bank_info.get('bank_branch', ''), key=f"branch_{selected_hash}")
+            form_data['Bank Details'] = bank_info
 
         # ---------- Line Items ----------
         with tabs[3]:
             editor_key = f"item_editor_{selected_hash}"
-
-            item_rows = data.get('Itemized Data', [])
+            item_rows = form_data.get('Itemized Data', [])
             df = pd.DataFrame(item_rows)
             for col in ["Description", "Quantity", "Unit Price", "Amount", "Tax", "Line Total"]:
                 if col not in df.columns:
                     df[col] = ""
-
-            st.write("✏️ Edit line items below. Changes are saved automatically.")
-            st.data_editor(
+            st.write("✏️ Edit line items below. Press Enter or click outside a cell to confirm each edit.")
+            edited_df = st.data_editor(
                 df,
                 num_rows="dynamic",
                 key=editor_key,
                 use_container_width=True,
-                on_change=_update_line_items,
-                args=(editor_key,)
             )
-
-            if len(df) == 0:
+            if len(edited_df) == 0:
                 st.info("No line items found in the invoice.")
 
-        # Save button (per file) — now optional since edits auto-save
-        if st.button("💾 Save Edits for This File (Optional)", key=f"save_{selected_hash}"):
-            # This is now redundant but kept for user reassurance
-            st.success(f"✅ Edits already auto-saved for {current['file_name']}")
+        # Save button (per file)
+        if st.button("💾 Save Edits for This File", key=f"save_{selected_hash}"):
+            # Update line items from data editor
+            form_data['Itemized Data'] = edited_df.to_dict('records')
+            # Update session state with complete form data
+            st.session_state.batch_results[selected_hash]["edited_data"] = form_data
+            st.success(f"✅ Edits saved for {current['file_name']}")
 
         # Download buttons (per file)
         st.markdown("---")
         col_a, col_b, col_c = st.columns([1, 1, 1])
+        #with col_a:
+            #jsonl_str = json.dumps(data, ensure_ascii=False, indent=2)
+            #st.download_button(
+               # "📥 Download JSON",
+                #jsonl_str.encode("utf-8"),
+                #file_name=f"{Path(current['file_name']).stem}_extracted.json",
+                #mime="application/json",
+                #key=f"dl_json_{selected_hash}"
+            #)
         with col_b:
             # ✅ Flatten entire invoice into rows (one per line item)
-            rows = flatten_invoice_to_rows(data)
+            rows = flatten_invoice_to_rows(form_data)
             full_df = pd.DataFrame(rows)
             
             # Optional: Reorder columns for better readability
@@ -1076,6 +870,11 @@ elif len(st.session_state.batch_results) > 0:
                 mime="text/csv",
                 key=f"dl_csv_{selected_hash}"
             )
+# Global Download All — produce a single Excel file (concatenated rows) and trigger direct download
+
+
+# ---------------------------
+# PROCESSING STATE
 
 # ---------------------------
 # PROCESSING STATE — Show progress