# backend/app.py — updated to use JWT auth (via backend/auth.py)

import os
import json
import csv
from pathlib import Path
from datetime import datetime, timedelta
from typing import List, Dict, Any, Optional, Set, Union
import inspect  # << added for signature-checking when calling DB helpers
import re
import logging

from fastapi import FastAPI, Request, Depends, HTTPException, Query, Body
from fastapi.responses import RedirectResponse, FileResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from openpyxl import load_workbook, Workbook
import boto3  # <-- S3 fallback

# Load .env if present
try:
    from dotenv import load_dotenv  # type: ignore
    load_dotenv()
except Exception:
    pass

# ===== JWT auth wiring =====
# Requires: backend/auth.py (provided) + .env with ADMIN_USER, ADMIN_HASHED_PW, JWT_SECRET, ...
from auth import router as auth_router, require_user, decode_token, COOKIE_NAME

logger = logging.getLogger(__name__)

# =========================================================
# Paths
# =========================================================
BASE_DIR = Path(__file__).resolve().parent
STATIC_DIR = BASE_DIR / "static"
DATA_DIR = BASE_DIR / "data"
CONV_DIR = BASE_DIR / "conv"

# Template config (triple header)
TEMPLATE_DEFAULT = STATIC_DIR / "templates" / "ARDEM_VITU_TEMPLATE.xlsx"
XLSX_TEMPLATE_PATH = Path(os.getenv("XLSX_TEMPLATE_PATH", str(TEMPLATE_DEFAULT)))
XLSX_TEMPLATE_SHEET = os.getenv("XLSX_TEMPLATE_SHEET", "")  # empty -> active sheet
HEADER_ROWS = int(os.getenv("XLSX_HEADER_ROWS", "3"))       # triple header by default
DATA_START_ROW = int(os.getenv("XLSX_DATA_START_ROW", str(HEADER_ROWS + 1)))

# New: templates directory + global fallback path (keeps existing fallback)
TEMPLATE_DIR = Path(os.getenv("XLSX_TEMPLATES_DIR", str(STATIC_DIR / "templates")))
GLOBAL_XLSX_TEMPLATE_PATH = Path(os.getenv("XLSX_TEMPLATE_PATH", str(TEMPLATE_DEFAULT)))

# Header synonyms for key columns
FILENAME_SYNONYMS: Set[str] = {
    "filename", "file name", "file", "document", "pdf"
}
TXN_SYNONYMS: Set[str] = {
    "transaction id", "transactionid", "txn id", "txn",
    "transaction #", "transaction number", "transaction no",
    "trans id", "trans #"
}

app = FastAPI(title="ARDEM VITU Backend", version="2.3")

# Static frontend
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")

# Include auth routes (/auth/login, /auth/logout, /auth/me)
app.include_router(auth_router, prefix="/auth", tags=["auth"])

# =========================================================
# DB toggle (minimal switch)
# =========================================================
USE_DB = os.getenv("DATA_SOURCE", "files").lower() == "db"
if USE_DB:
    # db_repo.py should live next to this file (or adjust import if placed in utils/)
    from db_repo import (
        list_files as db_list_files,
        get_fields as db_get_fields,
        list_audit as db_list_audit,
        get_s3_key as db_get_s3_key,         # used by /file/pdf
        upsert_output as db_upsert_output,   # used by /api/file/update
        list_rejected as db_list_rejected,   # used by /api/audit/rejected
        # ProcessData helpers
        get_processdata_by_date_filename,
        get_processdata_flags_by_date_filename,   # << NEW import
        upsert_processdata,
        get_guid_by_date_filename,
        get_processdata_by_guid,
    )

# =========================================================
# Helper functions
# =========================================================
def _norm_date(s: Optional[str]) -> Optional[str]:
    """Flexible date parsing (YYYY-MM-DD or MM/DD/YYYY) -> YYYY-MM-DD"""
    if not s:
        return None
    s = s.strip()
    for fmt in ("%Y-%m-%d", "%m/%d/%Y"):
        try:
            return datetime.strptime(s, fmt).date().isoformat()
        except Exception:
            pass
    raise HTTPException(status_code=400, detail=f"Invalid date format: {s}. Use YYYY-MM-DD or MM/DD/YYYY.")

def _read_json(path: Path, default: Any = None) -> Any:
    if path.exists():
        try:
            with path.open("r", encoding="utf-8") as f:
                return json.load(f)
        except Exception:
            return default
    return default

def _read_jsonl(path: Path) -> List[dict]:
    if not path.exists():
        return []
    rows = []
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                try:
                    rows.append(json.loads(line))
                except Exception:
                    continue
    return rows

def _read_csv(path: Path) -> List[dict]:
    if not path.exists():
        return []
    with path.open("r", encoding="utf-8") as f:
        return list(csv.DictReader(f))

def _date_folder(date: str) -> Path:
    return DATA_DIR / date

def _iter_dates(from_str: str, to_str: str) -> List[str]:
    d0 = datetime.strptime(from_str, "%Y-%m-%d").date()
    d1 = datetime.strptime(to_str, "%Y-%m-%d").date()
    if d1 < d0:
        d0, d1 = d1, d0
    out = []
    cur = d0
    while cur <= d1:
        out.append(cur.isoformat())
        cur += timedelta(days=1)
    return out

# ---------------- Text normalization helpers for header matching ----------------
def _normalize_header(s: Any) -> str:
    """Lowercase + collapse whitespace for robust header matching."""
    if s is None:
        return ""
    t = str(s).strip().lower()
    return " ".join(t.split())

def _collect_bottom_headers(ws) -> Dict[str, int]:
    """
    Read the bottom header row (HEADER_ROWS) and build {normalized_header -> column_index}.
    Assumes final labels live in the lowest header row.
    """
    row = HEADER_ROWS
    headers: Dict[str, int] = {}
    for col in range(1, ws.max_column + 1):
        val = ws.cell(row=row, column=col).value
        key = _normalize_header(val)
        if key:
            headers[key] = col
    return headers

def _find_header_col_by_synonyms(ws, synonyms: Set[str]) -> Optional[int]:
    """Locate a header column by scanning all header rows + bottom header map."""
    # scan all header rows (1..HEADER_ROWS)
    for r in range(1, min(HEADER_ROWS, ws.max_row) + 1):
        for c in range(1, ws.max_column + 1):
            val = ws.cell(row=r, column=c).value
            if _normalize_header(val) in synonyms:
                return c
    # check bottom headers map too
    bottom = _collect_bottom_headers(ws)
    for name in synonyms:
        if name in bottom:
            return bottom[name]
    return None

def _get_field_by_synonyms(fields: Dict[str, Any], synonyms: Set[str]) -> Optional[Any]:
    """Find a field value by any of the normalized key synonyms."""
    for k, v in (fields or {}).items():
        if _normalize_header(k) in synonyms:
            return v
    return None

def _coalesce_value(v: Any) -> Any:
    """Keep primitives as-is; JSON-encode other types."""
    return v if isinstance(v, (str, int, float, bool)) or v is None else json.dumps(v, ensure_ascii=False)


# ---------------- Excel helpers for in-place updates (file-mode compatibility) ----------------
def _excel_find_header(ws):
    headers = {}
    if ws.max_row < 1:
        ws.append(["filename"])
    if ws.max_column < 1:
        ws.cell(row=1, column=1).value = "filename"
    for col in range(1, ws.max_column + 1):
        val = ws.cell(row=1, column=col).value
        if val:
            headers[str(val).strip().lower()] = col
    return headers

def _excel_ensure_column(ws, headers, header_name: str) -> int:
    key = header_name.strip().lower()
    if key in headers:
        return headers[key]
    new_col = ws.max_column + 1
    ws.cell(row=1, column=new_col).value = header_name
    headers[key] = new_col
    return new_col

def _excel_find_or_create_row_for_filename(ws, headers, filename: str) -> int:
    if "filename" not in headers:
        _excel_ensure_column(ws, headers, "filename")
    col = headers["filename"]
    for r in range(2, ws.max_row + 1):
        val = ws.cell(row=r, column=col).value
        if (str(val or "")).strip().lower() == filename.strip().lower():
            return r
    new_row = ws.max_row + 1
    ws.cell(row=new_row, column=col).value = filename
    return new_row

def _to_excel_cell_value(v: Any):
    return v if isinstance(v, (str, int, float, bool)) or v is None else json.dumps(v, ensure_ascii=False)

def _update_excel_in_place(xlsx_path: Path, filename: str, fields: Dict[str, Any]):
    """
    Simple "file-mode" editor: single-row, single-header row sheets.
    Kept for compatibility with /api/file/update when not rebuilding at download time.
    """
    xlsx_path.parent.mkdir(parents=True, exist_ok=True)
    if xlsx_path.exists():
        wb = load_workbook(str(xlsx_path))
    else:
        wb = Workbook()
    ws = wb.active

    headers = _excel_find_header(ws)
    row_idx = _excel_find_or_create_row_for_filename(ws, headers, filename)

    for k, v in (fields or {}).items():
        if k is None:
            continue
        col_idx = _excel_ensure_column(ws, headers, str(k))
        ws.cell(row=row_idx, column=col_idx).value = _to_excel_cell_value(v)

    wb.save(str(xlsx_path))

