#!/usr/bin/env python3
"""Turn written business requirements into SQL data quality rules, test the rules, and keep the evidence.

Each requirement below is one sentence a business owner would sign, a SQL condition that is TRUE for a
bad row, the share of bad rows the owner tolerates, and example rows that prove the SQL does what the
sentence says. The script (1) profiles every column of the data group by group (here: election cycle by
election cycle), (2) tests every rule against its own good and bad example rows, (3) runs every rule over
the data, and (4) writes one evidence row per rule and group: the requirement, a fingerprint of the exact
SQL, rows tested, rows failed, the tolerance, the verdict and the first failing row NUMBERS.

    python3 requirements_to_rules.py --data-dir <empty folder>
    python3 requirements_to_rules.py --data-dir <another empty folder> --out-dir planted/ \\
        --plant duplicate-rows,unknown-type,blank-recipient,broken-rule,unused-code
    python3 requirements_to_rules.py --csv your_file.csv --group-column your_year_column     # your own data

Needs Python 3 and the duckdb module (tested on Python 3.12, DuckDB 1.5.6). The demo data is two closed
cycles of the FEC "contributions from committees to candidates and other committees" bulk file (pas2). It prints
each file's byte size, Last-Modified header and the first 16 characters of its SHA-256, then counts, shares, row numbers and code values
(such as a transaction type), never a name or address. If you add a code list for a column that holds personal
data, a failed list check prints the unlisted values: list only columns that hold codes. Every file it writes goes to --out-dir (default: this script's own folder). It exits
1, and writes no evidence CSV, if one of its own checks fails.

The FEC data stands in for a company's own data. FEC contributor names and addresses may not be sold or
used for solicitation or commercial purposes: https://www.fec.gov/updates/sale-or-use-contributor-information/
"""
import argparse
import csv
import hashlib
import os
import pathlib
import shutil
import sys
import urllib.request
import zipfile

import duckdb

BASE = "https://www.fec.gov/files/bulk-downloads"
CYCLES = [2018, 2020]
HEADER_URL = f"{BASE}/data_dictionaries/pas2_header_file.csv"

# What each TRANSACTION_TP code in this file means (FEC: "Transaction type codes").
# ---- YOUR LOOKUPS: code lists your rules refer to. Use {} if your rules need none. ----
# name -> the column the codes belong to, whether every code must appear in the data, and the codes.
LOOKUPS = {
    "txn_types": {"column": "txn_type", "must_all_appear": True, "codes": {
        "24A": "Independent expenditure opposing election of candidate",
        "24C": "Coordinated party expenditure",
        "24E": "Independent expenditure advocating election of candidate",
        "24F": "Communication cost for candidate (only for Form 7 filer)",
        "24K": "Contribution made to nonaffiliated committee",
        "24N": "Communication cost against candidate (only for Form 7 filer)",
        "24P": "Contribution made to possible federal candidate including in-kind contributions",
        "24R": "Election recount disbursement",
        "24Z": "In-kind contribution made to registered filer"}},
    # USPS state, territory and military codes. Not every code has to appear in a file.
    "us_codes": {"column": "state", "must_all_appear": False, "codes": {c: "USPS code" for c in (
        "AL AK AZ AR CA CO CT DE DC FL GA HI ID IL IN IA KS KY LA ME MD MA MI MN MS MO MT NE NV NH NJ NM NY "
        "NC ND OH OK OR PA RI SC SD TN TX UT VT VA WA WV WI WY AS FM GU MH MP PW PR VI AA AE AP").split()}},
}

# ---- YOUR BASELINE: one row that breaks no requirement, as SQL values for your columns. ----
# {i} becomes the example row's number, so every example row gets its own id.
BASELINE = {
    "grp": "2018", "row_no": "{i}", "cand_id": "'H8NY00000'", "other_id": "'C00000001'",
    "txn_date": "DATE '2018-06-01'", "amount": "CAST(100.00 AS DECIMAL(14,2))", "sub_id": "'S{i}'",
    "txn_type": "'24K'", "state": "'NY'",
}

