#!/usr/bin/env python3
"""Rewrite the intraday Parquet files of the lake with small row groups (65 536 rows) — resumable, file by file.

Why: the files written by frd_downloader.py before 2026-09 have 1 000 000-row groups (3 groups for a 2.5 M-row
1-minute file). A 5-minute window forces DuckDB to decode a whole third of the file; with 65 536-row groups
(≈ one week of 1-minute bars) the min/max statistics let it skip everything but one group.

Each file is rewritten with DuckDB (`COPY (SELECT * FROM read_parquet(src)) TO tmp (FORMAT PARQUET, COMPRESSION
ZSTD, ROW_GROUP_SIZE 65536)`), verified (row count identical) and swapped in atomically with `os.replace`.
Progress is kept in a JSON manifest so the job can be stopped and resumed (nightly PM2 task).

Usage (production node M3U96b — do NOT run while frd_downloader.py rewrites the same directory):
    cd ~/hfmarketdata && venv/bin/python scripts/rewrite_row_groups.py --data-root ~/firstratedata --dry-run
    … --asset stock --timeframe 1min                 one directory family
    … --workers 3 --row-group-size 65536             parallel files (each worker uses 2 DuckDB threads)
    … --max-files 500                                bounded nightly batch
    … --manifest ~/firstratedata/state/row_groups.json

Author: Simon-Pierre Boucher <contact@spboucher.ai>
"""
from __future__ import annotations

import argparse
import json
import logging
import os
import sys
import tempfile
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from pathlib import Path

import duckdb

log = logging.getLogger("rewrite_row_groups")

ASSETS_WITH_ADJ = ("stock", "etf", "crypto", "index", "fx", "futures", "futures_contracts")
INTRADAY = ("1min", "5min", "30min", "1hour")
DEFAULT_ROW_GROUP = 65_536


# ---- manifest -------------------------------------------------------------------------------------------

class Manifest:
    """{"files": {"<relative path>": {"status": "done"|"failed"|"skipped", "row_groups": n, "rows": n, "at": iso, ...}}}"""

    def __init__(self, path: Path) -> None:
        self.path = path
        self.lock = threading.Lock()
        self.data: dict = {"version": 1, "row_group_size": None, "files": {}}
        if path.is_file():
            try:
                self.data = json.loads(path.read_text())
            except json.JSONDecodeError:
                log.warning("manifest %s unreadable, starting over", path)
        self.data.setdefault("files", {})
        self._dirty = 0

    def status(self, rel: str) -> str | None:
        return (self.data["files"].get(rel) or {}).get("status")

    def mark(self, rel: str, status: str, **extra) -> None:
        with self.lock:
            self.data["files"][rel] = {"status": status, "at": datetime.now(timezone.utc).isoformat(timespec="seconds"), **extra}
            self._dirty += 1
            if self._dirty >= 20:
                self._flush()

    def _flush(self) -> None:
        self.path.parent.mkdir(parents=True, exist_ok=True)
        tmp = self.path.with_suffix(".tmp")
        tmp.write_text(json.dumps(self.data, indent=1, sort_keys=True))
        os.replace(tmp, self.path)
        self._dirty = 0

    def flush(self) -> None:
        with self.lock:
            self._flush()


# ---- discovery --------------------------------------------------------------------------------------------

def candidate_files(parquet_root: Path, assets: list[str] | None, timeframes: list[str]) -> list[Path]:
    out: list[Path] = []
    for asset in assets or ASSETS_WITH_ADJ:
        for tf in timeframes:
            base = parquet_root / asset / tf
            if not base.is_dir():
                continue
            for adj_dir in sorted(p for p in base.iterdir() if p.is_dir()):
                with os.scandir(adj_dir) as it:
                    out.extend(Path(e.path) for e in it if e.name.endswith(".parquet") and e.is_file())
    return sorted(out)


def parquet_layout(con: duckdb.DuckDBPyConnection, path: Path) -> tuple[int, int]:
    """(row_groups, rows) from the file footer — cheap."""
    p = str(path).replace("'", "''")
    rg = con.execute(f"SELECT count(DISTINCT row_group_id), coalesce(sum(row_group_num_rows), 0) "
                     f"FROM parquet_metadata('{p}') WHERE column_id = 0").fetchone()
    return int(rg[0]), int(rg[1])


# ---- rewrite ----------------------------------------------------------------------------------------------