# ------------------ Helpers for QC/suspect detection & recomputing totals -------------
# Simple heuristics used by UI + server to decide suspectness of an entered value.
NUMERIC_VALID_RE = re.compile(r'^(?:\d{1,3}(?:,\d{3})*|\d+)\.\d{2}$')
TXN_RE = re.compile(r'^\d{8}$')
PROCESS_DATE_RE = re.compile(r'^(0[1-9]|1[0-2])\/(0[1-9]|[12]\d|3[01])\/\d{4}$')
STATE_RE = re.compile(r'^[A-Z]{2}$')

NUMERIC_FIELDS = {
    "Taxable Selling Price / Purchase Price",
    "Sales Tax Credit Reciprocity",
    "Sales Tax",
    "Sales Tax Rate",
    "Total State Taxes",
    "Registration Fee",
    "Title Fee",
    "Plate Fee",
    "Lien Fee",
    "Transfer Fee",
    "Motor Vehicle Fees",
    "Local Use / County Fee",
    "Increase Fee",
    "Exchange Fee",
    "Specialty Plate Fee",
    "Additional Charge",
    "Total Title & Registration Fee",
    "Total Penalties",
    "Total Other Fees",
    "Total Sales Taxes + Reg Fee + Penalties",
    # GA-specific numeric fields
    "Taxable Selling Price/ Purchase",
    "Trade in Value",
    "Title Ad Valorem Tax",
    "GA Sales Tax",
    "Annual Ad Valorem",
    "Title Ad Valorem Dealer Penalty",
    "Title Ad Valorem Casual Interest",
    "Tag Fee",
    "Alternative Fuel Vehicle Fee",
    "Mail Fee",
    "Highway Impact Fee",
    "Title Late Penalty",
    "Title Ad Valorem Casual Penalty"
}

FEE_FIELDS_FOR_TITLE_REG = [
    "Registration Fee",
    "Title Fee",
    "Plate Fee",
    "Lien Fee",
    "Transfer Fee",
    "Motor Vehicle Fees",
    "Local Use / County Fee",
    "Increase Fee",
    "Exchange Fee",
    "Specialty Plate Fee",
    "Additional Charge",
    # GA-specific fee fields
    "Tag Fee",
    "Alternative Fuel Vehicle Fee",
    "Mail Fee",
    "Highway Impact Fee"
]

def _is_value_suspect(label: str, value: Any) -> bool:
    """
    Lightweight heuristic (patched):
      - Transaction ID → now allows 8–10 digits
      - Process Date → allows both MM/DD/YYYY and YYYY-MM-DD
    """

    if value is None:
        return False

    s = str(value).strip()

    if label in NUMERIC_FIELDS:
        if s == "":
            return True

        tokens = re.findall(r'[\d,\.]+', s)
        if len(tokens) > 1:
            return True

        if re.search(r'[A-Za-z]', s):
            return True

        if re.match(r'^\d+,\d{2}$', s) and not NUMERIC_VALID_RE.match(s):
            return True

        if not NUMERIC_VALID_RE.match(s):
            return True

        return False

    # 🔥 FIX 1: Transaction ID (8–10 digits instead of 8)
    if label == "Transaction ID":
        if s == "":
            return True

        if not re.match(r"^\d{8,10}$", s):
            return True

        return False

    # 🔥 FIX 2: Process Date (support both formats)
    if label == "Process Date":
        if s == "":
            return True

        if not (
            re.match(r"^\d{2}/\d{2}/\d{4}$", s) or
            re.match(r"^\d{4}-\d{2}-\d{2}$", s)
        ):
            return True

        return False

    if label == "Titling State":
        if s == "":
            return True

        if not STATE_RE.match(s.upper()):
            return True

        return False

    return False

def _to_float_safe(x):
    try:
        if x is None or x == "":
            return 0.0
        s = str(x).replace(",", "").strip()
        return float(s)
    except Exception:
        return 0.0

def _format_with_commas(value: float) -> str:
    """
    Format a float with thousands separators and two decimal places.
    Example: 25400.0 -> "25,400.00"
    """
    try:
        return f"{value:,.2f}"
    except Exception:
        return f"{value:.2f}"

def recompute_totals(fields: Dict[str, Any]) -> None:
    """
    Mutates `fields` in-place:
      - Recomputes "Total Title & Registration Fee" as sum of fee components.
      - Sets "Total Penalties" and "Total Other Fees" to "0.00" if missing.
      - Ensures "State Receipt Total" is blank string.
      - Recomputes "Total Sales Taxes + Reg Fee + Penalties" as Sales Tax + title_reg_total + penalties.
    """
    # ensure numeric fields exist as strings
    for nf in NUMERIC_FIELDS:
        if nf not in fields:
            fields[nf] = "" if nf in ("State Receipt Total",) else "0.00"

    # compute title/reg sum
    total_title_reg = 0.0
    for ff in FEE_FIELDS_FOR_TITLE_REG:
        total_title_reg += _to_float_safe(fields.get(ff, "0.00"))
    fields["Total Title & Registration Fee"] = f"{total_title_reg:.2f}"

    # forced fields
    fields["Total Penalties"] = fields.get("Total Penalties", "0.00") or "0.00"
    fields["Total Other Fees"] = fields.get("Total Other Fees", "0.00") or "0.00"
    fields["State Receipt Total"] = ""  # always blank

    # Handle state-specific sales tax fields
    sales_tax_val = 0.0
    if "GA Sales Tax" in fields and fields.get("GA Sales Tax"):
        # GA state uses "GA Sales Tax"
        sales_tax_val = _to_float_safe(fields.get("GA Sales Tax", "0.00"))
    else:
        # Default/CT uses "Sales Tax"
        sales_tax_val = _to_float_safe(fields.get("Sales Tax", "0.00"))
    
    penalties_val = _to_float_safe(fields.get("Total Penalties", "0.00"))
    total_sales_plus_reg_plus_penalties = sales_tax_val + total_title_reg + penalties_val
    fields["Total Sales Taxes + Reg Fee + Penalties"] = f"{total_sales_plus_reg_plus_penalties:.2f}"

# ---------------- Unified loader: DB or files ----------------
def _load_rows(date: Optional[str], from_date: Optional[str], to_date: Optional[str], state: Optional[str]) -> List[dict]:
    date = _norm_date(date) if date else None
    from_date = _norm_date(from_date) if from_date else None
    to_date = _norm_date(to_date) if to_date else None

    if USE_DB:
        if from_date and to_date:
            return db_list_files(from_date, to_date, state)
        elif date:
            return db_list_files(date, date, state)
        return []

    # file mode
    rows: List[dict] = []
    if from_date and to_date:
        for d in _iter_dates(from_date, to_date):
            snap = _read_json(_date_folder(d) / f"files_{d}.json", [])
            if snap:
                for r in snap:
                    r["date"] = d
                rows.extend(snap)
    elif date:
        snap = _read_json(_date_folder(date) / f"files_{date}.json", [])
        if snap:
            for r in snap:
                r["date"] = date
            rows.extend(snap)

    if state:
        rows = [
            r for r in rows
            if r.get("state") == state
            or Path(r.get("filename","")).name.startswith(f"{state}-")
        ]
    return rows

# ----------------------- S3 helpers (fallback for PDFs) -----------------------
def _s3_client():
    return boto3.client(
        "s3",
        aws_access_key_id=os.getenv("AWS_ACCESS_KEY"),
        aws_secret_access_key=os.getenv("AWS_SECRET_KEY"),
        region_name=os.getenv("AWS_REGION"),
    )

def _presign_s3(key: str, expires: int = 600) -> str:
    bucket = os.getenv("BUCKET_NAME")
    if not bucket:
        raise HTTPException(status_code=500, detail="BUCKET_NAME not configured")
    try:
        return _s3_client().generate_presigned_url(
            "get_object", Params={"Bucket": bucket, "Key": key}, ExpiresIn=expires
        )
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"S3 presign failed: {e}")

# ---- NEW: build S3 prefix for an arbitrary date + optional state ----
def _build_s3_prefix_for(date_iso: str, state: Optional[str]) -> str:
    """
    Build prefix matching your pipeline layout:
      VITU/{YYYY}/{MONTH_UPPER}/{MM.DD}_Files/{STATE_MM.DD}/   (if state provided)
      VITU/{YYYY}/{MONTH_UPPER}/{MM.DD}_Files/                 (if state not provided)
    Example:
      date_iso = "2025-09-17", state="PA" -> "VITU/2025/SEPTEMBER/09.17_Files/PA_09.17/"
    """
    try:
        dt = datetime.fromisoformat(date_iso)
    except Exception:
        # fallback: return empty prefix so S3 call will count nothing
        return ""
    year = dt.strftime("%Y")
    month_upper = dt.strftime("%B").upper()
    mmdd = dt.strftime("%m.%d")
    folder = f"{mmdd}_Files"
    if state:
        st = state.strip().upper()
        return f"VITU/{year}/{month_upper}/{folder}/{st}_{mmdd}/"
    return f"VITU/{year}/{month_upper}/{folder}/"