# ---- YOUR REQUIREMENTS ----
# id        a short name the evidence carries
# text      the business sentence an owner signs
# bad_when  a SQL condition that is TRUE for a bad row. {t} is the table name (use it in subqueries).
#           A condition that comes out NULL counts as bad: an unknown is never a pass.
# tolerance the most percent of bad rows the owner accepts in one group
# good, bad example rows, as overrides of BASELINE; the rule must flag every bad row and no good one
REQUIREMENTS = [
    {"id": "R1", "text": "Every payment names the candidate or committee that received it.",
     "bad_when": "cand_id IS NULL AND other_id IS NULL", "tolerance": 0.1,
     "good": [{}, {"cand_id": "NULL"}, {"other_id": "NULL"}],
     "bad": [{"cand_id": "NULL", "other_id": "NULL"}]},
    {"id": "R2", "text": "Every payment has an amount that is a number.",
     "bad_when": "amount IS NULL", "tolerance": 0.0,
     "good": [{}, {"amount": "CAST(-5.00 AS DECIMAL(14,2))"}],
     "bad": [{"amount": "NULL"}]},
    {"id": "R3", "text": "Every payment is dated inside the two-year election cycle it is filed under.",
     "bad_when": "txn_date IS NULL OR txn_date < MAKE_DATE(grp - 1, 1, 1) OR txn_date > MAKE_DATE(grp, 12, 31)",
     "tolerance": 1.0,
     "good": [{}, {"txn_date": "DATE '2017-01-01'"}, {"txn_date": "DATE '2018-12-31'"}],
     "bad": [{"txn_date": "DATE '2016-12-31'"}, {"txn_date": "DATE '2019-01-01'"}, {"txn_date": "NULL"}]},
    {"id": "R4", "text": "Every payment can be found by its own unique id.",
     "bad_when": "sub_id IS NULL OR (grp, sub_id) IN (SELECT grp, sub_id FROM {t} GROUP BY grp, sub_id HAVING COUNT(*) > 1)",
     "tolerance": 0.0,
     "good": [{}, {}],
     "bad": [{"sub_id": "'DUP'"}, {"sub_id": "'DUP'"}, {"sub_id": "NULL"}]},
    {"id": "R5", "text": "Every payment has a transaction type the business recognizes.",
     "bad_when": "txn_type IS NULL OR txn_type NOT IN (SELECT code FROM txn_types)", "tolerance": 0.0,
     "good": [{}, {"txn_type": "'24A'"}],
     "bad": [{"txn_type": "'99'"}, {"txn_type": "NULL"}]},
    {"id": "R6", "text": "Every payment carries a valid US state or territory code.",
     "bad_when": "state IS NULL OR state NOT IN (SELECT code FROM us_codes)", "tolerance": 2.0,
     "good": [{}, {"state": "'PR'"}],
     "bad": [{"state": "'ZZ'"}, {"state": "NULL"}, {"state": "'ny'"}]},
]

# ---- YOUR DATA: leave both as None to run the FEC demo with --data-dir. ----
DATA_CSV = None        # the path of your CSV file, which must have a header row
GROUP_COLUMN = None    # the column of that file that holds the year or period of each row

PLANTS = {"duplicate-rows", "unknown-type", "blank-recipient", "broken-rule", "unused-code"}
RUN_LOG = []


def say(text=""):
    print(text)
    RUN_LOG.append(str(text))


def fetch(url, dest):
    """Download url to dest. Print its size, Last-Modified header and the first 16 characters of its SHA-256."""
    head = urllib.request.urlopen(urllib.request.Request(url, method="HEAD"), timeout=60)
    modified = head.headers.get("Last-Modified")
    with urllib.request.urlopen(url, timeout=120) as src, open(dest, "wb") as dst:
        shutil.copyfileobj(src, dst)
    digest = hashlib.sha256(pathlib.Path(dest).read_bytes()).hexdigest()
    say(f"source  {url}  bytes={os.path.getsize(dest)}  last-modified={modified}  sha256={digest[:16]}")
    return {"url": url, "bytes": os.path.getsize(dest), "last_modified": modified, "sha256": digest}


