# critwire Learn: /learn/labs/e5.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 e5.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")

# %% e5-job
# q: e5-job
# status: ran
# after: py-connect
# A simulation: both "weeks" run in this one process, on TEMP tables that vanish when it ends. A scheduler runs a job as
# a new process each time, against tables that persist; e5-refresh writes that job, and e5-runs runs it.
import hashlib
import json
import logging
import os
import sys
import urllib.request

BASE = os.environ.get("LEARN_BASE", "https://critwire.com")
FILE = "datasets.csv"
# A scheduled job writes a timestamp on every line (add %(asctime)s); this lab leaves it out so its output can be checked.
logging.basicConfig(stream=sys.stdout, level=logging.INFO, format="%(levelname)s %(name)s: %(message)s")
log = logging.getLogger("refresh")


def read(url):
    with urllib.request.urlopen(url, timeout=60) as r:
        return r.read()


def refresh(conn):
    """One scheduled run: nothing to do if the published file is unchanged; else reload it in one transaction, checked."""
    entry = json.loads(read(f"{BASE}/learn/data/manifest.json"))["files"][FILE]
    last = conn.execute("SELECT sha256 FROM refresh_state WHERE file = %s", (FILE,)).fetchone()
    if last is not None and last[0] == entry["sha256"]:
        log.info("%s unchanged (sha256 %s): nothing to do", FILE, entry["sha256"][:12])
        return
    data = read(f"{BASE}/learn/data/{FILE}")
    if hashlib.sha256(data).hexdigest() != entry["sha256"]:
        raise RuntimeError(f"{FILE} doesn't match its manifest")
    with conn.transaction():  # a failed check below leaves the previous copy and state as they were
        conn.execute("TRUNCATE my_datasets")
        with conn.cursor().copy("COPY my_datasets FROM STDIN (FORMAT csv, HEADER true)") as copy:
            copy.write(data)
        rows = conn.execute("SELECT count(*) FROM my_datasets").fetchone()[0]
        if rows != entry["rows"]:
            raise RuntimeError(f"loaded {rows} rows of {FILE}; the manifest says {entry['rows']}")
        conn.execute("INSERT INTO refresh_state VALUES (%s, %s) ON CONFLICT (file) DO UPDATE SET sha256 = EXCLUDED.sha256",
                     (FILE, entry["sha256"]))
    log.info("%s loaded: %d rows, sha256 %s", FILE, rows, entry["sha256"][:12])


def main():
    """Exit 0 on success, 1 on any failure, with the traceback in the output. The scheduler sees only the exit code:
    systemd marks the unit failed; cron mails the output only if the crontab line doesn't redirect it to a file."""
    try:
        with connect() as conn:
            # TEMP tables: the simulation leaves learn_curriculum as it was; a real job's tables live in its own schema.
            conn.execute("CREATE TEMP TABLE my_datasets (LIKE datasets)")
            conn.execute("CREATE TEMP TABLE refresh_state (file text PRIMARY KEY, sha256 text NOT NULL)")
            for run in (1, 2):  # this week's run, then next week's with nothing new
                log.info("run %d", run)
                refresh(conn)
    except Exception:
        log.exception("refresh failed")
        return 1
    return 0


# A scheduled script ends with `raise SystemExit(main())`; here the lab goes on, so only a failure stops it.
exit_code = main()
if exit_code:
    raise SystemExit(exit_code)

# %% e5-refresh
# q: e5-refresh
# status: ran
from pathlib import Path