def _s3_count_objects_with_prefix(prefix: str) -> int:
    """
    Count S3 objects under the given prefix using list_objects_v2 with pagination.
    Returns total count (0 if bucket not configured or error).
    """
    bucket = os.getenv("BUCKET_NAME")
    if not bucket or not prefix:
        return 0
    client = _s3_client()
    paginator = client.get_paginator("list_objects_v2")
    total = 0
    try:
        for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
            if "Contents" in page:
                total += len(page["Contents"])
    except Exception as e:
        # Log server-side; return 0 as safe fallback (avoid crashing KPIs)
        print(f"[WARN] S3 count failed for prefix {prefix}: {e}")
        return 0
    return total

# =========================================================
# Auth-aware root (redirect based on JWT cookie)
# =========================================================
@app.get("/")
async def root(request: Request):
    token = request.cookies.get(COOKIE_NAME)
    if token:
        try:
            decode_token(token)  # will raise on invalid/expired
            return RedirectResponse(url="/static/dashboard.html")
        except Exception:
            pass
    return RedirectResponse(url="/static/index.html")

# =========================================================
# Data APIs (protected by JWT)
# =========================================================
@app.get("/api/kpis")
def get_kpis(
    request: Request,
    date: Optional[str] = None,
    from_date: Optional[str] = Query(None, alias="from"),
    to_date: Optional[str] = Query(None, alias="to"),
    state: Optional[str] = None,
    user=Depends(require_user),
):
    """
    DB-driven KPIs per your rules:
      - received: total DB rows for the date/range/state (DB-only)
      - processed: unique documents among those rows (priority: filename -> pdf_md5 -> txn_mdmd5)
                   duplicates are identified vs. previously existing DB rows (dates strictly before window)
                   and earlier rows within the same window (ordered by 'started')
      - delivered: same as processed
      - rejected: number of rows in the requested window with status == 'DUPLICATE' (fetched from DB when USE_DB)
    """
    # Normalize date inputs
    date = _norm_date(date) if date else None
    from_date = _norm_date(from_date) if from_date else None
    to_date = _norm_date(to_date) if to_date else None

    # Load rows for the requested window (DB or file-mode fallback)
    rows = _load_rows(date, from_date, to_date, state)

    # Received = DB row count for this query window (DB-only intent)
    received = len(rows)

    # Prepare sets for "seen" keys (historical + within-window)
    seen_filenames = set()
    seen_pdf_md5 = set()
    seen_txn_md5 = set()

    # Fetch historical rows (strictly before window start) when in DB mode
    historical = []
    if USE_DB:
        # determine window start date (earliest date in query)
        if date:
            window_start = date
        elif from_date:
            window_start = from_date
        else:
            window_start = None

        if window_start:
            try:
                # fetch rows from a very early date until day before window_start
                day_before = (datetime.fromisoformat(window_start).date() - timedelta(days=1)).isoformat()
                historical = db_list_files("1900-01-01", day_before, state)
            except Exception:
                historical = []
        else:
            historical = []

        # populate seen_* sets from historical rows
        for h in historical:
            fn = str(h.get("filename") or "").strip().lower()
            if fn:
                seen_filenames.add(fn)
            p = (h.get("pdf_md5") or h.get("pdfMd5") or "") or ""
            if p:
                seen_pdf_md5.add(str(p).strip().lower())
            t = (h.get("txn_md5") or h.get("txnMd5") or "") or ""
            if t:
                seen_txn_md5.add(str(t).strip().lower())

    # Sort current-window rows by 'started' (deterministic) to mark earlier files first
    def _row_started_key(r):
        try:
            st = r.get("started")
            if st:
                # if it's already a datetime object, this will be fine; else parse string
                if isinstance(st, datetime):
                    return st
                return datetime.fromisoformat(str(st))
        except Exception:
            pass
        # fallback: deterministic string key
        return str(r.get("filename") or "")

    rows_sorted = sorted(rows, key=_row_started_key)

    # Count processed = number of rows in window that introduce a NEW unique key
    processed_count = 0
    for r in rows_sorted:
        fn = str(r.get("filename") or "").strip().lower()
        p = (r.get("pdf_md5") or r.get("pdfMd5") or "") or ""
        p = str(p).strip().lower() if p else ""
        t = (r.get("txn_md5") or r.get("txnMd5") or "") or ""
        t = str(t).strip().lower() if t else ""

        is_duplicate = False
        # check filename first
        if fn and fn in seen_filenames:
            is_duplicate = True
        # then pdf_md5
        elif p and p in seen_pdf_md5:
            is_duplicate = True
        # then txn_md5
        elif t and t in seen_txn_md5:
            is_duplicate = True

        if not is_duplicate:
            processed_count += 1
            # add to seen sets
            if fn:
                seen_filenames.add(fn)
            if p:
                seen_pdf_md5.add(p)
            if t:
                seen_txn_md5.add(t)
        else:
            # still add keys to seen sets to handle chains within window
            if fn:
                seen_filenames.add(fn)
            if p:
                seen_pdf_md5.add(p)
            if t:
                seen_txn_md5.add(t)

    processed = processed_count
    delivered = processed

    # Rejected = rows in window whose status == 'DUPLICATE'
    # When USE_DB is true, prefer querying db_list_rejected (which returns rows with status='DUPLICATE' or rejection metadata).
    rejected = 0
    if USE_DB:
        try:
            rejected_rows = []
            if date:
                rejected_rows = db_list_rejected(date)
            elif from_date and to_date:
                # accumulate per-day rejected rows across range
                for d in _iter_dates(from_date, to_date):
                    try:
                        rejected_rows.extend(db_list_rejected(d))
                    except Exception:
                        # ignore day-level failures, continue
                        pass
            else:
                rejected_rows = []
            # Count those explicitly marked DUPLICATE
            rejected = sum(1 for r in rejected_rows if str(r.get("status", "")).upper() == "DUPLICATE")
        except Exception as e:
            print(f"[WARN] db_list_rejected failed in /api/kpis: {e}")
            # fallback to counting from loaded rows
            rejected = sum(1 for r in rows if str(r.get("status", "")).upper() == "DUPLICATE")
    else:
        rejected = sum(1 for r in rows if str(r.get("status", "")).upper() == "DUPLICATE")

    return {
        "received": received,
        "processed": processed,
        "delivered": delivered,
        "rejected": rejected
    }

