#!/usr/bin/env python3
"""
backfill_ct_annual_report_due.py — Populate entity_cache.annual_report_due_date
for Connecticut, so the /tools/corporation-check tool can show the
annual-report-overdue signal (the CT campaign's hook).

entity_cache has CT status but no annual-report clock. This pulls
annual_report_due_date from the live CT Socrata registry (data.ct.gov/n7gp-d28j),
keyed by accountnumber == entity_cache.entity_number, and writes it back.

Adds the column if missing. Idempotent: re-runnable.

    python3 scripts/backfill_ct_annual_report_due.py            # full CT backfill
    python3 scripts/backfill_ct_annual_report_due.py --limit 5000
"""
import argparse
import os
import sys
import time

import psycopg2
import psycopg2.extras
import requests

DB_URL = os.getenv("DATABASE_URL", "postgresql://pw@localhost/performancewest")
SOCRATA = "https://data.ct.gov/resource/n7gp-d28j.json"
PAGE = 50000


def fetch_ct_due_dates(limit=None):
    """Yield (accountnumber, annual_report_due_date) for CT entities that have both."""
    off = 0
    got = 0
    s = requests.Session()
    while True:
        params = {
            "$select": "accountnumber,annual_report_due_date",
            "$where": "annual_report_due_date IS NOT NULL AND accountnumber IS NOT NULL",
            "$limit": PAGE,
            "$offset": off,
            "$order": "accountnumber",
        }
        r = s.get(SOCRATA, params=params, timeout=60)
        r.raise_for_status()
        rows = r.json()
        if not rows:
            break
        for row in rows:
            acct = (row.get("accountnumber") or "").strip()
            due = (row.get("annual_report_due_date") or "")[:10]  # YYYY-MM-DD
            if acct and due:
                yield acct, due
                got += 1
                if limit and got >= limit:
                    return
        off += len(rows)
        print(f"  fetched {off:,} CT rows...", file=sys.stderr)
        if len(rows) < PAGE:
            break


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--limit", type=int, default=None)
    args = ap.parse_args()

    conn = psycopg2.connect(DB_URL)
    conn.autocommit = False
    cur = conn.cursor()

    # 1. add column if missing
    cur.execute("""
        ALTER TABLE entity_cache
        ADD COLUMN IF NOT EXISTS annual_report_due_date date
    """)
    conn.commit()
    print("column ensured: entity_cache.annual_report_due_date")

    # 2. stage the (accountnumber, due) pairs, then bulk-update by join
    cur.execute("DROP TABLE IF EXISTS _ct_due_stage")
    cur.execute("CREATE TEMP TABLE _ct_due_stage(entity_number text, due date)")

    batch = []
    total = 0
    t0 = time.time()
    for acct, due in fetch_ct_due_dates(args.limit):
        batch.append((acct, due))
        if len(batch) >= 10000:
            psycopg2.extras.execute_values(
                cur, "INSERT INTO _ct_due_stage(entity_number, due) VALUES %s", batch)
            total += len(batch); batch = []
    if batch:
        psycopg2.extras.execute_values(
            cur, "INSERT INTO _ct_due_stage(entity_number, due) VALUES %s", batch)
        total += len(batch)
    conn.commit()
    print(f"staged {total:,} CT due-date rows in {time.time()-t0:.0f}s")

    # 3. index + update entity_cache
    cur.execute("CREATE INDEX ON _ct_due_stage(entity_number)")
    cur.execute("""
        UPDATE entity_cache e
        SET annual_report_due_date = s.due
        FROM _ct_due_stage s
        WHERE e.state = 'CT' AND e.entity_number = s.entity_number
          AND (e.annual_report_due_date IS DISTINCT FROM s.due)
    """)
    updated = cur.rowcount
    conn.commit()
    print(f"updated {updated:,} CT entity_cache rows with annual_report_due_date")

    # 4. verify
    cur.execute("""
        SELECT count(*) FILTER (WHERE annual_report_due_date IS NOT NULL),
               count(*) FILTER (WHERE annual_report_due_date < now()),
               count(*)
        FROM entity_cache WHERE state = 'CT'
    """)
    have, overdue, tot = cur.fetchone()
    print(f"CT entity_cache: {have:,} with due date, {overdue:,} overdue, {tot:,} total")
    cur.close(); conn.close()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())