def sql_text(path):
    """A file path as a SQL string literal: an apostrophe in the path is doubled."""
    return "'" + str(path).replace("'", "''") + "'"


def load_fec(con, data_dir, plant):
    """Download the pinned FEC files into a new, empty folder and load them, unchanged, into table raw."""
    os.makedirs(data_dir, exist_ok=True)
    if os.listdir(data_dir):
        sys.exit(f"{data_dir} is not empty: give the script a new, empty folder")
    header_path = os.path.join(data_dir, "pas2_header_file.csv")
    sources = [fetch(HEADER_URL, header_path)]
    cols = pathlib.Path(header_path).read_text(encoding="utf-8").strip().split(",")
    colspec = ", ".join(f"'{c}': 'VARCHAR'" for c in cols)
    source_rows, first = {}, True
    for cy in CYCLES:
        zpath = os.path.join(data_dir, f"pas2{str(cy)[2:]}.zip")
        sources.append(fetch(f"{BASE}/{cy}/pas2{str(cy)[2:]}.zip", zpath))
        txt = os.path.join(data_dir, f"itpas2_{cy}.txt")
        with zipfile.ZipFile(zpath) as z:
            if z.namelist() != ["itpas2.txt"]:
                sys.exit(f"unexpected zip contents: {z.namelist()}")
            with z.open("itpas2.txt") as src, open(txt, "wb") as dst:
                n_lines = 0
                while chunk := src.read(1 << 20):
                    dst.write(chunk)
                    n_lines += chunk.count(b"\n")
        source_rows[str(cy)] = n_lines
        select = (f"SELECT {cy} AS grp, row_number() OVER () AS row_no, * FROM read_csv({sql_text(txt)}, delim='|', "
                  f"header=false, columns={{{colspec}}}, quote='', escape='', encoding='latin-1')")
        con.sql(("CREATE TABLE raw AS " if first else "INSERT INTO raw ") + select)
        first = False
    say(f"total downloaded: {sum(s['bytes'] for s in sources)} bytes")
    # The cleaned view the rules read: blanks become NULL, dates and amounts get types.
    con.sql("""CREATE VIEW data AS SELECT grp, row_no,
        NULLIF(TRIM(CAND_ID), '') AS cand_id, NULLIF(TRIM(OTHER_ID), '') AS other_id,
        TRY_STRPTIME(NULLIF(TRIM(TRANSACTION_DT), ''), '%m%d%Y')::DATE AS txn_date,
        TRY_CAST(NULLIF(TRIM(TRANSACTION_AMT), '') AS DECIMAL(14, 2)) AS amount,
        NULLIF(TRIM(SUB_ID), '') AS sub_id, NULLIF(TRIM(TRANSACTION_TP), '') AS txn_type,
        NULLIF(TRIM(STATE), '') AS state FROM raw""")
    # Planted faults, for the second run. Each one is invented so the clean data cannot already hold it.
    if "blank-recipient" in plant:
        elig = con.sql("SELECT COUNT(*) FROM data WHERE grp = 2018 AND (cand_id IS NOT NULL OR other_id IS NOT NULL)").fetchone()[0]
        if elig < 2000:
            sys.exit("could not plant blank-recipient: fewer than 2,000 rows with a recipient")
        con.sql("""UPDATE raw SET CAND_ID = '', OTHER_ID = '' WHERE grp = 2018 AND row_no IN (
                   SELECT row_no FROM data WHERE grp = 2018 AND (cand_id IS NOT NULL OR other_id IS NOT NULL)
                   ORDER BY row_no LIMIT 2000)""")
        say("PLANTED: the recipient blanked on 2,000 rows of the 2018 file")
    if "unknown-type" in plant:
        if con.sql("SELECT COUNT(*) FROM raw WHERE TRANSACTION_TP = '99'").fetchone()[0]:
            sys.exit("could not plant unknown-type: the data already holds the code '99'")
        con.sql("UPDATE raw SET TRANSACTION_TP = '99' WHERE grp = 2020 AND row_no <= 500")
        say("PLANTED: 500 rows of the 2020 file given the transaction type '99'")
    if "duplicate-rows" in plant:
        con.sql("INSERT INTO raw SELECT * FROM raw WHERE grp = 2018 ORDER BY row_no DESC LIMIT 1000")
        say("PLANTED: the last 1,000 rows of the 2018 file loaded twice")
    return source_rows, sources