@app.get("/api/charts")
def get_charts(
    date: Optional[str] = None,
    from_date: Optional[str] = Query(None, alias="from"),
    to_date: Optional[str] = Query(None, alias="to"),
    state: Optional[str] = None,
    user=Depends(require_user),
):
    """
    Charts endpoint now uses the same processed/rejected uniqueness logic as /api/kpis
    so dashboard tiles and charts remain consistent.
    """
    date = _norm_date(date) if date else None
    from_date = _norm_date(from_date) if from_date else None
    to_date = _norm_date(to_date) if to_date else None

    # Load rows for the requested window
    rows = _load_rows(date, from_date, to_date, state)

    # -------------------------
    # Compute processed (unique docs) using same logic as /api/kpis
    # -------------------------
    # Prepare seen sets and historical rows
    seen_filenames = set()
    seen_pdf_md5 = set()
    seen_txn_md5 = set()

    historical = []
    if USE_DB:
        if date:
            window_start = date
        elif from_date:
            window_start = from_date
        else:
            window_start = None

        if window_start:
            try:
                day_before = (datetime.fromisoformat(window_start).date() - timedelta(days=1)).isoformat()
                historical = db_list_files("1900-01-01", day_before, state)
            except Exception:
                historical = []
        else:
            historical = []

        for h in historical:
            fn = str(h.get("filename") or "").strip().lower()
            if fn:
                seen_filenames.add(fn)
            p = (h.get("pdf_md5") or h.get("pdfMd5") or "") or ""
            if p:
                seen_pdf_md5.add(str(p).strip().lower())
            t = (h.get("txn_md5") or h.get("txnMd5") or "") or ""
            if t:
                seen_txn_md5.add(str(t).strip().lower())

    # order rows by 'started' for deterministic within-window processing
    def _row_started_key(r):
        try:
            st = r.get("started")
            if st:
                if isinstance(st, datetime):
                    return st
                return datetime.fromisoformat(str(st))
        except Exception:
            pass
        return str(r.get("filename") or "")

    rows_sorted = sorted(rows, key=_row_started_key)

    processed_count = 0
    for r in rows_sorted:
        fn = str(r.get("filename") or "").strip().lower()
        p = (r.get("pdf_md5") or r.get("pdfMd5") or "") or ""
        p = str(p).strip().lower() if p else ""
        t = (r.get("txn_md5") or r.get("txnMd5") or "") or ""
        t = str(t).strip().lower() if t else ""

        is_duplicate = False
        if fn and fn in seen_filenames:
            is_duplicate = True
        elif p and p in seen_pdf_md5:
            is_duplicate = True
        elif t and t in seen_txn_md5:
            is_duplicate = True

        if not is_duplicate:
            processed_count += 1
            if fn:
                seen_filenames.add(fn)
            if p:
                seen_pdf_md5.add(p)
            if t:
                seen_txn_md5.add(t)
        else:
            if fn:
                seen_filenames.add(fn)
            if p:
                seen_pdf_md5.add(p)
            if t:
                seen_txn_md5.add(t)

    processed = processed_count

    # Rejected for charts: prefer DB-driven count (db_list_rejected), fallback to row scan
    rejected = 0
    if USE_DB:
        try:
            rejected_rows = []
            if date:
                rejected_rows = db_list_rejected(date)
            elif from_date and to_date:
                for d in _iter_dates(from_date, to_date):
                    try:
                        rejected_rows.extend(db_list_rejected(d))
                    except Exception:
                        pass
            else:
                rejected_rows = []
            rejected = sum(1 for r in rejected_rows if str(r.get("status", "")).upper() == "DUPLICATE")
        except Exception as e:
            print(f"[WARN] db_list_rejected failed in /api/charts: {e}")
            rejected = sum(1 for r in rows if str(r.get("status", "")).upper() == "DUPLICATE")
    else:
        rejected = sum(1 for r in rows if str(r.get("status", "")).upper() == "DUPLICATE")

    # Build time histogram same as before
    import numpy as np
    counts, bin_edges = ([], [])
    durations = [float(r.get("durationSec", 0)) for r in rows if "durationSec" in r and r.get("durationSec") is not None]
    if durations:
        counts, bin_edges = np.histogram(durations, bins=5)

    # Return status_pie with 'success' mapping to processed so frontend that expects 'success'
    # still receives meaningful data (processed vs rejected)
    return {
        "status_pie": {"success": processed, "rejected": rejected},
        "time_hist": {
            "bin_edges": bin_edges.tolist() if len(bin_edges) else [],
            "counts": counts.tolist() if len(counts) else [],
        },
    }

@app.get("/api/files")
def get_files(
    date: Optional[str] = None,
    from_date: Optional[str] = Query(None, alias="from"),
    to_date: Optional[str] = Query(None, alias="to"),
    state: Optional[str] = None,
    user=Depends(require_user),
):
    date = _norm_date(date) if date else None
    from_date = _norm_date(from_date) if from_date else None
    to_date = _norm_date(to_date) if to_date else None

    return _load_rows(date, from_date, to_date, state)

@app.get("/api/file")
def get_file(date: str, filename: str, user=Depends(require_user)):
    date = _norm_date(date)
    folder = _date_folder(date)

    # Fields from DB when enabled; fallback to *_output.json
    fields: Dict[str, Any] = {}
    if USE_DB:
        fields = dict(db_get_fields(date, filename) or {})
    else:
        detail_path = folder / f"{Path(filename).stem}_output.json"
        if detail_path.exists():
            fields = _read_json(detail_path, {}) or {}

    images = [str(p.as_posix()) for p in sorted(folder.glob(f"{Path(filename).stem}_page_*.png"))]

    # Fetch suspect info: prefer DB ProcessData row (IsSuspectData and SuspectExtractData columns)
    suspect = {}
    suspect_values = {}
    if USE_DB:
        try:
            # Use the parsed helper from db_repo which returns suspect_map/suspect_values and flags
            flags = get_processdata_flags_by_date_filename(date, filename) or {}
            suspect = flags.get("suspect_map", {}) or {}
            suspect_values = flags.get("suspect_values", {}) or {}
        except Exception as e:
            print(f"[WARN] Could not fetch ProcessData flags for suspect info: {e}")
            suspect = {}
            suspect_values = {}
    else:
        # file-mode fallback
        sfile = folder / f"{Path(filename).stem}_suspects.json"
        if sfile.exists():
            try:
                data = _read_json(sfile, {})
                suspect = data.get("IsSuspectData") or {}
                suspect_values = data.get("SuspectExtractData") or {}
            except Exception:
                suspect = {}
                suspect_values = {}

    return {"fields": fields, "images": images, "suspect": suspect or {}, "suspect_values": suspect_values or {}}

