Workflow-App/Workflow/scripts/SARS retun EMP201 push to PB.py

"""
SARS retun EMP201 push to PB.py

Reads the JSON output of `SARS Retun EMP201 single or all pty.py`
(the SARS capture script) and upserts one row per (company, period)
into the PocketBase collection `ysterbaard_SARS_EMP201_PAYE`.

The script does NOT touch the SARS website — that's the capture
script's job. This script is the "save what we already captured"
step in the workflow.

CLI:
    --input <path>            path to a specific JSON file
    --input latest            (default) newest *.json in sars_captures/
    --list                    list available JSON files, don't write
    --dry-run                 print what would be written, no writes
    --status-filter all|captured|not_filed
                              which periods to include (default: all)

Behaviour:
    * For each company in the JSON:
        * For each period in results[]  → POST or PATCH (Return_Filed=true)
        * For each period in not_filed[] → POST or PATCH (Return_Filed=false)
        * Periods in errors[] are SKIPPED (capture didn't succeed)
    * Upsert key: company_period_key = f"{pb_id}:{period}"
        → UNIQUE index in PB → at most one row per (company, period)
    * When a row already exists, PATCH in place (option A).
    * We do NOT write to: status (workflow app's job), documents,
      interest_d, penalties_d, date_paid, due_date,
      sars_reference, notes, is_active.
    * We DO write date_submitted (from the work page's Status Date cell,
      captured by the SARS capture script). Re-runs overwrite with the
      SARS-side value (it should be stable across runs — it's the
      original submission date, not the run date).

PB env:
    PB_URL, PB_TOKEN are read from `Workflow/.env.local`.
"""

import argparse
import json
import os
import re
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import List, Optional, Tuple

import requests


# ── Paths & env ──────────────────────────────────────────────────

SCRIPT_DIR = Path(__file__).resolve().parent
PROJECT_ROOT = SCRIPT_DIR.parent
ENV_PATH = PROJECT_ROOT / ".env.local"
CAPTURES_DIR = Path(r"C:\Users\SJ\Downloads\sars_captures")
PB_COLLECTION = "ysterbaard_SARS_EMP201_PAYE"


def load_env() -> None:
    """Load .env.local into os.environ (best-effort, no hard dep on dotenv)."""
    if not ENV_PATH.exists():
        return
    try:
        from dotenv import load_dotenv
        load_dotenv(ENV_PATH)
        return
    except ImportError:
        pass
    for line in ENV_PATH.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        k, v = line.split("=", 1)
        os.environ.setdefault(k.strip(), v.strip().strip('"').strip("'"))


# ── PB HTTP helpers ──────────────────────────────────────────────

def pb_headers() -> dict:
    token = os.environ.get("PB_TOKEN", "")
    if not token:
        raise RuntimeError(
            f"PB_TOKEN missing. Add it to {ENV_PATH}."
        )
    return {"Authorization": f"Bearer {token}"}


def pb_base() -> str:
    return os.environ.get("PB_URL", "https://pb.agent.apphub.co.za").rstrip("/")


def pb_find_by_key(key: str) -> Optional[dict]:
    """Return the existing record for `company_period_key`, or None."""
    # PB filter syntax: company_period_key="..."
    # The value needs to be URL-encoded but requests does that for params.
    r = requests.get(
        f"{pb_base()}/api/collections/{PB_COLLECTION}/records",
        headers=pb_headers(),
        params={"filter": f'company_period_key="{key}"', "perPage": 1},
        timeout=30,
    )
    r.raise_for_status()
    items = r.json().get("items", [])
    return items[0] if items else None


def pb_post(payload: dict) -> dict:
    r = requests.post(
        f"{pb_base()}/api/collections/{PB_COLLECTION}/records",
        headers={**pb_headers(), "Content-Type": "application/json"},
        json=payload,
        timeout=30,
    )
    if r.status_code >= 400:
        raise RuntimeError(
            f"PB POST {r.status_code}: {r.text[:500]}"
        )
    r.raise_for_status()
    return r.json()


def pb_patch(record_id: str, payload: dict) -> dict:
    r = requests.patch(
        f"{pb_base()}/api/collections/{PB_COLLECTION}/records/{record_id}",
        headers={**pb_headers(), "Content-Type": "application/json"},
        json=payload,
        timeout=30,
    )
    if r.status_code >= 400:
        raise RuntimeError(
            f"PB PATCH {r.status_code}: {r.text[:500]}"
        )
    r.raise_for_status()
    return r.json()


# ── Field mapping ───────────────────────────────────────────────

