# critwire Learn: /learn/labs/e3.py
# Run it after the lab loader (/learn/data/postgis/learn_curriculum.sql), on the same database, in the environment
# /learn/labs/requirements.txt pins (the setup page, /learn/setup/, gives the uv commands):
#   python e3.py
# It connects with libpq's PG* variables (PGHOST, PGPORT, PGDATABASE, PGUSER, PGPASSWORD) and reads the lab data from
# LEARN_BASE (default https://critwire.com). Each block runs once, in page order, after py-connect and the boot setup;
# an untested block is commented out under its reason.

# %% py-connect
# q: py-connect
# status: ran
import psycopg


def connect():
    """A connection to the lab database: plain table names find learn_curriculum's tables, then PostGIS in public."""
    return psycopg.connect(options="-c search_path=learn_curriculum,public")

# %% e3-rerun
# q: e3-rerun
# status: ran
# after: py-connect
import hashlib

SETUP = """
CREATE TEMP TABLE devil_records (id bigint PRIMARY KEY, year integer, longitude float8, latitude float8);
CREATE TEMP TABLE devil_records_naive (id bigint, year integer, longitude float8, latitude float8);
CREATE TEMP TABLE load_runs (run_id text PRIMARY KEY, input_rows integer NOT NULL)
"""


def checksum(cur, table):
    """md5 of every row in key order: two loads that leave the same rows give the same value."""
    return cur.execute(f"SELECT count(*), md5(coalesce(string_agg(t::text, E'\\n' ORDER BY t::text), '')) "
                       f"FROM {table} AS t").fetchone()


def load(cur, batch):
    """Idempotent: the run id is the input's sha256, so the same input is loaded once; rows upsert by key."""
    run_id = hashlib.sha256("\n".join(",".join(map(str, r)) for r in batch).encode()).hexdigest()
    if cur.execute("INSERT INTO load_runs VALUES (%s, %s) ON CONFLICT DO NOTHING RETURNING run_id",
                   (run_id, len(batch))).fetchone() is None:
        return run_id, "already loaded: skipped"
    cur.executemany("INSERT INTO devil_records VALUES (%s, %s, %s, %s) ON CONFLICT (id) DO UPDATE "
                    "SET year = EXCLUDED.year, longitude = EXCLUDED.longitude, latitude = EXCLUDED.latitude", batch)
    return run_id, f"loaded {len(batch)} rows"


with connect() as conn, conn.cursor() as cur:
    cur.execute(SETUP)
    batch = cur.execute("""SELECT o.id, o.year, o.longitude, o.latitude FROM occurrences AS o
                           JOIN species AS s ON s.id = o.species_id
                           WHERE s.scientific_name = 'Sarcophilus harrisii' ORDER BY o.id""").fetchall()
    for run in (1, 2):
        cur.executemany("INSERT INTO devil_records_naive VALUES (%s, %s, %s, %s)", batch)
        rows, md5 = checksum(cur, "devil_records_naive")
        print(f"append, run {run}: {rows} rows, md5 {md5[:12]}")
    for run in (1, 2):
        run_id, what = load(cur, batch)
        rows, md5 = checksum(cur, "devil_records")
        print(f"idempotent, run {run}: run {run_id[:12]}, {what}; {rows} rows, md5 {md5[:12]}")
    corrected = [(batch[0][0], 1999, *batch[0][2:]), *batch[1:]]  # next week the publisher adds one missing year
    run_id, what = load(cur, corrected)
    rows, md5 = checksum(cur, "devil_records")
    print(f"idempotent, run 3 (one year added): run {run_id[:12]}, {what}; {rows} rows, md5 {md5[:12]}")
    conn.rollback()

# %% e3-snapshot
# q: e3-snapshot
# status: ran
# after: py-connect
import hashlib
import json

from psycopg.types.json import Jsonb

# The target's DDL by SCHEMA_VERSION: version 2 adds a CHECK on the license.
SCHEMAS = {
    1: "CREATE TEMP TABLE devil_target (id bigint PRIMARY KEY, year integer, license text NOT NULL)",
    2: "CREATE TEMP TABLE devil_target (id bigint PRIMARY KEY, year integer, "
       "license text NOT NULL CHECK (license IN ('CC0 1.0', 'CC BY 4.0')))",
}


def md5_order(row):
    return hashlib.md5(str(row[0]).encode()).hexdigest(), row[0]


def transform_v1(rows, config):
    """The rows the config keeps: records under its licenses, the first `cap` by the md5 of their id."""
    return sorted((r for r in rows if r[2] in config["licenses"]), key=md5_order)[:config["cap"]]


def transform_v2(rows, config):
    """Version 2 also drops records without a year."""
    return transform_v1([r for r in rows if r[1] is not None], config)


TRANSFORMS = {1: transform_v1, 2: transform_v2}


def snapshot(rows):
    """A delivered file: one line per record, `id,year,license`, an empty year for none."""
    return "".join(f"{i},{'' if y is None else y},{lic}\n" for i, y, lic in rows).encode()


def parse(data):
    return [(int(i), int(y) if y else None, lic) for i, y, lic in (line.split(",") for line in data.decode().splitlines())]


def job_key(parts):
    """One sha256 over the four parts, as canonical JSON (sorted keys, no spaces), so equal parts give one key."""
    return hashlib.sha256(json.dumps(parts, sort_keys=True, separators=(",", ":")).encode()).hexdigest()