@app.post("/api/file/update")
def update_file_fields(payload: Dict[str, Any] = Body(...), user=Depends(require_user)):
    """
    Body: { "date": "...", "filename": "...", "fields": {k:v, ...}, "changes": {k:v,...}, "qc": true? }
    - Accepts either legacy 'fields' (full overwrite) or 'changes' (partial update).
    - Merges changes into existing ExtractedData and persists.
    - After persisting, ALWAYS re-run lightweight suspect heuristics on the merged canonical fields,
      persist the updated suspect metadata into ProcessData (or file-mode suspects JSON), and return
      the updated suspect status to the caller.
    - If this save is identified as a QC correction (qc: true OR updated keys intersect suspect keys
      OR user-entered values are deemed suspect by heuristics), update ProcessData.QCCorrectedData as well.
    """
    date_in = payload.get("date")
    filename = payload.get("filename")
    # prefer 'changes' (partial) but accept legacy 'fields' (full)
    changes = payload.get("changes")
    fields_full = payload.get("fields")
    qc_flag = bool(payload.get("qc", False))

    if not date_in or not filename or (not isinstance(changes, dict) and not isinstance(fields_full, dict)):
        # allow case where changes provided as dict OR fields provided
        if not (isinstance(changes, dict) or isinstance(fields_full, dict)):
            raise HTTPException(status_code=400, detail="Missing or invalid date/filename/fields or changes")

    # Normalize inputs
    date_iso = _norm_date(date_in)
    fn = filename if filename.lower().endswith(".pdf") else f"{filename}.pdf"
    stem = Path(fn).stem
    folder = _date_folder(date_iso)

    # Determine update payload: if 'changes' provided use that (partial), else use full fields
    if isinstance(changes, dict) and len(changes) > 0:
        to_apply = {k: v for k, v in changes.items()}
        partial_update = True
    else:
        to_apply = {k: v for k, v in (fields_full or {}).items()}
        partial_update = False

    # 1) Load existing extracted data (DB or file-mode), merge with to_apply
    merged_fields: Dict[str, Any] = {}
    if USE_DB:
        try:
            existing = db_get_fields(date_iso, fn) or {}
            merged_fields.update(existing)
        except Exception:
            merged_fields = {}
    else:
        out = folder / f"{stem}_output.json"
        if out.exists():
            try:
                existing = _read_json(out, {}) or {}
                merged_fields.update(existing)
            except Exception:
                merged_fields = {}

    # Merge (apply partial changes on top)
    merged_fields.update(to_apply)

    # Recompute totals so saved JSON contains recomputed values
    try:
        recompute_totals(merged_fields)
    except Exception:
        # swallow errors, continue to persist merged_fields anyway
        pass

    # --- Special-case: preserve comma thousands formatting for Taxable Selling Price / Purchase Price ---
    TAXABLE_KEY = "Taxable Selling Price / Purchase Price"
    try:
        if TAXABLE_KEY in merged_fields and merged_fields[TAXABLE_KEY] not in (None, ""):
            val_float = _to_float_safe(merged_fields[TAXABLE_KEY])
            merged_fields[TAXABLE_KEY] = _format_with_commas(val_float)
    except Exception:
        # Do not block save if formatting fails; keep whatever value is present
        pass

    # 2) Persist merged ExtractedData (DB or file-mode)
    if USE_DB:
        try:
            # db_upsert_output expects the full fields dict
            db_upsert_output(date_iso, fn, merged_fields)
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"DB update failed: {e}")
    else:
        out = folder / f"{stem}_output.json"
        try:
            out.parent.mkdir(parents=True, exist_ok=True)
            with out.open("w", encoding="utf-8") as f:
                json.dump(merged_fields, f, indent=2, ensure_ascii=False)
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"Local JSON update failed: {e}")

    # 2.5) Update Excel convenience copy (best-effort)
    try:
        state = Path(fn).stem.split("-", 1)[0].upper()
        mmddyyyy = datetime.fromisoformat(date_iso).strftime("%m%d%Y")
        xlsx_name = f"ARDEM_VITU_{state}_{mmddyyyy}_Output.xlsx"
        target = _date_folder(date_iso) / xlsx_name
        _update_excel_in_place(target, fn, merged_fields)
    except Exception:
        pass

    # 3) Determine existing suspect keys (from ProcessData or suspects file)
    existing_suspect_keys = set()
    existing_suspect_map = {}
    existing_suspect_values = {}
    existing_qc_corrected = {}

    if USE_DB:
        try:
            # Prefer the parsed flags helper which returns decoded JSON structures
            flags = get_processdata_flags_by_date_filename(date_iso, fn) or {}
            existing_suspect_map = flags.get("suspect_map") or {}
            existing_suspect_values = flags.get("suspect_values") or {}
            existing_qc_corrected = flags.get("qc_corrected") or {}
            existing_suspect_keys = set(existing_suspect_map.keys()) if isinstance(existing_suspect_map, dict) else set()
        except Exception as e:
            print(f"[WARN] Could not fetch ProcessData row when deciding QC: {e}")
            existing_suspect_map = {}
            existing_suspect_values = {}
            existing_qc_corrected = {}
            existing_suspect_keys = set()
    else:
        sfile = folder / f"{stem}_suspects.json"
        if sfile.exists():
            try:
                prev = _read_json(sfile, {}) or {}
                existing_suspect_map = prev.get("IsSuspectData") or {}
                existing_suspect_values = prev.get("SuspectExtractData") or {}
                existing_qc_corrected = prev.get("QCCorrectedData") or {}
                existing_suspect_keys = set(existing_suspect_map.keys()) if isinstance(existing_suspect_map, dict) else set()
            except Exception:
                existing_suspect_map = {}
                existing_suspect_values = {}
                existing_qc_corrected = {}
                existing_suspect_keys = set()

    # 4) Decide whether this update is a QC action:
    is_qc = qc_flag
    if not is_qc:
        # if any of the changed keys intersect with existing suspect keys, treat as QC
        if existing_suspect_keys and any(k in existing_suspect_keys for k in to_apply.keys()):
            is_qc = True

    # If user supplied values that our heuristic considers suspect, treat as QC and make sure we record suspect markers.
    newly_suspect_keys = set()
    for k, v in to_apply.items():
        # Skip structural keys not part of canonical fields
        try:
            if _is_value_suspect(k, v):
                newly_suspect_keys.add(k)
        except Exception:
            # Any failure should default to not marking as suspect here
            pass
    if newly_suspect_keys:
        is_qc = True

    # 5) Build the canonical suspect map based on the merged_fields (always re-run heuristics)
    #    We create a simple reason code "format_mismatch" for flagged values; you can refine codes if desired.
    new_suspect_map: Dict[str, str] = {}
    new_suspect_values: Dict[str, Any] = {}
    try:
        for k, v in merged_fields.items():
            try:
                if _is_value_suspect(k, v):
                    # Provide a generic reason; keep existing reason if present
                    reason = existing_suspect_map.get(k) if k in existing_suspect_map else "format_mismatch"
                    new_suspect_map[k] = reason
                    new_suspect_values[k] = v
            except Exception:
                continue
    except Exception:
        new_suspect_map = {}
        new_suspect_values = {}

    # 6) Persist ProcessData / suspects JSON: ensure DB or file-mode has the latest suspect view and QC metadata.
    now_ts = datetime.now().isoformat(timespec="seconds")
    is_suspect_validated = "Yes" if not new_suspect_map else "No"
    is_qc_completed = "Yes" if not new_suspect_map else "No"

    if is_qc:
        # If flagged QC above, we recorded QCCorrectedData in the earlier QC branch handling.
        # To keep behavior consistent, we will set QCCorrectedData to merged_fields (overwrite) when QC.
        qc_payload = merged_fields.copy()
    else:
        qc_payload = existing_qc_corrected or None

    if USE_DB:
        # Try to find GUID
        guid = None
        try:
            proc = proc if 'proc' in locals() and proc else get_processdata_by_date_filename(date_iso, fn)
        except Exception:
            proc = None
        try:
            if proc and proc.get("GUID"):
                guid = proc.get("GUID")
            else:
                guid = get_guid_by_date_filename(date_iso, fn)
        except Exception:
            guid = None

        # Build proc_row for upsert_processdata
        proc_row = {
            "GUID": guid,
            "FileName": fn,
            "State_Name": Path(fn).stem.split("-", 1)[0].upper() if "-" in Path(fn).stem else Path(fn).stem[:2].upper(),
            "DataExtractedDateTime": now_ts,
            "SuspectCheckingDateTime": now_ts,
            "IsDataExtracted": "Yes",
            "IsSuspectValidated": is_suspect_validated,
            "IsQCCompleted": is_qc_completed,
            "IsDelivered": False,
            "AWSS3Path": None,
            "ExtractedData": json.dumps(merged_fields, ensure_ascii=False),
            "IsSuspectData": json.dumps(new_suspect_map, ensure_ascii=False) if new_suspect_map else None,
            "SuspectExtractData": json.dumps(new_suspect_values, ensure_ascii=False) if new_suspect_values else None,
            # Store QCCorrectedData when QC; otherwise preserve existing if available
            "QCCorrectedData": json.dumps(qc_payload, ensure_ascii=False) if qc_payload else None,
            "QCCheckingDateTime": now_ts
        }
        try:
            upsert_processdata(proc_row)
        except Exception as e:
            print(f"[WARN] upsert_processdata failed in /api/file/update: {e}")
            # fallback: write file-mode suspects in date folder
            try:
                sfile = folder / f"{stem}_suspects.json"
                prev = _read_json(sfile, {}) or {}
                prev["ExtractedData"] = merged_fields
                prev["IsSuspectData"] = new_suspect_map if new_suspect_map else None
                prev["SuspectExtractData"] = new_suspect_values or {}
                prev["QCCorrectedData"] = qc_payload or prev.get("QCCorrectedData")
                prev["IsSuspectValidated"] = is_suspect_validated
                prev["IsQCCompleted"] = is_qc_completed
                prev["QCCheckingDateTime"] = now_ts
                sfile.parent.mkdir(parents=True, exist_ok=True)
                sfile.write_text(json.dumps(prev, ensure_ascii=False), encoding="utf-8")
            except Exception as ee:
                print(f"[ERROR] fallback suspects write also failed: {ee}")
    else:
        # file-mode: write suspects JSON reflecting latest state
        sfile = folder / f"{stem}_suspects.json"
        try:
            prev = _read_json(sfile, {}) or {}
            prev["ExtractedData"] = merged_fields
            prev["IsSuspectData"] = new_suspect_map if new_suspect_map else None
            prev["SuspectExtractData"] = new_suspect_values or {}
            prev["QCCorrectedData"] = qc_payload or prev.get("QCCorrectedData")
            prev["IsSuspectValidated"] = is_suspect_validated
            prev["IsQCCompleted"] = is_qc_completed
            prev["QCCheckingDateTime"] = now_ts
            sfile.parent.mkdir(parents=True, exist_ok=True)
            sfile.write_text(json.dumps(prev, ensure_ascii=False), encoding="utf-8")
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"Could not write suspects QC file: {e}")

    # Return structured response including new suspect status so frontend can immediately refresh view
    return {
        "ok": True,
        "qc": bool(is_qc),
        "has_suspect": bool(new_suspect_map),
        "suspect_map": new_suspect_map,
        "suspect_values": new_suspect_values
    }

@app.get("/api/audit")
def get_audit(
    from_date: str = Query(..., alias="from"),
    to_date: str = Query(..., alias="to"),
    state: Optional[str] = Query(None, description="Optional state code filter, e.g. PA"),
    user=Depends(require_user),
):
    from_date = _norm_date(from_date)
    to_date   = _norm_date(to_date)
    state_norm = state.strip().upper() if state else None

    def _normalize_audit_row(r):
        rr = dict(r)
        rr["processed"] = int(rr.get("processed") or 0)
        rr["failed"] = int(rr.get("failed") or 0)
        rr["rejected"] = int(rr.get("rejected") or rr.get("failed", 0))
        return rr

    def _collect_run_ids_from_files(file_rows):
        ids = set()
        for f in (file_rows or []):
            if not f:
                continue
            # prefer run_id/runId; also accept numeric types
            for key in ("run_id", "runId", "batch_id", "batch"):
                if key in f and f.get(key) is not None:
                    ids.add(str(f.get(key)).strip())
        return ids

    # ---------------- DB mode ----------------
    if USE_DB:
        try:
            rows = db_list_audit(from_date, to_date)
        except Exception as e:
            print(f"[ERROR] db_list_audit failed: {e}")
            return []

        audit_rows = [_normalize_audit_row(r) for r in (rows or [])]

        if state_norm:
            # Try to fetch files directly filtered by state if helper supports it
            files_for_state = []
            try:
                sig = inspect.signature(db_list_files)
                params = list(sig.parameters)
                if "state" in params or len(params) >= 3:
                    # db_list_files(from, to, state)
                    files_for_state = db_list_files(from_date, to_date, state_norm)
                else:
                    # helper doesn't accept state -> fetch all files and filter in-memory
                    files_all = db_list_files(from_date, to_date)
                    files_for_state = [
                        f for f in (files_all or [])
                        if (str(f.get("state") or "").strip().upper() == state_norm)
                        or (str(f.get("filename") or "").strip().upper().startswith(state_norm + "-"))
                    ]
            except Exception:
                # fallback to calling without state and then filtering
                try:
                    files_all = db_list_files(from_date, to_date)
                    files_for_state = [
                        f for f in (files_all or [])
                        if (str(f.get("state") or "").strip().upper() == state_norm)
                        or (str(f.get("filename") or "").strip().upper().startswith(state_norm + "-"))
                    ]
                except Exception as ex:
                    print(f"[WARN] db_list_files probing failed: {ex}")
                    files_for_state = []

            run_ids = _collect_run_ids_from_files(files_for_state)

            # Strict: if we cannot find any run_ids for the state, return empty list
            if not run_ids:
                return []

            filtered = [r for r in audit_rows if str(r.get("run_id") or "").strip() in run_ids]
            return filtered

        # No state requested
        return audit_rows

    # ---------------- File-mode ----------------
    rows = _read_jsonl(DATA_DIR / "audit_log.jsonl")
    audit_rows = [r for r in rows if from_date <= r.get("date", "") <= to_date]
    audit_rows = [_normalize_audit_row(r) for r in audit_rows]

    if state_norm:
        # _load_rows supports file-mode state filtering already
        files_for_state = _load_rows(None, from_date, to_date, state_norm)
        run_ids = _collect_run_ids_from_files(files_for_state)
        if not run_ids:
            return []
        filtered = [r for r in audit_rows if str(r.get("run_id") or "").strip() in run_ids]
        return filtered

    return audit_rows

