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):
    # Prefer the repo working directory if writable, otherwise use /tmp
    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}")

# Ensure the .streamlit folder exists under HOME so Streamlit won't try to write to '/'
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}")

import json
from io import BytesIO
from datetime import datetime
from pathlib import Path
import hashlib
import zipfile
from typing import Optional, Dict, Any

import streamlit as st
import pdf2image
import pandas as pd
from PIL import Image
from huggingface_hub import login

# ---------------------------
# UI: main
# ---------------------------
st.set_page_config(page_title="Invoice Extractor (Donut) - Batch Mode", layout="wide")
st.title("Invoice Extraction")

# Reduce top margin and tighten layout
st.markdown(
    """
    <style>
        /* Reduce top padding of main block */
        .stApp {
            background-color: #E8E8E8 !important;
        }
        div.block-container {
            padding-top: 2rem;
            padding-bottom: 1rem;
        }
        /* Tighten title spacing */
        h1 {
            margin-top: 0.4rem !important;
            margin-bottom: 0.4rem !important;
            background-color: #E8E8E8 !important;
        }
        /* Reduce gap between columns */
        [data-testid="column"] {
            padding-top: 0.5rem;
            background-color: #E8E8E8 !important;
        }
    </style>
    """,
    unsafe_allow_html=True
)

# --- Secure token handling: prefer session-state -> env var -> Streamlit secrets; never hardcode or commit token ---
def _get_hf_token():
    if st.session_state.get("_hf_token"):
        return st.session_state.get("_hf_token"), "session"

    env_tok = os.getenv("HF_TOKEN")
    if env_tok:
        return env_tok, "env"

    try:
        project_secrets = Path(".streamlit/secrets.toml")
        user_secrets = Path.home() / ".streamlit" / "secrets.toml"
        if project_secrets.exists() or user_secrets.exists():
            sec = st.secrets.get("HF_TOKEN")
            if sec:
                return sec, "secrets"
    except Exception:
        pass

    return None, None

hf_token, hf_token_source = _get_hf_token()

# --- Interactive login fallback (development) ---
if hf_token is None:
    st.subheader("Login Token🔑")
    token_input = st.text_input("Enter your Login token (starts with 'hf_'):", type="password")
    if token_input:
        if not token_input.startswith("hf_"):
            st.error("Invalid token format. Token must start with 'hf_'.")
            st.stop()
        try:
            login(token_input)
            st.session_state["_hf_token"] = token_input
            st.session_state.logged_in = True
            st.success("Logged in successfully. Loading model...")
            st.rerun()
        except Exception as e:
            st.error(f"Failed to log in: {e}")
            st.stop()
    else:
        st.warning("Provide a token via the UI or set HF_TOKEN as an environment variable.")
        st.stop()
else:
    try:
        login(hf_token)
        st.session_state.logged_in = True
    except Exception as e:
        st.error(f"Failed to log in with {hf_token_source or 'unknown'} token: {e}")
        st.stop()

# ---------------------------
# Configuration (edit these)
# ---------------------------
HF_MODEL_ID = "Bhuvi13/model-V7"
TASK_PROMPT = "<s_cord-v2>"

# ---------------------------
# Helper: load model & processor (cached)
# ---------------------------
@st.cache_resource(show_spinner=False)
def load_model_and_processor(hf_model_id: str, task_prompt: str):
    try:
        import torch
        from transformers import VisionEncoderDecoderModel, DonutProcessor
    except Exception as e:
        raise RuntimeError(f"Failed to import ML libraries: {e}")

    try:
        processor = DonutProcessor.from_pretrained(hf_model_id)
        model = VisionEncoderDecoderModel.from_pretrained(hf_model_id)
    except Exception as e:
        raise RuntimeError(
            f"Failed to load model/processor from Hugging Face ({hf_model_id}). "
            "Make sure your HF token is available and model id is correct.\n"
            f"Original error: {e}"
        )

    model.eval()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    with torch.no_grad():
        decoder_input_ids = processor.tokenizer(
            task_prompt,
            add_special_tokens=False,
            return_tensors="pt"
        ).input_ids.to(device)

    return processor, model, device, decoder_input_ids