def rewrite_one(path: Path, row_group_size: int, threads: int, dry_run: bool) -> dict:
    con = duckdb.connect()
    con.execute(f"SET threads TO {threads}")
    con.execute("SET memory_limit = '2GB'")
    try:
        groups, rows = parquet_layout(con, path)
        if groups and rows and rows / groups <= row_group_size * 1.5:
            return {"status": "skipped", "reason": "row groups already small", "row_groups": groups, "rows": rows}
        if dry_run:
            return {"status": "dry-run", "row_groups": groups, "rows": rows,
                    "target_groups": -(-rows // row_group_size) if rows else 0}
        src = str(path).replace("'", "''")
        fd, tmp_name = tempfile.mkstemp(prefix=f".{path.stem}.", suffix=".parquet.tmp", dir=str(path.parent))
        os.close(fd)
        tmp = Path(tmp_name)
        try:
            t0 = time.perf_counter()
            dst = str(tmp).replace("'", "''")
            con.execute(f"COPY (SELECT * FROM read_parquet('{src}')) TO '{dst}' "
                        f"(FORMAT PARQUET, COMPRESSION ZSTD, ROW_GROUP_SIZE {int(row_group_size)})")
            new_groups, new_rows = parquet_layout(con, tmp)
            if new_rows != rows:
                raise RuntimeError(f"row count mismatch after rewrite: {rows} → {new_rows}")
            old_size, new_size = path.stat().st_size, tmp.stat().st_size
            os.replace(tmp, path)          # atomic on the same filesystem; readers see old or new, never partial
            return {"status": "done", "row_groups": new_groups, "rows": rows, "old_row_groups": groups,
                    "bytes_before": old_size, "bytes_after": new_size, "seconds": round(time.perf_counter() - t0, 2)}
        finally:
            if tmp.exists():
                tmp.unlink(missing_ok=True)
    finally:
        con.close()


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--data-root", type=Path, default=Path(os.environ.get("HFMD_DATA_ROOT", "/Volumes/ssd/firstratedata")))
    ap.add_argument("--asset", action="append", help="restrict to an asset directory (repeatable): stock, etf, futures, futures_contracts, crypto, index, fx")
    ap.add_argument("--timeframe", action="append", help="restrict to a timeframe (repeatable): 1min, 5min, 30min, 1hour (default: all intraday)")
    ap.add_argument("--row-group-size", type=int, default=DEFAULT_ROW_GROUP)
    ap.add_argument("--workers", type=int, default=2, help="files rewritten in parallel (default 2)")
    ap.add_argument("--threads-per-worker", type=int, default=2)
    ap.add_argument("--max-files", type=int, default=0, help="stop after N files rewritten (0 = no bound)")
    ap.add_argument("--manifest", type=Path, help="progress file (default <data-root>/state/row_groups.json)")
    ap.add_argument("--retry-failed", action="store_true", help="re-attempt files marked failed in the manifest")
    ap.add_argument("--dry-run", action="store_true", help="list what would be rewritten, touch nothing")
    ap.add_argument("--verbose", "-v", action="store_true")
    args = ap.parse_args()
    logging.basicConfig(level=logging.INFO if args.verbose or args.dry_run else logging.WARNING,
                        format="%(asctime)s %(levelname)s %(message)s", stream=sys.stdout)

    parquet_root = args.data_root / "parquet"
    if not parquet_root.is_dir():
        log.error("no parquet directory under %s", args.data_root)
        return 2
    tfs = args.timeframe or list(INTRADAY)
    bad = [t for t in tfs if t not in INTRADAY]
    if bad:
        log.error("only intraday timeframes are rewritten (%s); got %s", ", ".join(INTRADAY), ", ".join(bad))
        return 2
    manifest = Manifest(args.manifest or (args.data_root / "state" / "row_groups.json"))
    manifest.data["row_group_size"] = args.row_group_size

    files = candidate_files(parquet_root, args.asset, tfs)
    todo = []
    for f in files:
        rel = str(f.relative_to(parquet_root))
        st = manifest.status(rel)
        if st in ("done", "skipped") or (st == "failed" and not args.retry_failed):
            continue
        todo.append((f, rel))
    if args.max_files:
        todo = todo[:args.max_files]
    log.info("%d candidate files, %d to process (%s)", len(files), len(todo), "dry run" if args.dry_run else "rewrite")

    counts = {"done": 0, "skipped": 0, "failed": 0, "dry-run": 0}
    t0 = time.time()

    def job(item):
        f, rel = item
        try:
            return rel, rewrite_one(f, args.row_group_size, args.threads_per_worker, args.dry_run)
        except Exception as e:  # noqa: BLE001 — one bad file must not stop the batch
            return rel, {"status": "failed", "error": f"{e.__class__.__name__}: {str(e)[:300]}"}

    try:
        with ThreadPoolExecutor(max_workers=max(1, args.workers)) as ex:
            for rel, res in (fut.result() for fut in as_completed([ex.submit(job, it) for it in todo])):
                counts[res["status"]] = counts.get(res["status"], 0) + 1
                if res["status"] == "dry-run":
                    print(f"would rewrite {rel}: {res['row_groups']} groups → {res['target_groups']} ({res['rows']:,} rows)")
                    continue
                manifest.mark(rel, **res)
                if res["status"] == "failed":
                    log.error("FAILED %s: %s", rel, res["error"])
                else:
                    log.info("%s %s (%s)", res["status"], rel, ", ".join(f"{k}={v}" for k, v in res.items() if k not in ("status",)))
    finally:
        if not args.dry_run:
            manifest.flush()
    summary = {"processed": len(todo), **counts, "seconds": round(time.time() - t0, 1), "manifest": str(manifest.path)}
    print(json.dumps(summary))
    return 1 if counts.get("failed") else 0


if __name__ == "__main__":
    sys.exit(main())