def _to_float(s) -> Optional[float]:
    """Parse a number string from the JSON. Returns None for empty/None/non-numeric."""
    if s is None:
        return None
    s = str(s).strip()
    if not s:
        return None
    try:
        return float(s)
    except ValueError:
        return None


def _pb_date(dt: datetime) -> str:
    """Format a datetime in PocketBase's expected `date` field format:
    "YYYY-MM-DD HH:MM:SS.sssZ" (space separator, millis, trailing Z).
    PB silently drops ISO "T"-separated strings on date fields, so this
    helper is required."""
    return dt.strftime("%Y-%m-%d %H:%M:%S.") + f"{dt.microsecond // 1000:03d}Z"


def build_payload(period: str, pb_id_company: str, fields: dict,
                  return_filed: bool) -> dict:
    """Build the PB row payload. The fields dict is the row from
    `companies[].results[]` in the JSON; for not_filed rows, pass an
    empty dict (amounts will be None).

    NOTE: `status` is a required PB select. We set a sensible default
    on first write (submitted for captured, draft for not_filed) and
    the workflow app takes it from there (submitted → paid once the
    EFT is done). On subsequent PATCHes we deliberately do NOT touch
    status — if the user already moved it to "paid", our re-run won't
    drag it back to "submitted".
    """
    payload = {
        "company_period_key": f"{pb_id_company}:{period}",
        "company":           pb_id_company,
        "period":            period,
        "Return_Filed":      return_filed,
        # PB's `date` field expects "YYYY-MM-DD HH:MM:SS.sssZ" (with the
        # space separator and trailing Z, NOT the ISO "T" form we
        # produced by default with .isoformat()).
        "last_updated":      _pb_date(datetime.now()),
        "paye_amount":       _to_float(fields.get("paye_payable"))  if return_filed else None,
        "sdl_amount":        _to_float(fields.get("sdl_payable"))   if return_filed else None,
        "uif_amount":        _to_float(fields.get("uif_payable"))   if return_filed else None,
        "total_payable":     _to_float(fields.get("total_payable")) if return_filed else None,
        "payment_reference": (fields.get("payment_ref_no") or "").strip() or None if return_filed else None,
        "date_submitted":    fields.get("date_submitted") or None if return_filed else None,
        "status":            "submitted" if return_filed else "draft",
    }
    return payload


def build_patch_payload(fields: dict, return_filed: bool) -> dict:
    """Fields we update on PATCH. Deliberately does NOT include
    `status` — we don't want to overwrite a "paid" the user manually
    set. `date_submitted` IS included — it's the SARS-side submission
    date, stable across re-runs, so overwriting is safe."""
    return {
        "Return_Filed":      return_filed,
        "last_updated":      _pb_date(datetime.now()),
        "paye_amount":       _to_float(fields.get("paye_payable"))  if return_filed else None,
        "sdl_amount":        _to_float(fields.get("sdl_payable"))   if return_filed else None,
        "uif_amount":        _to_float(fields.get("uif_payable"))   if return_filed else None,
        "total_payable":     _to_float(fields.get("total_payable")) if return_filed else None,
        "payment_reference": (fields.get("payment_ref_no") or "").strip() or None if return_filed else None,
        "date_submitted":    fields.get("date_submitted") or None if return_filed else None,
    }


# ── Capture file resolution ─────────────────────────────────────

def find_latest_capture() -> Path:
    if not CAPTURES_DIR.exists():
        raise SystemExit(
            f"Captures dir not found: {CAPTURES_DIR}. "
            "Run the SARS capture script first."
        )
    candidates = sorted(CAPTURES_DIR.glob("*.json"),
                         key=lambda p: p.stat().st_mtime,
                         reverse=True)
    if not candidates:
        raise SystemExit(
            f"No *.json in {CAPTURES_DIR}. Run the capture script first."
        )
    return candidates[0]


def list_captures() -> List[Tuple[Path, float]]:
    if not CAPTURES_DIR.exists():
        return []
    out = []
    for p in sorted(CAPTURES_DIR.glob("*.json"),
                    key=lambda p: p.stat().st_mtime,
                    reverse=True):
        out.append((p, p.stat().st_mtime))
    return out


# ── Upsert loop ─────────────────────────────────────────────────

