#!/usr/bin/env python3
"""1-second candles for a list of windows (seconds_windows.csv: day,symbol,start,end in New York time), as traded.
Saved as MASSIVE_ROOT/seconds/<day>/<SYMBOL>_<HHMM>.json.gz = [[epoch ms, open, high, low, close, volume], ...].

    python massive_seconds.py --windows seconds_windows.csv [--workers 8]
"""
from __future__ import annotations

import argparse
import csv
import gzip
import json
import os
import sys
import time
import urllib.parse
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from zoneinfo import ZoneInfo

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from massive_download import ROOT, get  # noqa: E402

ET = ZoneInfo("America/New_York")


def fetch(day: str, sym: str, start: str, end: str) -> int:
    out = ROOT / "seconds" / day / f"{sym}_{start.replace(':', '')}.json.gz"
    if out.exists():
        return -2
    y, m, d = map(int, day.split("-"))
    a = int(datetime(y, m, d, int(start[:2]), int(start[3:5]), tzinfo=ET).timestamp() * 1000)
    b = int(datetime(y, m, d, int(end[:2]), int(end[3:5]), tzinfo=ET).timestamp() * 1000)
    url = f"/v2/aggs/ticker/{urllib.parse.quote(sym, safe='')}/range/1/second/{a}/{b}?adjusted=false&sort=asc&limit=50000"
    rows = []
    while url:
        js = get(url)
        if js.get("http") == 403:
            return -1
        rows += [[r["t"], r["o"], r["h"], r["l"], r["c"], r.get("v", 0)] for r in js.get("results") or []]
        url = js.get("next_url")
    out.parent.mkdir(parents=True, exist_ok=True)
    tmp = out.with_suffix(".tmp")
    with gzip.open(tmp, "wt") as fh:
        json.dump(rows, fh, separators=(",", ":"))
    os.replace(tmp, out)
    return len(rows)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--windows", default="seconds_windows.csv")
    ap.add_argument("--workers", type=int, default=8)
    a = ap.parse_args()
    jobs = [(r["day"], r["symbol"], r["start"], r["end"]) for r in csv.DictReader(open(a.windows))]
    print(f"{len(jobs)} windows", flush=True)
    t0, done, rows, denied = time.time(), 0, 0, 0
    with ThreadPoolExecutor(a.workers) as pool:
        for fu in as_completed([pool.submit(fetch, *j) for j in jobs]):
            n = fu.result()
            denied += n == -1
            rows += max(n, 0)
            done += 1
            if done % 200 == 0:
                print(f"  {done}/{len(jobs)}, {rows:,} seconds, {(time.time() - t0) / 60:.1f} min", flush=True)
    print(f"done: {done} windows, {rows:,} seconds, not allowed by the plan {denied}", flush=True)


if __name__ == "__main__":
    main()