def refresh(conn, data, config, transform, schema):
    """Make the target hold `data` under this transform, config and schema, unless the state says it already does.
    One transaction: the target's rows and the state row change together, or neither does."""
    parts = {"input": hashlib.sha256(data).hexdigest(), "transform": transform, "config": config, "schema": schema}
    key = job_key(parts)
    with conn.transaction(), conn.cursor() as cur:
        state = cur.execute("SELECT job_key, parts FROM job_state WHERE job = 'devil' FOR UPDATE").fetchone()
        if state is not None and state[0] == key:
            return "nothing changed: skipped"
        if state is None or state[1]["schema"] != schema:
            cur.execute("DROP TABLE IF EXISTS devil_target")
            cur.execute(SCHEMAS[schema])
        else:
            cur.execute("TRUNCATE devil_target")
        rows = TRANSFORMS[transform](parse(data), config)
        cur.executemany("INSERT INTO devil_target VALUES (%s, %s, %s)", rows)
        cur.execute("INSERT INTO job_state VALUES ('devil', %s, %s) "
                    "ON CONFLICT (job) DO UPDATE SET job_key = EXCLUDED.job_key, parts = EXCLUDED.parts",
                    (key, Jsonb(parts)))
    changed = "first load" if state is None else " and ".join(p for p in sorted(parts) if parts[p] != state[1][p]) + " changed"
    return f"{changed}: loaded {len(rows)} rows"


def ledger(conn, data, config):
    """e3-rerun's rule: the run id is the input's checksum alone, and a seen run id is skipped."""
    with conn.transaction(), conn.cursor() as cur:
        run_id = hashlib.sha256(data).hexdigest()
        if cur.execute("INSERT INTO load_runs VALUES (%s) ON CONFLICT DO NOTHING RETURNING run_id",
                       (run_id,)).fetchone() is None:
            return "run id seen before: skipped"
        rows = transform_v1(parse(data), config)
        cur.executemany("INSERT INTO ledger_target VALUES (%s, %s, %s) ON CONFLICT (id) DO UPDATE "
                        "SET year = EXCLUDED.year, license = EXCLUDED.license", rows)
        return f"loaded {len(rows)} rows"


def holds(conn, table, inputs, config, transform):
    """The input whose rows, under these settings, are exactly the table's rows."""
    rows = set(conn.execute(f"SELECT id, year, license FROM {table}").fetchall())
    return next(name for name, data in inputs.items() if rows == set(TRANSFORMS[transform](parse(data), config)))


with connect() as conn:
    conn.autocommit = True  # each `with conn.transaction()` is then one transaction of its own
    conn.execute("CREATE TEMP TABLE job_state (job text PRIMARY KEY, job_key text NOT NULL, parts jsonb NOT NULL)")
    conn.execute("CREATE TEMP TABLE load_runs (run_id text PRIMARY KEY)")
    conn.execute(SCHEMAS[1].replace("devil_target", "ledger_target"))
    try:
        rows = conn.execute("""SELECT o.id, o.year, o.license FROM occurrences AS o
                               JOIN species AS s ON s.id = o.species_id
                               WHERE s.scientific_name = 'Sarcophilus harrisii' ORDER BY o.id""").fetchall()
        fixed = min((r for r in rows if r[1] is not None), key=md5_order)  # next week the publisher corrects one year
        inputs = {"A": snapshot(rows), "B": snapshot([(i, y - 1, lic) if i == fixed[0] else (i, y, lic) for i, y, lic in rows])}
        default = {"cap": 1000, "licenses": ["CC0 1.0", "CC BY 4.0"]}
        smaller = {**default, "cap": 500}
        runs = [("A", default, 1, 1), ("A", default, 1, 1), ("B", default, 1, 1), ("A", default, 1, 1),
                ("A", smaller, 1, 1), ("A", smaller, 2, 1), ("A", smaller, 2, 2), ("A", smaller, 2, 2)]
        for n, (name, config, transform, schema) in enumerate(runs, 1):
            what = refresh(conn, inputs[name], config, transform, schema)
            settings = f"cap {config['cap']}, transform {transform}, schema {schema}"
            print(f"run {n} · {name} ({settings}) · {what}; the target holds "
                  f"{holds(conn, 'devil_target', inputs, config, transform)}")
            if n <= 4:
                what = ledger(conn, inputs[name], default)
                print(f"  ledger rule · {name} · {what}; its target holds {holds(conn, 'ledger_target', inputs, default, 1)}")
    finally:
        conn.execute("DROP TABLE IF EXISTS job_state, load_runs, ledger_target, devil_target")

# %% e3-ex
# q: e3-ex
# status: ran
# after: py-connect
import hashlib

EXPORTS = {
    "without ORDER BY": "SELECT id, year, license FROM devil",
    "ORDER BY id": "SELECT id, year, license FROM devil ORDER BY id",
}


def export(cur, query):
    """The query's rows as a CSV file's bytes, and their sha256."""
    with cur.copy(f"COPY ({query}) TO STDOUT (FORMAT csv, HEADER true)") as copy:
        return hashlib.sha256(b"".join(copy)).hexdigest()


with connect() as conn, conn.cursor() as cur:
    cur.execute("""CREATE TEMP TABLE devil AS
                   SELECT o.id, o.year, o.license FROM occurrences AS o JOIN species AS s ON s.id = o.species_id
                   WHERE s.scientific_name = 'Sarcophilus harrisii' ORDER BY o.id""")
    before = {name: export(cur, query) for name, query in EXPORTS.items()}
    cur.execute("UPDATE devil SET year = year WHERE id = (SELECT min(id) FROM devil)")  # changes no value
    for name, query in EXPORTS.items():
        after = export(cur, query)
        print(f"{name}: sha256 {before[name][:12]} then {after[:12]}: {'same file' if after == before[name] else 'a different file'}")
    conn.rollback()