JOB = r'''"""refresh.py: makes <REFRESH_SCHEMA>.datasets hold the published datasets.csv, under this code and config.
Exit 0: loaded, or nothing to do. Exit 1: failed, and nothing changed. The reason goes to stdout, the traceback to stderr."""
import hashlib
import json
import logging
import os
import sys
import traceback
import urllib.request
from pathlib import Path

import psycopg
from psycopg import sql
from psycopg.types.json import Jsonb

FILE = "datasets.csv"
TRANSFORM_VERSION = 1  # bump when this code's output changes
SCHEMA_VERSION = 1  # bump when the target's DDL changes
SETTINGS = ("PGHOST", "PGPORT", "PGUSER", "PGDATABASE")

# A scheduled job writes a timestamp on every line (add %(asctime)s); this lab leaves it out so its output can be checked.
logging.basicConfig(stream=sys.stdout, level=logging.INFO, format="%(levelname)s %(name)s: %(message)s")
log = logging.getLogger("refresh")


def missing_settings():
    """The connection settings this environment lacks, by name: no value, the password least of all, reaches the log."""
    missing = [name for name in SETTINGS if not os.environ.get(name)]
    default = (Path(os.environ.get("APPDATA", "")) / "postgresql" / "pgpass.conf" if os.name == "nt"
               else Path.home() / ".pgpass")
    if not os.environ.get("PGPASSWORD") and not Path(os.environ.get("PGPASSFILE", default)).is_file():
        missing.append("PGPASSWORD or a password file")
    return missing


def read(url):
    with urllib.request.urlopen(url, timeout=60) as r:
        return r.read()


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


def refresh(conn, base, schema, licenses):
    """One run in one transaction: the lock, the comparison, the replacement and the new state commit together, or
    nothing does. Returns the log line."""
    entry = json.loads(read(f"{base}/learn/data/manifest.json"))["files"][FILE]
    parts = {"input": entry["sha256"], "transform": TRANSFORM_VERSION,
             "config": {"file": FILE, "licenses": licenses}, "schema": SCHEMA_VERSION}
    key = job_key(parts)
    target, state = sql.Identifier(schema, "datasets"), sql.Identifier(schema, "refresh_state")
    with conn.transaction():
        # The transaction's own lock: COMMIT or ROLLBACK frees it, however the job ends. Held means another run is at work.
        if not conn.execute("SELECT pg_try_advisory_xact_lock(hashtext(%s))", (f"refresh:{schema}",)).fetchone()[0]:
            return "another refresh holds the lock: skipped"
        last = conn.execute(sql.SQL("SELECT job_key, parts FROM {} WHERE job = %s").format(state), (FILE,)).fetchone()
        if last is not None and last[0] == key:
            return "nothing changed: skipped"
        data = read(f"{base}/learn/data/{FILE}")
        if hashlib.sha256(data).hexdigest() != entry["sha256"]:
            raise RuntimeError(f"{FILE} doesn't match its manifest's sha256")
        conn.execute(sql.SQL("CREATE TEMP TABLE stage (LIKE {}) ON COMMIT DROP").format(target))
        with conn.cursor().copy("COPY pg_temp.stage FROM STDIN (FORMAT csv, HEADER true)") as copy:
            copy.write(data)
        staged = conn.execute("SELECT count(*) FROM pg_temp.stage").fetchone()[0]
        if staged != entry["rows"]:
            raise RuntimeError(f"staged {staged} rows; the manifest says {entry['rows']}")
        conn.execute(sql.SQL("TRUNCATE {}").format(target))
        loaded = conn.execute(sql.SQL("INSERT INTO {} SELECT * FROM pg_temp.stage WHERE license = ANY(%s)").format(target),
                              (licenses,)).rowcount
        conn.execute(sql.SQL("INSERT INTO {} VALUES (%s, %s, %s) ON CONFLICT (job) DO UPDATE "
                             "SET job_key = EXCLUDED.job_key, parts = EXCLUDED.parts").format(state),
                     (FILE, key, Jsonb(parts)))
    if last is None:
        return f"first load: loaded {loaded} rows"
    changed = " and ".join(p for p in sorted(parts) if parts[p] != last[1].get(p))
    return f"{changed} changed: loaded {loaded} rows"


def main():
    missing = missing_settings()
    if missing:
        log.error("failed: not set in this environment: %s", ", ".join(missing))
        return 1
    base = os.environ.get("LEARN_BASE", "https://critwire.com")
    schema = os.environ.get("REFRESH_SCHEMA", "etl")
    licenses = sorted(name.strip() for name in os.environ.get("REFRESH_LICENSES", "CC0 1.0,CC BY 4.0").split(","))
    try:
        with psycopg.connect(autocommit=True) as conn:  # autocommit: conn.transaction() is then the one transaction
            log.info(refresh(conn, base, schema, licenses))
    except Exception as e:
        log.error("failed: %s", e)
        traceback.print_exc()
        return 1
    return 0


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

Path("refresh.py").write_text(JOB, encoding="utf-8")

# %% e5-runs
# q: e5-runs
# status: ran
# after: py-connect, e5-refresh
import csv
import hashlib
import io
import json
import os
import subprocess
import sys
import tempfile
import urllib.request
from pathlib import Path

BASE = os.environ.get("LEARN_BASE", "https://critwire.com")
SCHEMA = "learn_e5"
LICENSES = {"CC0 1.0", "CC BY 4.0"}  # refresh.py's default config


def rows_of(data):
    return [tuple(r) for r in csv.reader(io.StringIO(data.decode("utf-8"), newline=""))][1:]


def csv_file(header, rows):
    out = io.StringIO()
    csv.writer(out, lineterminator="\n").writerows([header, *rows])
    return out.getvalue().encode("utf-8")


def publish(folder, data, rows):
    """A copy of /learn/data/ on this machine holding `data`, its manifest saying `rows` rows; returns its LEARN_BASE."""
    d = Path(folder) / "learn" / "data"
    d.mkdir(parents=True)
    (d / "datasets.csv").write_bytes(data)
    entry = {"sha256": hashlib.sha256(data).hexdigest(), "rows": rows}
    (d / "manifest.json").write_text(json.dumps({"files": {"datasets.csv": entry}}), encoding="utf-8")
    return Path(folder).as_uri()


def start(base=None, licenses=None, without=None):
    """refresh.py as a new process, with this block's environment and these changes."""
    env = {k: v for k, v in os.environ.items() if k != without}
    env["REFRESH_SCHEMA"] = SCHEMA
    if base:
        env["LEARN_BASE"] = base
    if licenses:
        env["REFRESH_LICENSES"] = licenses
    return subprocess.Popen([sys.executable, "-I", "refresh.py"], env=env, text=True,
                            stdout=subprocess.PIPE, stderr=subprocess.PIPE)


def finish(proc):
    out, err = proc.communicate(timeout=120)
    return proc.returncode, out.strip(), " · traceback on stderr" if "Traceback" in err else ""


def show(n, label, proc):
    code, out, err = finish(proc)
    print(f"run {n} · {label} · exit {code} · {out}{err}")


def holds(conn, inputs):
    """Which input the target's rows and the state's key come from: read, not assumed."""
    rows = set(conn.execute(f"SELECT CAST(dataset_key AS text), title, publisher, coalesce(doi, ''), license "
                            f"FROM {SCHEMA}.datasets").fetchall())
    target = next((n for n, data in inputs.items() if rows == {r for r in rows_of(data) if r[4] in LICENSES}), "no input")
    sha = conn.execute(f"SELECT parts->>'input' FROM {SCHEMA}.refresh_state WHERE job = 'datasets.csv'").fetchone()[0]
    state = next((n for n, data in inputs.items() if hashlib.sha256(data).hexdigest() == sha), "no input")
    print(f"  the target holds {target}; the state's input is {state}")


with connect() as conn, tempfile.TemporaryDirectory() as tmp:
    conn.autocommit = True  # each statement commits, so the job's own sessions see the tables
    conn.execute(f"CREATE SCHEMA {SCHEMA}")
    try:
        conn.execute(f"CREATE TABLE {SCHEMA}.datasets (dataset_key uuid PRIMARY KEY, title text NOT NULL, "
                     "publisher text NOT NULL, doi text, license text NOT NULL)")
        conn.execute(f"CREATE TABLE {SCHEMA}.refresh_state (job text PRIMARY KEY, job_key text NOT NULL, parts jsonb NOT NULL)")
        with urllib.request.urlopen(f"{BASE}/learn/data/datasets.csv", timeout=60) as r:
            a = r.read()
        header, *rows = csv.reader(io.StringIO(a.decode("utf-8"), newline=""))
        b_rows = [[*rows[0][:1], rows[0][1] + " (corrected)", *rows[0][2:]], *rows[1:]]
        made_up = ["00000000-0000-4000-8000-000000000001", "A made-up dataset", "Nobody", "", "CC0 1.0"]
        inputs = {"A": a, "B": csv_file(header, b_rows), "C": csv_file(header, [*b_rows, made_up])}
        b_base = publish(Path(tmp) / "b", inputs["B"], len(b_rows))
        c_base = publish(Path(tmp) / "c", inputs["C"], len(b_rows))  # one row more than its manifest says

        show(1, "A", start())
        show(2, "A again", start())
        show(3, "A, REFRESH_LICENSES=CC0 1.0", start(licenses="CC0 1.0"))
        show(4, "B, one title corrected", start(base=b_base))
        show(5, "C, one row added, its manifest still says 400 rows", start(base=c_base))
        holds(conn, inputs)
        lock = f"refresh:{SCHEMA}"
        conn.execute("SELECT pg_advisory_lock(hashtext(%s))", (lock,))  # as a run that is still going holds it
        try:
            show(6, "A, the lock held by another session", start())
        finally:
            conn.execute("SELECT pg_advisory_unlock(hashtext(%s))", (lock,))
        holds(conn, inputs)
        both = [finish(p) for p in [start(), start()]]
        # Which one wins, and whether the other meets the lock or the new state, varies from run to run.
        outcomes = sorted("loaded" if ": loaded " in out else "not loaded" for _, out, _ in both)
        print(f"run 7 · A, two runs started together · exit {' and '.join(str(c) for c, _, _ in both)} · {', '.join(outcomes)}")
        holds(conn, inputs)
        show(8, "A, without PGDATABASE", start(without="PGDATABASE"))
    finally:
        conn.execute(f"DROP SCHEMA {SCHEMA} CASCADE")

# %% e5-ex
# q: e5-ex
# status: ran
# after: py-connect, e5-refresh
import os
import subprocess
import sys
import tempfile
from pathlib import Path

SCHEMA = "learn_e5ex"
JOB = Path("refresh.py").read_text(encoding="utf-8")


def edit(text, old, new):
    """`text` with its one `old` replaced by `new`: an edit that finds nothing to change fails, it doesn't pass silently."""
    if text.count(old) != 1:
        raise ValueError(f"expected one {old!r} in refresh.py")
    return text.replace(old, new)


# The new transform: each DOI stored as a link, https://doi.org/<doi>. Same input, same config, a different output.
LINKS = edit(JOB, "SELECT * FROM pg_temp.stage WHERE",
             "SELECT dataset_key, title, publisher, 'https://doi.org/' || doi, license FROM pg_temp.stage WHERE")


def run(conn, n, label, script):
    """`script` as a new process, with this block's environment, then the target's DOIs that are links."""
    done = subprocess.run([sys.executable, "-I", str(script)], env={**os.environ, "REFRESH_SCHEMA": SCHEMA},
                          text=True, capture_output=True, timeout=120)
    links = conn.execute(f"SELECT count(*) FROM {SCHEMA}.datasets WHERE starts_with(doi, 'https://doi.org/')").fetchone()[0]
    print(f"run {n} · {label} · exit {done.returncode} · {done.stdout.strip()} · {links} DOIs as links")


with connect() as conn, tempfile.TemporaryDirectory() as tmp:
    conn.autocommit = True  # each statement commits, so the job's own sessions see the tables
    unbumped, bumped = Path(tmp) / "refresh_links.py", Path(tmp) / "refresh_v2.py"
    unbumped.write_text(LINKS, encoding="utf-8")
    bumped.write_text(edit(LINKS, "TRANSFORM_VERSION = 1", "TRANSFORM_VERSION = 2"), encoding="utf-8")
    conn.execute(f"CREATE SCHEMA {SCHEMA}")
    try:
        conn.execute(f"CREATE TABLE {SCHEMA}.datasets (dataset_key uuid PRIMARY KEY, title text NOT NULL, "
                     "publisher text NOT NULL, doi text, license text NOT NULL)")
        conn.execute(f"CREATE TABLE {SCHEMA}.refresh_state (job text PRIMARY KEY, job_key text NOT NULL, parts jsonb NOT NULL)")
        run(conn, 1, "the old job", "refresh.py")
        run(conn, 2, "links, TRANSFORM_VERSION still 1", unbumped)
        run(conn, 3, "links, TRANSFORM_VERSION = 2", bumped)
        run(conn, 4, "links, TRANSFORM_VERSION = 2 again", bumped)
    finally:
        conn.execute(f"DROP SCHEMA {SCHEMA} CASCADE")