def load_csv(con, path, group_column):
    """Your own data: load the CSV unchanged into raw and add grp and row_no. DuckDB reads the whole file before it
    picks each column's type, so a column with a value that is not a number or a date is read as text. The rules
    read the view data, where a text value that is blank or only spaces becomes NULL, as on the demo data."""
    path = pathlib.Path(path)
    sources = [{"url": str(path), "bytes": path.stat().st_size, "last_modified": "",
                "sha256": hashlib.sha256(path.read_bytes()).hexdigest()}]
    say(f"source  {path}  bytes={sources[0]['bytes']}  sha256={sources[0]['sha256'][:16]}")
    try:
        con.sql(f"CREATE TABLE raw AS SELECT row_number() OVER () AS row_no, \"{group_column}\" AS grp, * "
                f"FROM read_csv({sql_text(path)}, header=true, sample_size=-1)")
    except duckdb.Error as e:
        sys.exit(f"could not read {path} as a CSV file with a header row: {str(e).splitlines()[0]}")
    cols = con.sql("DESCRIBE raw").fetchall()
    say("columns read as: " + ", ".join(f"{c} {t}" for c, t, *_ in cols if c not in ("grp", "row_no")))
    con.sql("CREATE VIEW data AS SELECT " + ", ".join(
        f"NULLIF(TRIM(\"{c}\"), '') AS \"{c}\"" if t == "VARCHAR" and c != "row_no" else f'"{c}"'
        for c, t, *_ in cols) + " FROM raw")
    # Count records the way DuckDB loads them: it skips blank lines, so they are not counted here either.
    with open(path, newline="", encoding="utf-8") as f:
        records = sum(1 for row in csv.reader(f) if row) - 1
    return {"__all__": records}, sources


def profile(con):
    """Share of blank values (NULL or only spaces) in every column of the loaded file, group by group."""
    groups = [r[0] for r in con.sql("SELECT DISTINCT grp FROM raw ORDER BY 1").fetchall()]
    cols = [r[0] for r in con.sql("DESCRIBE raw").fetchall() if r[0] not in ("grp", "row_no")]
    say("column profile: percent of rows with a blank value, by group")
    say("column".ljust(18) + "".join(str(g).rjust(10) for g in groups))
    out = {}
    for c in cols:
        blank = f'NULLIF(TRIM(CAST("{c}" AS VARCHAR)), \'\') IS NULL'
        rows = dict(con.sql(f'SELECT grp, 100.0 * COUNT(*) FILTER (WHERE {blank}) / COUNT(*) FROM raw GROUP BY grp').fetchall())
        out[c] = rows
        say(c.ljust(18) + "".join(f"{rows[g]:.2f}".rjust(10) for g in groups))
    return out


def sql_of(req):
    return " ".join(req["bad_when"].split())


def fingerprint(req):
    return hashlib.sha256(sql_of(req).encode()).hexdigest()[:12]


def fixture_table(con, req, rows):
    """Build the table fx: one SELECT per example row, from BASELINE with that row's overrides."""
    parts = []
    for i, (_kind, over) in enumerate(rows, start=1):
        vals = {**BASELINE, **over}
        parts.append("SELECT " + ", ".join(f"{v.replace('{i}', str(i))} AS {k}" for k, v in vals.items()))
    con.sql("DROP TABLE IF EXISTS fx")
    con.sql("CREATE TABLE fx AS " + " UNION ALL ".join(parts))


