#!/usr/bin/env python
"""Build training shards from the archive and stream them to the training host.

    dataset_build.py --days 60 --out /tmp/shards --push htpc@192.168.1.66:/var/home/htpc/suvi

This is a ~260,000-frame traversal of a filesystem that has repeatedly run out of file
handles under exactly this kind of load, on a VM with 184 GB free on the mount and 3.8 GB
on root against a dataset that measures 168 GB.  Both facts shape the design:

* every read goes through :class:`suvi.vfs.Reliever`, which hands handles back as it goes;
* a shard is written to `--out` (point it at tmpfs), pushed, and deleted before the next
  one starts, so peak local usage is one day rather than the whole set.

Days are chosen from the index, not the filesystem, and the split is written first.  If
extraction dies half way the manifest still describes the intended experiment, and a
resumed run fills in the shards that are missing rather than choosing different days.
"""

import argparse
import datetime
import json
import os
import subprocess
import sys
import time
from multiprocessing import Pool

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from suvi import dataset, db, fitsio, paths, vfs

#: The bench's ground-truth window.  Excluded from training so it stays a clean holdout.
HOLDOUT = ("2024-05-08", "2024-05-16")
#: GOES-18 joins the archive in 2022-08; before that there is no second satellite.
DEFAULT_FROM = "2022-09-01"
DEFAULT_TO = "2025-03-20"


def _day_bounds(day):
    start = datetime.datetime.fromisoformat(day).replace(tzinfo=datetime.timezone.utc)
    return int(start.timestamp()), int(start.timestamp()) + 86400


def holdout_days(first=HOLDOUT[0], last=HOLDOUT[1]):
    start = datetime.date.fromisoformat(first)
    end = datetime.date.fromisoformat(last)
    days = []
    while start <= end:
        days.append(start.isoformat())
        start += datetime.timedelta(days=1)
    return days


def _extract(job):
    """Worker: one slot -> encoded block.  Runs in a subprocess, so it opens its own files."""
    root, time_start, rows = job
    fitsio.quiet_astropy()
    try:
        return time_start, dataset.read_slot(root, rows)
    except (OSError, ValueError):
        return time_start, None


def build_day(conn, root, day, satellite, out_dir, workers, reliever):
    """Extract one (day, satellite) shard.  Returns (path, bytes, slots_with_frames)."""
    start, end = _day_bounds(day)
    by_slot = {}
    for row in conn.execute(
        """
        SELECT wavelength, t_start, path FROM frame
        WHERE satellite = ? AND t_start >= ? AND t_start < ?
        """,
        (satellite, start, end),
    ):
        by_slot.setdefault(row["t_start"], {})[row["wavelength"]] = row["path"]

    jobs = [(root, t, rows) for t, rows in sorted(by_slot.items())]
    records = {}
    with Pool(workers) as pool:
        for time_start, block in pool.imap_unordered(_extract, jobs, chunksize=1):
            records[time_start] = block
            reliever.tick(len(paths.WAVELENGTHS))

    path = os.path.join(out_dir, dataset.shard_name(day, satellite))
    size = dataset.write_shard(path, day, satellite, paths.WAVELENGTHS, records)
    return path, size, sum(1 for block in records.values() if block is not None)


def push(path, destination, retries=3):
    """Copy a shard to the training host and remove the local copy."""
    for attempt in range(retries):
        result = subprocess.run(
            ["rsync", "-q", "--partial", "--inplace", path, destination],
            stderr=subprocess.PIPE,
        )
        if result.returncode == 0:
            os.remove(path)
            return True
        print(f"    rsync failed (attempt {attempt + 1}): "
              f"{result.stderr.decode(errors='replace').strip()[:200]}")
        time.sleep(5)
    return False


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--days", type=int, default=60)
    parser.add_argument("--out", default="/tmp/shards",
                        help="staging directory; put this on tmpfs")
    parser.add_argument("--push", default=None,
                        help="rsync destination, e.g. host:/path (omit to keep locally)")
    parser.add_argument("--from", dest="from_day", default=DEFAULT_FROM)
    parser.add_argument("--to", dest="to_day", default=DEFAULT_TO)
    parser.add_argument("--satellites", default="16,18")
    parser.add_argument("--workers", type=int, default=6)
    parser.add_argument("--db", default=None)
    parser.add_argument("--manifest", default=None,
                        help="reuse an existing manifest instead of choosing days afresh")
    args = parser.parse_args(argv)

    os.makedirs(args.out, exist_ok=True)
    satellites = tuple(int(s) for s in args.satellites.split(","))
    conn = db.connect(args.db, readonly=True)
    root = paths.data_root()

    manifest_path = args.manifest or os.path.join(args.out, "manifest.json")
    if args.manifest and os.path.exists(args.manifest):
        manifest = json.load(open(args.manifest))
        splits = manifest["splits"]
        print(f"Reusing manifest {args.manifest} (digest {manifest['digest']})")
    else:
        chosen = dataset.choose_days(
            conn, args.days,
            *[_day_bounds(d)[0] for d in (args.from_day, args.to_day)],
            exclude=holdout_days(), satellites=satellites,
        )
        if len(chosen) < args.days:
            print(f"WARNING: only {len(chosen)} complete days available, wanted {args.days}")
        splits = dataset.split_days(chosen)
        manifest = dataset.manifest(splits, {"satellites": list(satellites),
                                             "holdout": list(HOLDOUT)})
        with open(manifest_path, "w") as handle:
            json.dump(manifest, handle, indent=2)
        print(f"Wrote {manifest_path} (digest {manifest['digest']})")

    adjacent = dataset.check_split(splits)
    for name, group in sorted(splits.items()):
        print(f"  {name:>5}: {len(group)} days  {group[0] if group else '-'}"
              f" .. {group[-1] if group else '-'}")
    if adjacent:
        print(f"  NOTE: {len(adjacent)} split boundaries fall on consecutive days; "
              "frames four minutes apart sit either side of them.")

    if args.push:
        subprocess.run(["rsync", "-q", manifest_path, args.push], check=False)

    days = sorted(sum(splits.values(), []))
    total_bytes = 0
    total_slots = 0
    started = time.time()
    reliever = vfs.Reliever(label="dataset")
    try:
        for index, day in enumerate(days, 1):
            for satellite in satellites:
                name = dataset.shard_name(day, satellite)
                path, size, slots = build_day(
                    conn, root, day, satellite, args.out, args.workers, reliever
                )
                total_bytes += size
                total_slots += slots
                elapsed = time.time() - started
                done = (index - 1) * len(satellites) + satellites.index(satellite) + 1
                remaining = (len(days) * len(satellites) - done) * elapsed / max(done, 1)
                print(f"  [{done:>3}/{len(days) * len(satellites)}] {name}: "
                      f"{slots} slots, {size / 1e6:.0f} MB, "
                      f"{total_bytes / 1e9:.1f} GB total, "
                      f"eta {remaining / 3600:.1f} h", flush=True)
                if args.push and not push(path, args.push):
                    raise SystemExit(f"could not push {name}; stopping rather than "
                                     "filling the staging disk")
    finally:
        reliever.finish()

    print(f"\n{total_slots} slots, {total_bytes / 1e9:.1f} GB, "
          f"{(time.time() - started) / 3600:.2f} h")
    return 0


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