def upsert_period(period: str, pb_id: str, fields: dict,
                  return_filed: bool, dry_run: bool,
                  stats: dict) -> None:
    """One (company, period) row. Updates stats. Never raises on
    a single-row failure (logged and counted in `errored`)."""
    if not pb_id:
        msg = f"{period}: missing company pb_id, skipping"
        print(f"  {msg}", flush=True)
        stats["skipped"] += 1
        stats["messages"].append(msg)
        return

    payload = build_payload(period, pb_id, fields, return_filed)
    key = payload["company_period_key"]

    if dry_run:
        existing = "(would PATCH)" if pb_find_by_key(key) else "(would POST)"
        print(f"  {period}: {return_filed=!s:5}  {existing}  {key}",
              flush=True)
        stats["dry_run"] += 1
        return

    try:
        existing = pb_find_by_key(key)
        if existing:
            # PATCH: only update the data fields, leave `status` alone
            # so we don't drag "paid" back to "submitted".
            patch_payload = build_patch_payload(fields, return_filed)
            pb_patch(existing["id"], patch_payload)
            print(f"  {period}: PATCHED  {key}", flush=True)
            stats["patched"] += 1
        else:
            pb_post(payload)
            print(f"  {period}: POSTED   {key}", flush=True)
            stats["posted"] += 1
    except Exception as e:
        msg = f"{period}: {type(e).__name__}: {e}"
        print(f"  {msg}", flush=True)
        stats["errored"] += 1
        stats["messages"].append(msg)


# ── Main ─────────────────────────────────────────────────────────

def main():
    ap = argparse.ArgumentParser(
        description="Push SARS capture JSON → PocketBase ysterbaard_SARS_EMP201_PAYE."
    )
    ap.add_argument("--input", default="latest",
                    help='Path to capture JSON, or "latest" (default: latest)')
    ap.add_argument("--list", action="store_true",
                    help="List available capture JSON files and exit")
    ap.add_argument("--dry-run", action="store_true",
                    help="Print what would be written, no PB writes")
    ap.add_argument("--status-filter", default="all",
                    choices=["all", "captured", "not_filed"],
                    help="Which period states to include (default: all)")
    args = ap.parse_args()

    load_env()

    if args.list:
        print("Available capture files (newest first):", flush=True)
        for p, mtime in list_captures():
            print(f"  {datetime.fromtimestamp(mtime).isoformat(timespec='seconds')}  "
                  f"{p.name}  ({p.stat().st_size} bytes)", flush=True)
        return

    # Resolve input JSON
    if args.input == "latest":
        in_path = find_latest_capture()
    elif os.path.isfile(args.input):
        in_path = Path(args.input)
    else:
        # Maybe they passed a filename relative to captures dir
        candidate = CAPTURES_DIR / args.input
        if candidate.is_file():
            in_path = candidate
        else:
            raise SystemExit(
                f"Input not found: {args.input!r} (also tried {candidate})"
            )

    print(f"Reading:        {in_path}", flush=True)
    print(f"Dry run:        {args.dry_run}", flush=True)
    print(f"Status filter:  {args.status_filter}", flush=True)

    with open(in_path, "r", encoding="utf-8") as f:
        capture = json.load(f)

    companies = capture.get("companies", [])
    if not companies:
        raise SystemExit("No 'companies' array in the capture JSON.")

    # Stats
    stats = {
        "posted": 0, "patched": 0,
        "errored": 0, "skipped": 0, "dry_run": 0,
        "messages": [],
    }

    started = time.time()
    for ci, co in enumerate(companies, 1):
        pb_id = co.get("pb_id") or ""
        name = co.get("name") or "?"
        print(f"\n[{ci}/{len(companies)}] {name}  pb_id={pb_id}", flush=True)

        # captured periods
        if args.status_filter in ("all", "captured"):
            for r in co.get("results", []):
                period = r.get("period")
                if not period:
                    continue
                upsert_period(period, pb_id, r, return_filed=True,
                              dry_run=args.dry_run, stats=stats)

        # not_filed periods
        if args.status_filter in ("all", "not_filed"):
            for period in co.get("not_filed", []):
                if not period:
                    continue
                upsert_period(period, pb_id, {}, return_filed=False,
                              dry_run=args.dry_run, stats=stats)

    elapsed = time.time() - started

    print("\n" + "=" * 60, flush=True)
    if args.dry_run:
        print(f"DRY RUN — no writes performed.", flush=True)
        print(f"  would-write rows:  {stats['dry_run']}", flush=True)
    else:
        print(f"PB writes:", flush=True)
        print(f"  POSTED (new):     {stats['posted']}", flush=True)
        print(f"  PATCHED (updated): {stats['patched']}", flush=True)
    if stats["errored"]:
        print(f"  ERRORED:          {stats['errored']}", flush=True)
        for m in stats["messages"]:
            print(f"    {m}", flush=True)
    if stats["skipped"]:
        print(f"  SKIPPED:          {stats['skipped']}", flush=True)
    print(f"\nElapsed: {elapsed:.1f}s", flush=True)
    print(f"Source:  {in_path.name}", flush=True)
    print("=" * 60, flush=True)


if __name__ == "__main__":
    main()