# === New endpoint: return rejected rows for a given date (used by audit UI) ===
@app.get("/api/audit/rejected")
def get_rejected(
    date: str = Query(..., description="YYYY-MM-DD or MM/DD/YYYY"),
    run_id: Optional[str] = Query(None, description="Optional run_id to filter rejected rows to a specific run"),
    state: Optional[str] = Query(None, description="Optional state code to filter rejected rows"),
    user=Depends(require_user)
):
    date_iso = _norm_date(date)
    state_norm = state.strip().upper() if state else None

    def _matches_state_field(r, st):
        if not r or not st:
            return False
        # check state field or filename prefix
        if str(r.get("state") or "").strip().upper() == st:
            return True
        fn = str(r.get("filename") or r.get("file") or "").strip()
        if fn and "-" in fn and fn.split("-", 1)[0].strip().upper() == st:
            return True
        return False

    if USE_DB:
        try:
            sig = inspect.signature(db_list_rejected)
            params = list(sig.parameters)
            if len(params) >= 2:
                rows = db_list_rejected(date_iso, run_id or None)
            else:
                rows = db_list_rejected(date_iso)
        except Exception as e:
            print(f"[WARN] db_list_rejected failed: {e}")
            rows = []

        if state_norm:
            rows = [r for r in rows if _matches_state_field(r, state_norm)]
        return rows

    # File-mode fallback
    out = []
    folder = _date_folder(date_iso)
    snap = _read_json(folder / f"files_{date_iso}.json", [])
    if snap:
        for r in snap:
            status = str(r.get("status", "")).upper()
            if status in ("FAILED", "DUPLICATE") or r.get("rejection_reason") or r.get("rejected_for"):
                out.append(r)

    if state_norm:
        out = [r for r in out if _matches_state_field(r, state_norm)]

    return out

# ----------------------- Excel rebuild from DB + template -----------------------

def resolve_xlsx_template_for_state(state: Optional[str]) -> Path:
    """
    Resolve the XLSX template Path for a given state code.
    Priority:
      1) Environment variable: XLSX_TEMPLATE_PATH_{STATE} (e.g. XLSX_TEMPLATE_PATH_CT)
      2) File in TEMPLATE_DIR named with state suffix/prefix (several common names)
      3) Global XLSX template fallback (GLOBAL_XLSX_TEMPLATE_PATH)
    """
    s = (state or "").strip().upper()
    # 1) env var per state
    if s:
        env_var = f"XLSX_TEMPLATE_PATH_{s}"
        env_val = os.getenv(env_var)
        if env_val:
            p = Path(env_val)
            if p.exists():
                logger.info("Using XLSX template from env var %s -> %s", env_var, p)
                return p
            else:
                logger.warning("Env var %s set to %s but file not found", env_var, env_val)

    # 2) look for common filename patterns in TEMPLATE_DIR
    if s:
        candidates = [
            f"ARDEM_VITU_TEMPLATE_{s}.xlsx",
            f"{s}_ARDEM_VITU_TEMPLATE.xlsx",
            f"{s}_TEMPLATE.xlsx",
            f"{s}.xlsx",
            f"ARDEM_VITU_{s}_Template.xlsx",
        ]
        for name in candidates:
            p = TEMPLATE_DIR / name
            if p.exists():
                logger.info("Using state-specific template: %s", p)
                return p

    # 3) fallback to global
    if GLOBAL_XLSX_TEMPLATE_PATH.exists():
        logger.info("Using global XLSX template: %s", GLOBAL_XLSX_TEMPLATE_PATH)
        return GLOBAL_XLSX_TEMPLATE_PATH

    # nothing found — return expected global path (caller should check existence)
    logger.error("No XLSX template found for state=%s; falling back to expected path: %s", s or "NONE", GLOBAL_XLSX_TEMPLATE_PATH)
    return GLOBAL_XLSX_TEMPLATE_PATH

def _rebuild_excel_from_db_with_template(date_iso: str, state: str) -> Path:
    """
    Build ARDEM_VITU_{STATE}_{MMDDYYYY}_Output.xlsx in data/<date>/ using:
      - Triple-header template (rows 1..HEADER_ROWS) preserved
      - Data starts at DATA_START_ROW
      - Columns mapped by bottom header row (row HEADER_ROWS)
      - Accepts either a 'Transaction ID' column or a 'Filename' column in the template
      - Only fields whose keys match template bottom headers are written (case-insensitive)

    IMPORTANT: Prefer writing fields from ProcessData.QCCorrectedData if available,
    otherwise fall back to db_get_fields (ExtractedData).
    """
    if not USE_DB:
        raise HTTPException(status_code=400, detail="Rebuild from DB is only available in DB mode")

    # Resolve template path for this state
    template_path = resolve_xlsx_template_for_state(state)
    if not template_path.exists():
        raise HTTPException(status_code=500, detail=f"Template not found at {template_path}")

    # 1) Collect the files for this date/state
    files_rows = db_list_files(date_iso, date_iso, state)
    filenames = [r["filename"] for r in files_rows if r.get("filename")]

    # 2) Load template
    wb = load_workbook(str(template_path))
    ws = wb[XLSX_TEMPLATE_SHEET] if XLSX_TEMPLATE_SHEET and XLSX_TEMPLATE_SHEET in wb.sheetnames else wb.active

    # 3) Header maps
    bottom_headers = _collect_bottom_headers(ws)  # normalized header -> col idx
    filename_col = _find_header_col_by_synonyms(ws, FILENAME_SYNONYMS)
    txn_col = _find_header_col_by_synonyms(ws, TXN_SYNONYMS)

    if not filename_col and not txn_col:
        raise HTTPException(
            status_code=500,
            detail="Template must contain either a 'Transaction ID' (or synonym) column or a 'Filename' (or synonym) column in the header rows."
        )

    # 4) Clear existing data rows, keep the triple header exactly
    if ws.max_row >= DATA_START_ROW:
        ws.delete_rows(DATA_START_ROW, ws.max_row - DATA_START_ROW + 1)

    # 5) Write rows
    row_idx = DATA_START_ROW
    for fn in filenames:
        # Prefer ProcessData.QCCorrectedData when present (so Excel reflects QC corrections)
        fields = {}
        try:
            proc = get_processdata_by_date_filename(date_iso, fn)
            if proc:
                qc_raw = proc.get("QCCorrectedData")
                if qc_raw:
                    try:
                        fields = json.loads(qc_raw) if isinstance(qc_raw, str) else dict(qc_raw)
                    except Exception:
                        fields = {}
            # fallback to db_get_fields if qc not present or parse failed
            if not fields:
                fields = db_get_fields(date_iso, fn) or {}
        except Exception:
            # on any failure, fallback to db_get_fields
            try:
                fields = db_get_fields(date_iso, fn) or {}
            except Exception:
                fields = {}

        # Prefer writing Transaction ID if that column exists and value present; else fall back to filename
        txn_value = _get_field_by_synonyms(fields, TXN_SYNONYMS)

        if txn_col and (txn_value is not None) and (str(txn_value).strip() != ""):
            ws.cell(row=row_idx, column=txn_col).value = _coalesce_value(txn_value)
        elif filename_col:
            ws.cell(row=row_idx, column=filename_col).value = fn
        # else: no ID column to write; still fill other fields

        # write fields that match template bottom headers (case-insensitive)
        for k, v in fields.items():
            normk = _normalize_header(k)
            col = bottom_headers.get(normk)
            if not col:
                continue  # ignore fields not present in template bottom headers
            ws.cell(row=row_idx, column=col).value = _coalesce_value(v)

        row_idx += 1

    # 6) Save into date folder
    mmddyyyy = datetime.fromisoformat(date_iso).strftime("%m%d%Y")
    out_name = f"ARDEM_VITU_{state}_{mmddyyyy}_Output.xlsx"
    out_path = _date_folder(date_iso) / out_name
    out_path.parent.mkdir(parents=True, exist_ok=True)
    wb.save(str(out_path))
    return out_path