def test_rule(con, req):
    """Run the rule on its own example rows. Returns a problem text, or '' if it flags exactly the bad ones."""
    rows = [("good", o) for o in req["good"]] + [("bad", o) for o in req["bad"]]
    fixture_table(con, req, rows)
    expected = [i for i, (kind, _o) in enumerate(rows, start=1) if kind == "bad"]
    cond = sql_of(req).replace("{t}", "fx")
    got = [r[0] for r in con.sql(f"SELECT row_no FROM fx WHERE COALESCE(({cond}), TRUE) ORDER BY row_no").fetchall()]
    if got != expected:
        return f"{req['id']} flagged example rows {got}, expected exactly {expected}"
    return ""


def run_rule(con, req):
    cond = sql_of(req).replace("{t}", "data")
    tested = dict(con.sql("SELECT grp, COUNT(*) FROM data GROUP BY grp").fetchall())
    failed = dict(con.sql(f"SELECT grp, COUNT(*) FROM data WHERE COALESCE(({cond}), TRUE) GROUP BY grp").fetchall())
    result = {}
    for g in sorted(tested, key=str):
        first = [r[0] for r in con.execute(f"SELECT row_no FROM data WHERE grp IS NOT DISTINCT FROM ? "
                                           f"AND COALESCE(({cond}), TRUE) "
                                           f"ORDER BY row_no LIMIT 3", [g]).fetchall()]
        n, bad = tested[g], failed.get(g, 0)
        pct = 100.0 * bad / n
        result[str(g)] = {"tested": n, "failed": bad, "pct": pct, "tolerance": req["tolerance"],
                          "verdict": "Pass" if pct <= req["tolerance"] else "Fail", "first": first}
    return result


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data-dir", help="a new, empty folder for the FEC downloads (the demo data)")
    ap.add_argument("--csv", default=DATA_CSV, help="your own data: a CSV file with a header row (or set DATA_CSV)")
    ap.add_argument("--group-column", default=GROUP_COLUMN,
                    help="with --csv: the column that holds the year or period of each row (or set GROUP_COLUMN)")
    ap.add_argument("--out-dir", default=str(pathlib.Path(__file__).parent), help="folder for every file written")
    ap.add_argument("--log", help="also save everything printed to this file")
    ap.add_argument("--plant", default="", help="comma list of test faults (demo data only): " + ", ".join(sorted(PLANTS)))
    a = ap.parse_args()
    plant = [x for x in a.plant.split(",") if x]
    if set(plant) - PLANTS:
        ap.error(f"unknown --plant value(s) {sorted(set(plant) - PLANTS)}; use {sorted(PLANTS)}")
    if bool(a.data_dir) == bool(a.csv):
        ap.error("give exactly one of --data-dir (the demo) or --csv (your data)")
    if a.csv and (not a.group_column or plant):
        ap.error("--csv needs --group-column (or GROUP_COLUMN), and --plant only works on the demo data")
    try:
        sys.exit(run(a, plant))
    finally:
        if a.log:
            pathlib.Path(a.log).write_text("\n".join(RUN_LOG) + "\n")


