"""Shared helpers for the FIRC-0001 Stage D read-only audits (run on M3U96b).

Every script opens the two DuckDB files with read_only=True, SET threads=4,
SET memory_limit='24GB', writes CSV/JSON under ~/firc0001_review/results/<name>/
and a manifest.json (script sha256, DuckDB mtimes, start/end, row counts).
Nothing is ever written under ~/UQO_RESEARCH_DB.
"""
from __future__ import annotations

import csv
import datetime as dt
import hashlib
import json
import os
import sys
import time
from pathlib import Path

import duckdb

HOME = Path(os.environ.get("HOME", "~")).expanduser()
RDB = HOME / "UQO_RESEARCH_DB"
SPID_DB = RDB / "project-db" / "data" / "db" / "spid.duckdb"
EDGAR_DB = RDB / "data" / "duckdb" / "edgar_research.duckdb"
WORK = HOME / "firc0001_review"
RESULTS = WORK / "results"
THREADS = int(os.environ.get("FIRC_THREADS", "4"))
MEMLIMIT = os.environ.get("FIRC_MEMLIMIT", "24GB")
SEED = 20261105


def connect(path: Path = SPID_DB) -> duckdb.DuckDBPyConnection:
    con = duckdb.connect(str(path), read_only=True)
    con.execute(f"SET threads={THREADS}")
    con.execute(f"SET memory_limit='{MEMLIMIT}'")
    con.execute("SET preserve_insertion_order=false")
    # spill files must never land next to the databases (default = <db>.tmp)
    (WORK / "tmp").mkdir(parents=True, exist_ok=True)
    con.execute(f"SET temp_directory='{WORK / 'tmp'}'")
    return con


def attach_edgar(con: duckdb.DuckDBPyConnection, alias: str = "edgar") -> None:
    con.execute(f"ATTACH '{EDGAR_DB}' AS {alias} (READ_ONLY)")


def sha256_file(p: Path) -> str:
    h = hashlib.sha256()
    with open(p, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


class Run:
    """Context for one audit script: output dir, timing, manifest."""

    def __init__(self, name: str, script: str | None = None, out: str | None = None):
        self.name = name
        self.script = Path(script or sys.argv[0]).resolve()
        self.out = Path(out) if out else RESULTS / name
        self.out.mkdir(parents=True, exist_ok=True)
        self.t0 = time.time()
        self.started = dt.datetime.now().isoformat(timespec="seconds")
        self.rows: dict[str, int] = {}
        self.notes: dict = {}

    def write_csv(self, fname: str, cols: list[str], rows) -> int:
        rows = list(rows)
        with open(self.out / fname, "w", newline="") as f:
            w = csv.writer(f)
            w.writerow(cols)
            w.writerows(rows)
        self.rows[fname] = len(rows)
        print(f"[{self.name}] {fname}: {len(rows)} rows", flush=True)
        return len(rows)

    def dump(self, fname: str, con, sql: str, params=None):
        cur = con.execute(sql, params) if params else con.execute(sql)
        cols = [d[0] for d in cur.description]
        return self.write_csv(fname, cols, cur.fetchall())

    def write_json(self, fname: str, obj) -> None:
        with open(self.out / fname, "w") as f:
            json.dump(obj, f, indent=2, default=str)
        print(f"[{self.name}] {fname} written", flush=True)

    def finish(self, extra: dict | None = None) -> None:
        man = {
            "package": self.name,
            "script": str(self.script),
            "script_sha256": sha256_file(self.script) if self.script.exists() else None,
            "started": self.started,
            "ended": dt.datetime.now().isoformat(timespec="seconds"),
            "wall_seconds": round(time.time() - self.t0, 1),
            "node": os.uname().nodename,
            "python": sys.version.split()[0],
            "duckdb": duckdb.__version__,
            "threads": THREADS,
            "memory_limit": MEMLIMIT,
            "seed": SEED,
            "databases": {
                str(SPID_DB): {"bytes": SPID_DB.stat().st_size,
                               "mtime": dt.datetime.fromtimestamp(SPID_DB.stat().st_mtime).isoformat()},
                str(EDGAR_DB): {"bytes": EDGAR_DB.stat().st_size,
                                "mtime": dt.datetime.fromtimestamp(EDGAR_DB.stat().st_mtime).isoformat()},
            },
            "read_only": True,
            "row_counts": self.rows,
            "notes": self.notes,
        }
        if extra:
            man.update(extra)
        with open(self.out / "manifest.json", "w") as f:
            json.dump(man, f, indent=2, default=str)
        print(f"[{self.name}] done in {man['wall_seconds']} s", flush=True)