# ======================== Excel download endpoint ============================
@app.get("/api/audit/download")
def audit_download(
    date: str = Query(..., description="YYYY-MM-DD or MM/DD/YYYY"),
    state: Optional[str] = Query(None, description="State code in filename, e.g. PA"),
    run_id: Optional[str] = Query(None, description="Optional run_id to restrict to a specific run"),
    user=Depends(require_user),
):
    """
    Serve daily Excel following the naming:
      ARDEM_VITU_{STATE}_{MMDDYYYY}_Output.xlsx

    If run_id is provided → only include rows from that run with status='SUCCESS'.
    When building Excel from DB, prefer QCCorrectedData (full JSON) for fields if available.
    """
    date_iso = _norm_date(date)
    dt = datetime.fromisoformat(date_iso)
    mmddyyyy = dt.strftime("%m%d%Y")

    if USE_DB:
        # Resolve state if not provided
        st = (state or "").strip().upper() or None
        if not st:
            # Infer states from filenames in DB for that date
            rows = db_list_files(date_iso, date_iso, None)
            states = set()
            for r in rows:
                fn = str(r.get("filename", ""))
                if "-" in fn:
                    states.add(fn.split("-", 1)[0].upper())
            if len(states) == 0:
                raise HTTPException(status_code=404, detail=f"No records in DB for {date_iso}")
            if len(states) > 1:
                raise HTTPException(status_code=400, detail=f"Multiple states for {date_iso}: {', '.join(sorted(states))}. Specify ?state=")
            st = next(iter(states))

        # If run_id provided → filter DB rows to that run, only SUCCESS
        if run_id:
            rows = db_list_files(date_iso, date_iso, st)
            rows = [r for r in rows if str(r.get("run_id")) == run_id and str(r.get("status", "")).upper() == "SUCCESS"]

            if not rows:
                raise HTTPException(status_code=404, detail=f"No processed rows found for {date_iso} run_id={run_id}")

            # Resolve template path for this state and load
            template_path = resolve_xlsx_template_for_state(st)
            if not template_path.exists():
                raise HTTPException(status_code=500, detail=f"Template not found at {template_path}")

            wb = load_workbook(str(template_path))
            ws = wb[XLSX_TEMPLATE_SHEET] if XLSX_TEMPLATE_SHEET and XLSX_TEMPLATE_SHEET in wb.sheetnames else wb.active

            bottom_headers = _collect_bottom_headers(ws)  # normalized header -> col idx
            filename_col = _find_header_col_by_synonyms(ws, FILENAME_SYNONYMS)
            txn_col = _find_header_col_by_synonyms(ws, TXN_SYNONYMS)

            if not filename_col and not txn_col:
                raise HTTPException(
                    status_code=500,
                    detail="Template must contain either a 'Transaction ID' or 'Filename' column."
                )

            # clear data rows but preserve header rows
            if ws.max_row >= DATA_START_ROW:
                ws.delete_rows(DATA_START_ROW, ws.max_row - DATA_START_ROW + 1)

            row_idx = DATA_START_ROW
            for r in rows:
                fn = str(r.get("filename") or "")
                # Prefer ProcessData.QCCorrectedData
                fields = {}
                try:
                    proc = get_processdata_by_date_filename(date_iso, fn)
                    if proc:
                        qc_raw = proc.get("QCCorrectedData")
                        if qc_raw:
                            try:
                                fields = json.loads(qc_raw) if isinstance(qc_raw, str) else dict(qc_raw)
                            except Exception:
                                fields = {}
                    if not fields:
                        fields = db_get_fields(date_iso, fn) or {}
                except Exception:
                    try:
                        fields = db_get_fields(date_iso, fn) or {}
                    except Exception:
                        fields = {}

                txn_value = _get_field_by_synonyms(fields, TXN_SYNONYMS)
                if txn_col and (txn_value is not None) and (str(txn_value).strip() != ""):
                    ws.cell(row=row_idx, column=txn_col).value = _coalesce_value(txn_value)
                elif filename_col:
                    ws.cell(row=row_idx, column=filename_col).value = fn

                for k, v in fields.items():
                    normk = _normalize_header(k)
                    col = bottom_headers.get(normk)
                    if col:
                        ws.cell(row=row_idx, column=col).value = _coalesce_value(v)

                row_idx += 1

            out_name = f"ARDEM_VITU_{st}_{mmddyyyy}_{run_id}_Output.xlsx"
            out_path = _date_folder(date_iso) / out_name
            out_path.parent.mkdir(parents=True, exist_ok=True)
            wb.save(str(out_path))

            return FileResponse(
                str(out_path),
                media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
                filename=out_path.name,
            )

        # Default: fall back to original behavior (date-level rebuild)
        out_path = _rebuild_excel_from_db_with_template(date_iso, st)
        return FileResponse(
            str(out_path),
            media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
            filename=out_path.name,
        )

    # ---- Legacy File Mode ----
    # unchanged
    search_dirs = [
        _date_folder(date_iso),
        DATA_DIR,
        BASE_DIR,
    ]

    if state:
        code = "".join(ch for ch in state.upper() if ch.isalpha())
        fname = f"ARDEM_VITU_{code}_{mmddyyyy}_Output.xlsx"
        for d in search_dirs:
            p = d / fname
            if p.exists():
                return FileResponse(
                    str(p),
                    media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
                    filename=p.name,
                )
        raise HTTPException(status_code=404, detail=f"Excel not found for {date_iso} and state {code}")

    matches: List[Path] = []
    pattern = f"ARDEM_VITU_*_{mmddyyyy}_Output.xlsx"
    for d in search_dirs:
        if d.exists():
            matches.extend(d.glob(pattern))

    uniq = list(dict.fromkeys([m.resolve() for m in matches if m.is_file()]))
    if not uniq:
        raise HTTPException(status_code=404, detail=f"No Excel found for {date_iso}")

    if len(uniq) > 1:
        target_date_dir = _date_folder(date_iso).resolve()
        preferred = next((p for p in uniq if p.parent == target_date_dir), None)
        chosen = preferred or max(uniq, key=lambda x: x.stat().st_mtime)
        return FileResponse(
            str(chosen),
            media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
            filename=chosen.name,
        )

    p = uniq[0]
    return FileResponse(
        str(p),
        media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        filename=p.name,
    )

# ============================================================================

@app.get("/file/pdf")
def serve_pdf(date: str, filename: str, user=Depends(require_user)):
    # normalize inputs
    date = _norm_date(date)
    fn = filename if filename.lower().endswith(".pdf") else f"{filename}.pdf"
    # fix accidental brace if present
    if fn.endswith('}'):
        fn = fn[:-1]
    folder = _date_folder(date)
    path = folder / fn

    # 1) Local file → stream inline (no download)
    if path.exists():
        f = path.open("rb")
        return StreamingResponse(
            f,
            media_type="application/pdf",
            headers={"Content-Disposition": f'inline; filename="{fn}"'}
        )

    # Legacy pipeline fallback: older state workers persisted PDFs at the
    # repository-level data/<date>/ folder rather than Backend_db/data/<date>/.
    legacy_path = BASE_DIR.parent / "data" / date / fn
    if legacy_path.exists():
        f = legacy_path.open("rb")
        return StreamingResponse(
            f,
            media_type="application/pdf",
            headers={"Content-Disposition": f'inline; filename="{fn}"'}
        )

    # 2) S3 fallback via DB key → stream inline (no redirect/CORS/download)
    s3_key = None
    if USE_DB:
        try:
            s3_key = db_get_s3_key(date, fn)
        except Exception:
            s3_key = None

    if s3_key:
        bucket = os.getenv("BUCKET_NAME")
        if not bucket:
            raise HTTPException(status_code=500, detail="BUCKET_NAME not configured")
        try:
            obj = _s3_client().get_object(Bucket=bucket, Key=s3_key)
            body = obj["Body"]  # botocore.response.StreamingBody

            def _iter():
                for chunk in iter(lambda: body.read(1024 * 1024), b""):
                    yield chunk

            return StreamingResponse(
                _iter(),
                media_type="application/pdf",
                headers={"Content-Disposition": f'inline; filename="{fn}"'}
            )
        except Exception as e:
            raise HTTPException(status_code=502, detail=f"S3 fetch failed: {e}")

    # 3) Not found
    raise HTTPException(
        status_code=404,
        detail="PDF not found locally and no S3 key stored for this file/date"
    )