def run_inference_on_image(image: Image.Image, processor, model, device, decoder_input_ids):
    import torch

    pixel_values = processor(images=image, return_tensors="pt").pixel_values.to(device)

    gen_kwargs = dict(
        pixel_values=pixel_values,
        decoder_input_ids=decoder_input_ids,
        max_length=1536,
        num_beams=4,
        early_stopping=False,
    )

    with torch.no_grad():
        generated_ids = model.generate(**gen_kwargs)

    raw_pred = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
    cleaned = (raw_pred
               .replace(processor.tokenizer.eos_token or "", "")
               .replace(processor.tokenizer.pad_token or "", "")
               .strip())
    token2json_out = processor.token2json(cleaned)

    if isinstance(token2json_out, str):
        try:
            pred_dict = json.loads(token2json_out)
        except Exception:
            pred_dict = token2json_out
    else:
        pred_dict = token2json_out

    return pred_dict

# ---------------------------
# Helper: map donut output to our UI schema
# ---------------------------
def map_prediction_to_ui(pred):
    import json, re
    from collections import defaultdict

    # --- 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
        return None

    # --- normalize numeric strings like "1,800.00" -> float ---
    def clean_number(x):
        if x is None:
            return 0.0
        if isinstance(x, (int, float)):
            return float(x)
        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 ("", ".", "-", "-."):
            return 0.0
        try:
            return float(s)
        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

    # create empty UI template
    ui = {
        "Invoice Number": "",
        "Invoice Date": "",
        "Due Date": "",
        "Currency": "",
        "Subtotal": 0.0,
        "Tax Percentage": 0.0,
        "Total Tax": 0.0,
        "Total Amount": 0.0,
        "Sender": {"Name": "", "Address": ""},
        "Recipient": {"Name": "", "Address": ""},
        "Sender Name": "",
        "Sender Address": "",
        "Recipient Name": "",
        "Recipient Address": "",
        "Bank Details": {},
        "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

    # 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 ""

    # 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

    # 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 = []

    # --- Primary approach: detect explicit list-of-dicts candidates first (unchanged) ---
    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

    # --- Secondary approach: if parsed is a single dict that itself contains the item fields
    # This is important because your model sometimes emits a single item as a top-level dict
    # (e.g. {"descriptions":"...","quantity":"1.00","unit_price":"35,000.00",...}).
    # We must map that directly (do NOT rely on finding a list named "items").
    if not items_rows:
        single_candidate_keys = {k.strip().lower() for k in (parsed.keys() if isinstance(parsed, dict) else [])} if isinstance(parsed, dict) else set()
        # item-like keys we expect in the raw model output (explicitly include variants the model uses)
        item_like_keys = {"descriptions", "description", "desc", "item", "quantity", "qty", "unit_price", "unit price", "price", "amount", "line_total", "line total", "line_total", "line_total".lower(), "sku", "tax", "tax_amount"}
        if single_candidate_keys and single_candidate_keys.intersection(item_like_keys):
            # map the parsed dict as a single line item
            single_row = map_item_dict(parsed)
            if single_row is not None:
                items_rows.append(single_row)

    # 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"]}

    return ui

# ---------------------------
# Helper: flatten invoice to CSV rows
# ---------------------------
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.
    This version collects bank details from both:
      - invoice_data.get("Bank Details", {})  (nested dict style)
      - top-level keys in invoice_data that start with 'bank_'
    Ensures the expected bank_* columns always exist in the produced rows.
    """
    EXPECTED_BANK_FIELDS = [
        "bank_name",
        "bank_acc_no",
        "bank_acc_name",
        "bank_iban",
        "bank_swift",
        "bank_routing",
        "bank_branch"
    ]

    rows = []
    invoice_data = invoice_data or {}

    # Collect line items (if present)
    line_items = invoice_data.get("Itemized Data", []) or []

    # Collect bank details from nested dict (if any) and from top-level bank_ keys
    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

    # also collect flat top-level bank_* keys (these come from your form_data)
    for k, v in invoice_data.items():
        if isinstance(k, str) and k.lower().startswith("bank_"):
            bank_details[k] = v

    # ensure all expected bank fields are present (empty string if missing)
    for f in EXPECTED_BANK_FIELDS:
        bank_details.setdefault(f, "")

    # Helper to create base invoice row (shared for empty-items case and per-item rows)
    def base_invoice_info():
        return {
            "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 Name", "") or (invoice_data.get("Sender",{}) or {}).get("Name",""),
            "Sender Address": invoice_data.get("Sender Address", "") or (invoice_data.get("Sender",{}) or {}).get("Address",""),
            "Recipient Name": invoice_data.get("Recipient Name", "") or (invoice_data.get("Recipient",{}) or {}).get("Name",""),
            "Recipient Address": invoice_data.get("Recipient Address", "") or (invoice_data.get("Recipient",{}) or {}).get("Address",""),
        }

    # If no line items, emit a single invoice-only row (with empty item columns)
    if not line_items:
        row = base_invoice_info()
        # include all expected bank fields (consistent names)
        for k in EXPECTED_BANK_FIELDS:
            row[k] = bank_details.get(k, "")
        # item columns (empty)
        row.update({
            "Item Description": "",
            "Item Quantity": 0,
            "Item Unit Price": 0.0,
            "Item Amount": 0.0,
            "Item Tax": 0.0,
            "Item Line Total": 0.0,
        })
        rows.append(row)
        return rows

    # For each line item, create a row with all invoice context + bank fields
    for item in line_items:
        row = base_invoice_info()
        for k in EXPECTED_BANK_FIELDS:
            row[k] = bank_details.get(k, "")
        # try to read canonical item keys (safe .get)
        row.update({
            "Item Description": item.get("Description", "") if isinstance(item, dict) else "",
            "Item Quantity": item.get("Quantity", 0) if isinstance(item, dict) else 0,
            "Item Unit Price": item.get("Unit Price", 0.0) if isinstance(item, dict) else 0.0,
            "Item Amount": item.get("Amount", 0.0) if isinstance(item, dict) else 0.0,
            "Item Tax": item.get("Tax", 0.0) if isinstance(item, dict) else 0.0,
            "Item Line Total": item.get("Line Total", item.get("Amount", 0.0)) if isinstance(item, dict) else 0.0,
        })
        rows.append(row)

    return rows


# Load model once
try:
    with st.spinner("Loading model & processor (cached) ..."):
        processor, model, device, decoder_input_ids = load_model_and_processor(HF_MODEL_ID, TASK_PROMPT)
except Exception as e:
    st.error("Could not load model automatically. See details below.")
    st.exception(e)
    st.stop()

# Initialize batch state
if "batch_results" not in st.session_state:
    st.session_state.batch_results = {}  # {file_hash: {image, raw_pred, mapped_data, edited_data}}
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

# ---------------------------
# UPLOAD SECTION — Only shown if not processing and no results yet
# ---------------------------
if not st.session_state.is_processing_batch and len(st.session_state.batch_results) == 0:
    st.markdown("Upload one or more invoice images (png/jpg/jpeg/pdf). The app will process them one by one.")

    st.header("📤 Upload Invoices")

    uploaded_files = st.file_uploader(
        "Upload invoice images (png/jpg/jpeg/pdf)",
        type=["png", "jpg", "jpeg", "pdf"],
        accept_multiple_files=True
    )

    if uploaded_files and len(uploaded_files) > 0:
        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}")

            # Read and hash
            uploaded_bytes = uploaded_file.read()
            file_hash = hashlib.sha256(uploaded_bytes).hexdigest()

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

            # Convert to image
            image = None
            is_pdf = uploaded_file.name.lower().endswith('.pdf') or (hasattr(uploaded_file, 'type') and uploaded_file.type == 'application/pdf')
            if is_pdf:
                try:
                    from pdf2image import convert_from_bytes
                    pages = convert_from_bytes(uploaded_bytes, dpi=200)
                    if len(pages) > 0:
                        image = pages[0].convert("RGB")
                    else:
                        st.warning(f"PDF {uploaded_file.name} has no pages.")
                        continue
                except Exception as e:
                    st.warning(f"Could not render PDF {uploaded_file.name}. Ensure 'pdf2image' and poppler are installed.")
                    continue
            else:
                try:
                    image = Image.open(BytesIO(uploaded_bytes)).convert("RGB")
                except Exception as e:
                    st.warning(f"Failed to open {uploaded_file.name}.")
                    continue


            if image is None:
                continue

            # Run inference
            try:
                pred = run_inference_on_image(image, processor, model, device, decoder_input_ids)
                mapped = map_prediction_to_ui(pred)
            except Exception as e:
                st.warning(f"Error processing {uploaded_file.name}: {str(e)}")
                pred = None
                mapped = {}  # 👈 Ensure mapped is always a dict

            # ✅ SAFETY: Ensure mapped is a dict before copying
            safe_mapped = mapped if isinstance(mapped, dict) else {}

            # Save to session state
            st.session_state.batch_results[file_hash] = {
                "file_name": uploaded_file.name,
                "image": image,
                "raw_pred": pred,
                "mapped_data": mapped,
                "edited_data": safe_mapped.copy()  # editable copy — now safe
            }

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

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

# ---------------------------
# 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
    # ---------------------------
    if st.button("📦 Download All Results (Excel)", key="download_all"):
        # 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
            for r in rows:
                r["Source File"] = result.get("file_name", file_hash)
            all_rows.extend(rows)

        if len(all_rows) == 0:
            st.warning("No invoice data available to download.")
        else:
            full_df = pd.DataFrame(all_rows)

            # 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.
            buffer = BytesIO()
            dl_filename = "all_extracted_invoices.xlsx"
            tried_xlsx = False
            try:
                with pd.ExcelWriter(buffer, engine="openpyxl") as writer:
                    full_df.to_excel(writer, index=False, sheet_name="Invoices")
                tried_xlsx = True
                buffer.seek(0)
                file_bytes = buffer.read()
                mime = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
            except Exception:
                # Fallback to CSV
                buffer = BytesIO()
                csv_data = full_df.to_csv(index=False).encode("utf-8")
                buffer.write(csv_data)
                buffer.seek(0)
                file_bytes = buffer.read()
                dl_filename = "all_extracted_invoices.csv"
                mime = "text/csv"

            # 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()
            data_uri = f"data:{mime};base64,{b64}"

            auto_dl_html = f'''<html>
                <body>
                    <a id="dlLink" href="{data_uri}" download="{dl_filename}"></a>
                    <script>
                        const a = document.getElementById('dlLink');
                        a.click();
                    </script>
                </body>
            </html>'''

            components.html(auto_dl_html, height=0)

    # File selector
    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]
    st.session_state.current_file_hash = selected_hash

    # Back button
    if st.button("⬅️ Back to Upload"):
        st.session_state.batch_results.clear()
        st.session_state.current_file_hash = None
        st.session_state.is_processing_batch = False
        st.rerun()

    # Get current file data
    current = st.session_state.batch_results[selected_hash]
    image = current["image"]
    
    # ✅ FIX: Don't create a copy here - just reference the stored data
    form_data = current["edited_data"]

    # Layout
    left_col, right_col = st.columns([1, 1])

    # LEFT: Image + Raw Output
    with left_col:
        st.image(image, caption=current["file_name"], use_container_width=True)
        st.write(f"**File Hash:** {selected_hash[:8]}...")

        if current.get('raw_pred') is not None:
            with st.expander("🔍 Show raw model output"):
                st.json(current['raw_pred'])

    # RIGHT: Editable Form
    with right_col:
        st.subheader(f"Editable Invoice: {current['file_name']}")

        # ---------- Re-run (per-file) ----------
        if st.button("🔁 Re-Run", key=f"rerun_{selected_hash}"):
            # Re-run inference only for the selected file's image, update stored predictions and editable copy
            with st.spinner("Re-running inference for selected file..."):
                try:
                    pred = run_inference_on_image(image, processor, model, device, decoder_input_ids)
                    mapped = map_prediction_to_ui(pred)
                    safe_mapped = mapped if isinstance(mapped, dict) else {}

                    # Save updated results for this single file
                    st.session_state.batch_results[selected_hash]["raw_pred"] = pred
                    st.session_state.batch_results[selected_hash]["mapped_data"] = mapped
                    st.session_state.batch_results[selected_hash]["edited_data"] = safe_mapped.copy()

                    st.success("✅ Re-run complete — predictions updated for this file.")
                    # Refresh the UI so the new values appear in the form
                    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(
            """
            <style>
                div[data-testid="stTabs"] > div > div {
                    padding-bottom: 5px;
                    margin-top: 0;
                    background-color: #E8E8E8;
                }
                .stTextInput, .stNumberInput, .stSelectbox, .stTextArea, .stDateInput {
                    margin-bottom: 5px;
                    padding-bottom: 5px;
                }
                div[data-testid="stTabs"] {
                    background-color: #E8E8E8;
                }
                h3:first-of-type {
                    margin-top: 10px;
                }
            </style>
            """,
            unsafe_allow_html=True,
        )

        # ---------- Invoice Details ----------
        # ✅ FIX: Read values directly from widgets without assigning back to form_data
        with tabs[0]:
            with st.container():
                st.text_input("Invoice Number", value=form_data.get('Invoice Number', ''), key=f"invoice_number_{selected_hash}")
                st.text_input("Invoice Date", value=str(form_data.get('Invoice Date', '')).strip(), key=f"invoice_date_text_{selected_hash}")
                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 = form_data.get('Currency', 'USD')
                curr_index = curr_options.index(curr_value) if curr_value in curr_options else (len(curr_options) - 1)
                st.selectbox("Currency", options=curr_options, index=curr_index, key=f"currency_select_{selected_hash}")
                if st.session_state.get(f"currency_select_{selected_hash}") == 'Other':
                    st.text_input("Specify Currency", value=form_data.get('Currency', ''), key=f"custom_currency_{selected_hash}")
                
                safe_number_input("Subtotal", form_data.get('Subtotal', 0.0), f"subtotal_{selected_hash}")
                safe_number_input("Tax Percentage", form_data.get('Tax Percentage', 0.0), f"tax_pct_{selected_hash}")
                safe_number_input("Total Tax", form_data.get('Total Tax', 0.0), f"total_tax_{selected_hash}")
                safe_number_input("Total Amount", form_data.get('Total Amount', 0.0), f"total_amount_{selected_hash}")

        # ---------- Sender / Recipient ----------
        with tabs[1]:
            sender_name = form_data.get('Sender Name', '')
            sender_address = form_data.get('Sender Address', '')
            recipient_name = form_data.get('Recipient Name', '')
            recipient_address = form_data.get('Recipient Address', '')
            
            with st.container():
                st.text_input("Sender Name*", value=sender_name, key=f"sender_name_{selected_hash}")
                st.text_area("Sender Address*", value=sender_address, key=f"sender_address_{selected_hash}")
                st.text_input("Recipient Name*", value=recipient_name, key=f"recipient_name_{selected_hash}")
                st.text_area("Recipient Address*", value=recipient_address, key=f"recipient_address_{selected_hash}")
                
                if st.button("⇄ Swap", help="Swap sender and recipient information", key=f"swap_{selected_hash}"):
                    # Swap in session_state widget values
                    temp_name = st.session_state.get(f"sender_name_{selected_hash}", "")
                    temp_addr = st.session_state.get(f"sender_address_{selected_hash}", "")
                    
                    st.session_state[f"sender_name_{selected_hash}"] = st.session_state.get(f"recipient_name_{selected_hash}", "")
                    st.session_state[f"sender_address_{selected_hash}"] = st.session_state.get(f"recipient_address_{selected_hash}", "")
                    st.session_state[f"recipient_name_{selected_hash}"] = temp_name
                    st.session_state[f"recipient_address_{selected_hash}"] = temp_addr
                    st.rerun()

        # ---------- Bank Details ----------
        with tabs[2]:
            bank_details = form_data.get("Bank Details", {})
            if not isinstance(bank_details, dict):
                bank_details = {}
            
            bank_name = bank_details.get('bank_name', '')
            bank_acc_no = bank_details.get('bank_acc_no', '')
            bank_acc_name = bank_details.get('bank_acc_name', '')
            bank_iban = bank_details.get('bank_iban', '')
            bank_swift = bank_details.get('bank_swift', '')
            bank_routing = bank_details.get('bank_routing', '')
            bank_branch = bank_details.get('bank_branch', '')
            
            with st.container():
                st.text_input("Bank Name", value=bank_name, key=f"bank_name_{selected_hash}")
                st.text_input("Account Number", value=bank_acc_no, key=f"bank_acc_no_{selected_hash}")
                st.text_input("Bank Account Name", value=bank_acc_name, key=f"bank_acc_name_{selected_hash}")
                st.text_input("IBAN", value=bank_iban, key=f"iban_{selected_hash}")
                st.text_input("SWIFT Code", value=bank_swift, key=f"swift_code_{selected_hash}")
                st.text_input("Routing Number", value=bank_routing, key=f"routing_{selected_hash}")
                st.text_input("Branch", value=bank_branch, key=f"branch_{selected_hash}")

        # ---------- Line Items ----------
        with tabs[3]:
            editor_key = f"item_editor_{selected_hash}"
            item_rows = form_data.get('Itemized Data', []) or []

            # --- Normalize item keys produced by the model ---
            def normalize_item_keys(item):
                if not isinstance(item, dict):
                    return {
                        "Description": "",
                        "Quantity": "",
                        "Unit Price": "",
                        "Amount": "",
                        "Tax": "",
                        "Line Total": ""
                    }
                mapping = {
                    "Item Description": "Description",
                    "description": "Description",
                    "desc": "Description",
                    "Item Quantity": "Quantity",
                    "quantity": "Quantity",
                    "qty": "Quantity",
                    "Item Unit Price": "Unit Price",
                    "unit_price": "Unit Price",
                    "price": "Unit Price",
                    "Item Amount": "Amount",
                    "amount": "Amount",
                    "Item Tax": "Tax",
                    "tax": "Tax",
                    "Item Line Total": "Line Total",
                    "line_total": "Line Total",
                }
                new = {}
                for k, v in item.items():
                    key = mapping.get(k, mapping.get(str(k).lower(), k))
                    if key in ["Description", "Quantity", "Unit Price", "Amount", "Tax", "Line Total"]:
                        new[key] = v
                    else:
                        new[k] = v

                for kk in ["Description", "Quantity", "Unit Price", "Amount", "Tax", "Line Total"]:
                    if kk not in new:
                        new[kk] = ""

                return new

            normalized_items = [normalize_item_keys(it) for it in item_rows]
            df = pd.DataFrame(normalized_items)

            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. 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,
            )
            if len(edited_df) == 0:
                st.info("No line items found in the invoice.")

        # ✅ FIX: Save button now collects values from session_state widgets
        if st.button("💾 Save Edits for This File", key=f"save_{selected_hash}"):
            # Collect all values from session_state
            updated_data = {
                'Invoice Number': st.session_state.get(f"invoice_number_{selected_hash}", ""),
                'Invoice Date': st.session_state.get(f"invoice_date_text_{selected_hash}", ""),
                'Due Date': st.session_state.get(f"due_date_text_{selected_hash}", ""),
                'Currency': st.session_state.get(f"custom_currency_{selected_hash}", "") if st.session_state.get(f"currency_select_{selected_hash}") == 'Other' else st.session_state.get(f"currency_select_{selected_hash}", "USD"),
                'Subtotal': st.session_state.get(f"subtotal_{selected_hash}", 0.0),
                'Tax Percentage': st.session_state.get(f"tax_pct_{selected_hash}", 0.0),
                'Total Tax': st.session_state.get(f"total_tax_{selected_hash}", 0.0),
                '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_name_{selected_hash}", ""),
                    'bank_acc_no': st.session_state.get(f"bank_acc_no_{selected_hash}", ""),
                    'bank_acc_name': st.session_state.get(f"bank_acc_name_{selected_hash}", ""),
                    'bank_iban': st.session_state.get(f"iban_{selected_hash}", ""),
                    'bank_swift': st.session_state.get(f"swift_code_{selected_hash}", ""),
                    'bank_routing': st.session_state.get(f"routing_{selected_hash}", ""),
                    'bank_branch': st.session_state.get(f"branch_{selected_hash}", "")
                },
                'Itemized Data': edited_df.to_dict('records')
            }
            
            # Also set convenience fields
            updated_data['Sender'] = {"Name": updated_data['Sender Name'], "Address": updated_data['Sender Address']}
            updated_data['Recipient'] = {"Name": updated_data['Recipient Name'], "Address": updated_data['Recipient Address']}
            
            # Update session state
            st.session_state.batch_results[selected_hash]["edited_data"] = updated_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_b:
            # Use the saved edited_data (not the temporary form_data)
            rows = flatten_invoice_to_rows(st.session_state.batch_results[selected_hash]["edited_data"])
            full_df = pd.DataFrame(rows)
            
            # Optional: Reorder columns for better readability
            desired_col_order = [
                "Invoice Number", "Invoice Date", "Due Date", "Currency",
                "Subtotal", "Tax Percentage", "Total Tax", "Total Amount",
                "Sender Name", "Sender Address", "Recipient Name", "Recipient Address",
                "bank_name", "bank_acc_no", "bank_acc_name", "bank_iban", "bank_swift", "bank_routing", "bank_branch", 
                "Item Description", "Item Quantity", "Item Unit Price", "Item Amount", "Item Tax", "Item Line Total"
            ]
            # Keep only columns that exist
            existing_cols = [col for col in desired_col_order if col in full_df.columns]
            # Add any extra columns that weren't in desired order
            remaining_cols = [col for col in full_df.columns if col not in existing_cols]
            final_col_order = existing_cols + remaining_cols
            
            full_df = full_df[final_col_order]

            csv_bytes = full_df.to_csv(index=False).encode("utf-8")
            st.download_button(
                "📥 Download Full Invoice CSV",
                csv_bytes,
                file_name=f"{Path(current['file_name']).stem}_full.csv",
                mime="text/csv",
                key=f"dl_csv_{selected_hash}"
            )

# ---------------------------
# PROCESSING STATE — Show progress
# ---------------------------
elif st.session_state.is_processing_batch:
    st.info("⏳ Processing batch... Please wait.")
    st.progress(0)  # Placeholder — real progress handled in upload section

# ---------------------------
# DEFAULT — Nothing to show
# ---------------------------
else:
    pass