def run(a, plant):
    reqs = [dict(r) for r in REQUIREMENTS]
    lookups = {k: {**v, "codes": dict(v["codes"])} for k, v in LOOKUPS.items()}
    if "broken-rule" in plant:
        reqs[2]["bad_when"] = "FALSE"
        say("PLANTED: the SQL of rule R3 replaced with FALSE, so it can never fire")
    if "unused-code" in plant:
        lookups["txn_types"]["codes"]["24X"] = "A code no row carries"
        say("PLANTED: the code '24X' added to the transaction types")
    say(f"duckdb {duckdb.__version__}, python {sys.version.split()[0]}")
    out = pathlib.Path(a.out_dir)
    dest = out / ("rule_evidence-planted.csv" if plant else "rule_evidence.csv")
    dest.unlink(missing_ok=True)   # a failed run must not leave an earlier evidence file in place
    con = duckdb.connect()
    source_rows, sources = load_fec(con, a.data_dir, plant) if a.data_dir else load_csv(con, a.csv, a.group_column)
    for name, lk in lookups.items():
        con.sql(f"CREATE TABLE {name} (code VARCHAR, meaning VARCHAR)")
        con.executemany(f"INSERT INTO {name} VALUES (?, ?)", list(lk["codes"].items()))
    say()
    profile(con)
    say()

    failures = []
    # Check 1: the loaded rows equal the source files' rows. A duplicated or lost row trips it.
    loaded = {str(g): n for g, n in con.sql("SELECT grp, COUNT(*) FROM raw GROUP BY grp").fetchall()}
    if "__all__" in source_rows:
        if sum(loaded.values()) != source_rows["__all__"]:
            failures.append(f"loaded {sum(loaded.values())} rows but the file has {source_rows['__all__']}")
    else:
        for g, n in source_rows.items():
            if loaded.get(g) != n:
                failures.append(f"group {g}: loaded {loaded.get(g)} rows but the source file has {n} "
                                f"({(loaded.get(g) or 0) - n:+d})")
    # Check 2: every rule flags exactly its own bad example rows and none of its good ones.
    say("rules (SQL fingerprint, then the SQL that is TRUE for a bad row)")
    for r in reqs:
        say(f"  {r['id']}  {fingerprint(r)}  {sql_of(r)}")
        try:
            problem = test_rule(con, r)
        except duckdb.Error as e:
            problem = f"{r['id']} could not run on its example rows: {str(e).splitlines()[0]}"
        if problem:
            failures.append("example rows: " + problem)
    # Checks 3 and 4: every value is in its lookup, and every code in a strict lookup is in the data.
    for name, lk in lookups.items():
        col = lk["column"]
        if lk["must_all_appear"]:
            unmapped = con.sql(f"SELECT {col}, COUNT(*) FROM data WHERE {col} IS NOT NULL AND {col} NOT IN "
                               f"(SELECT code FROM {name} WHERE code IS NOT NULL) GROUP BY 1 ORDER BY 1").fetchall()
            if unmapped:
                failures.append(f"{name}: " + ", ".join(f"'{c}' x {k}" for c, k in unmapped)
                                + f" are not in the {name} list")
            unused = [c[0] for c in con.sql(f"SELECT code FROM {name} WHERE code IS NOT NULL AND code NOT IN "
                                            f"(SELECT DISTINCT {col} FROM data WHERE {col} IS NOT NULL) ORDER BY 1").fetchall()]
            if unused:
                failures.append(f"{name}: the codes {unused} match no row in the data")

    results = {}
    for r in reqs:
        try:
            results[r["id"]] = run_rule(con, r)
        except duckdb.Error as e:
            results[r["id"]] = {}
            failures.append(f"{r['id']} could not run on the data: {str(e).splitlines()[0]}")
    say()
    say("rule  group  tested  failed  failed%  tolerance%  verdict  first failing row numbers")
    evidence = []
    for r in reqs:
        for g, e in results[r["id"]].items():
            say(f"{r['id']:<5} {g:<6} {e['tested']:>7} {e['failed']:>7} {e['pct']:>8.3f} {e['tolerance']:>11.3f}  "
                f"{e['verdict']:<7}  {', '.join(map(str, e['first'])) or '-'}")
            evidence.append({"group": g, "rule_id": r["id"], "requirement": r["text"], "sql_sha256_12": fingerprint(r),
                             "rows_tested": e["tested"], "rows_failed": e["failed"], "pct_failed": f"{e['pct']:.3f}",
                             "tolerance_pct": f"{e['tolerance']:.3f}", "verdict": e["verdict"],
                             "first_failing_rows": " ".join(map(str, e["first"]))})
    say()
    if failures:
        say("SELF-CHECK FAILED: no evidence CSV written")
        for f in failures:
            say("  - " + f)
        return 1
    say(f"SELF-CHECK PASSED: row counts equal the source, all {len(reqs)} rules flag exactly their own bad examples, "
        f"every value of a required lookup is listed, and every listed code is used")
    out.mkdir(parents=True, exist_ok=True)
    with open(dest, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(evidence[0]), lineterminator="\n")
        w.writeheader()
        w.writerows(evidence)
    say(f"wrote {dest.name}")
    return 0


if __name__ == "__main__":
    main()