# ---------------------------------------------------------------------
# QC compatibility endpoints (frontend tries /api/qc/* — support those)
# ---------------------------------------------------------------------
from fastapi import APIRouter

# /api/qc/charts -> reuse existing get_charts implementation
@app.get("/api/qc/charts")
def qc_charts(
    date: Optional[str] = None,
    from_date: Optional[str] = Query(None, alias="from"),
    to_date: Optional[str] = Query(None, alias="to"),
    state: Optional[str] = None,
    user=Depends(require_user),
):
    """
    Compatibility wrapper so frontend can call /api/qc/charts.
    Delegates to existing get_charts implementation to preserve behavior.
    """
    # simply call the already-defined handler function
    return get_charts(date=date, from_date=from_date, to_date=to_date, state=state, user=user)


def _row_is_suspect_for_qc(r: dict) -> bool:
    """Mirror JS countSuspectInRow logic with conservative checks."""
    try:
        # Prefer explicit DB-friendly flags (multiple possible names)
        if r is None:
            return False
        # Several DB schemas may use different keys; treat truthy values as suspect
        for key in ("has_suspect", "is_suspect_data", "is_suspect", "HasSuspect", "IsSuspectData"):
            v = r.get(key)
            if v in (1, "1", True, "true", "TRUE", "Yes", "YES", "yes"):
                return True

        # suspect JSON blobs under common keys
        suspect_map = r.get("suspect_map") or r.get("IsSuspectData") or r.get("suspectMap")
        if isinstance(suspect_map, dict) and len(suspect_map):
            return True
        suspect_values = r.get("suspect_extract_data") or r.get("SuspectExtractData") or r.get("suspect_values")
        if isinstance(suspect_values, dict) and len(suspect_values):
            return True

        # Generic scan: any key name containing 'suspect' with non-empty non-numeric value is suspect
        for k, v in (r or {}).items():
            if "suspect" in str(k).lower():
                # dict means structured suspect info -> suspect
                if isinstance(v, dict) and len(v):
                    return True
                # non-empty string/other -> suspect
                if v and not isinstance(v, (int, float)):
                    if str(v).strip() != "":
                        return True
    except Exception:
        # Be conservative on failure and treat as not suspect
        pass
    return False


def _row_is_qc_completed_for_qc(r: dict) -> bool:
    """Mirror JS isQCCompletedInRow logic: true if QC complete flags or equivalent present."""
    try:
        if r is None:
            return False
        # Common keys
        for key in ("is_qc_completed", "IsQCCompleted", "isQCCompleted", "QCCompleted"):
            v = r.get(key)
            if v in (1, "1", True, "true", "TRUE", "Yes", "YES", "yes"):
                return True

        # older schema might use 'IsQCCompleted' as "Yes"/"No"
        v = r.get("IsQCCompleted")
        if isinstance(v, str):
            if v.strip().upper() == "YES":
                return True

        # fallback: processed + no suspect -> treat as QC-completed
        status = str(r.get("status") or "").upper()
        if status == "SUCCESS":
            # if there is no suspect evidence we consider it cleared/complete
            if not _row_is_suspect_for_qc(r):
                return True
    except Exception:
        pass
    return False


@app.get("/api/qc/summary")
def qc_summary(
    date: Optional[str] = None,
    from_date: Optional[str] = Query(None, alias="from"),
    to_date: Optional[str] = Query(None, alias="to"),
    state: Optional[str] = None,
    user=Depends(require_user),
):
    """
    Lightweight QC summary endpoint for frontend:
      - received : total rows in window
      - suspect  : files flagged suspect by DB flags or heuristics
      - cleared  : processed and QC completed OR processed with no suspects
      - rejected : failed / duplicate / rejected_for / rejection_reason
    Mirrors the logic used in the QC page JS so counts stay consistent.
    """
    # normalize dates like other endpoints
    date = _norm_date(date) if date else None
    from_date = _norm_date(from_date) if from_date else None
    to_date = _norm_date(to_date) if to_date else None

    rows = _load_rows(date, from_date, to_date, state)
    received = len(rows)

    suspect_count = 0
    cleared_count = 0
    rejected_count = 0

    for r in rows:
        status = str(r.get("status") or "").upper()
        # consider rejected if explicit failure/duplicate or rejection metadata present
        is_rejected = status in ("FAILED", "DUPLICATE") or bool(r.get("rejection_reason")) or bool(r.get("rejected_for"))
        if is_rejected:
            rejected_count += 1
            continue

        is_sus = _row_is_suspect_for_qc(r)
        if is_sus:
            suspect_count += 1

        is_processed = (status == "SUCCESS")
        qc_done = _row_is_qc_completed_for_qc(r)
        # cleared = processed AND (qc done OR not suspect)
        if is_processed and (qc_done or not is_sus):
            cleared_count += 1

    return {
        "received": received,
        "suspect": suspect_count,
        "cleared": cleared_count,
        "rejected": rejected_count
    }


import os
import pyodbc
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List

# ---------- State model ----------
class StateIn(BaseModel):
    state_name: str


# ---------- Build MSSQL connection string safely ----------
MSSQL_DRIVER = os.getenv("MSSQL_DRIVER", "ODBC Driver 17 for SQL Server")
MSSQL_SERVER = os.getenv("MSSQL_SERVER", "WORKSPACE-95\\SQLEXPRESS")  # Note: double backslash
MSSQL_DATABASE = os.getenv("MSSQL_DATABASE", "VITU")
MSSQL_USER = os.getenv("MSSQL_USER", "vitu_app")
MSSQL_PASSWORD = os.getenv("MSSQL_PASSWORD", "JrN3$9wz-LarkTide_482")

conn_str = (
    f"DRIVER={{{MSSQL_DRIVER}}};"
    f"SERVER={MSSQL_SERVER};"
    f"DATABASE={MSSQL_DATABASE};"
    f"UID={MSSQL_USER};"
    f"PWD={MSSQL_PASSWORD};"
    f"TrustServerCertificate=yes"
)


# ---------- POST /api/states ----------
@app.post("/api/states", status_code=201)
def create_state(state: StateIn):
    """Create a new state (two-letter initials) in dbo.states"""
    name = (state.state_name or "").strip().upper()

    if not name.isalpha() or len(name) != 2:
        raise HTTPException(
            status_code=400,
            detail="state_name must be exactly two letters (e.g., 'KA')"
        )

    try:
        with pyodbc.connect(conn_str, autocommit=True) as conn:
            cur = conn.cursor()
            try:
                cur.execute("INSERT INTO dbo.states (state_name) VALUES (?)", name)
            except pyodbc.IntegrityError:
                raise HTTPException(status_code=409, detail="State already exists")
            except Exception as ex:
                msg = str(ex).lower()
                if 'unique' in msg or 'duplicate' in msg:
                    raise HTTPException(status_code=409, detail="State already exists")
                raise
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"DB error: {str(e)}")

    return {"state_name": name}


# ---------- GET /api/states ----------
@app.get("/api/states", response_model=List[dict])
def list_states():
    """Return all states stored in dbo.states"""
    try:
        with pyodbc.connect(conn_str, autocommit=True) as conn:
            cur = conn.cursor()
            cur.execute("SELECT id, state_name FROM dbo.states ORDER BY state_name;")
            rows = cur.fetchall()
            return [{"id": int(r[0]), "state_name": (r[1] or '').upper()} for r in rows]
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"DB error: {str(e)}")

from fastapi import UploadFile, File
import shutil

@app.post("/api/upload")
async def upload_file(
    file: UploadFile = File(...),
    state: Optional[str] = None,           # optional form field sent by client
    user=Depends(require_user),
):
    """
    Accepts a single file (multipart/form-data). Optional form field 'state' allowed.
    Saves to DATA_DIR/<YYYY-MM-DD>/ (or DATA_DIR/<YYYY-MM-DD>/<STATE>/).
    Returns JSON with saved path.
    """
    # ensure we have a filename
    filename = file.filename or "uploaded_file"
    # sanitize a bit
    filename = filename.replace("/", "_").replace("\\", "_")

    # choose date folder (today) — you could accept a 'date' form field if you prefer
    date_iso = datetime.utcnow().date().isoformat()
    # if client provided a state, put in subfolder
    if state:
        state_token = str(state).strip().upper()
        target_dir = _date_folder(date_iso) / f"{state_token}"
    else:
        target_dir = _date_folder(date_iso)

    try:
        target_dir.mkdir(parents=True, exist_ok=True)
        out_path = target_dir / filename
        # write uploaded file to disk (streaming)
        with out_path.open("wb") as f:
            while True:
                chunk = await file.read(1024 * 1024)
                if not chunk:
                    break
                f.write(chunk)
        # Optionally, you may want to create an initial _output.json or metadata entry here.
        return {"ok": True, "filename": filename, "saved_to": str(out_path.relative_to(BASE_DIR))}
    except Exception as e:
        logger.exception("Upload failed")
        raise HTTPException(status_code=500, detail=f"Upload failed: {e}")
