#!/usr/bin/env python3
"""Offline synthetic data generation (T1.5).

Produces JSONL datasets by composing :mod:`tsmm.data` primitives in parallel:

* ``data/align.jsonl``  — alignment-style samples (default: many)
* ``data/sft.jsonl``    — Evol-Instruct rephrased (subset)
* ``data/eval_synth.jsonl`` — held-out eval set (fixed seed)

Usage::

    python scripts/gen_synthetic.py --n 1000 --out data/align.jsonl --seed 0 --workers 8

Each line is one JSON sample with keys: series, attributes, timestamps,
events, instruction, answer, labels, category, segments.
"""

from __future__ import annotations

import argparse
import json
import os
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path
from typing import Any

from tqdm import tqdm

from tsmm.data.gen_pipeline import _to_jsonable, gen_one_sample, write_jsonl
from tsmm.data.instruct import evolve_instruction


def _one(args):
    """Worker: generate a single sample (top-level for pickling)."""
    T, C, seed, anomaly_types, missing_rate, event_prob, do_evolve = args
    s = gen_one_sample(
        T=T, C=C, seed=seed,
        anomaly_types=anomaly_types,
        missing_rate=missing_rate,
        event_prob=event_prob,
    )
    if do_evolve:
        s["instruction"] = evolve_instruction(s["instruction"], seed=seed)
    return s


def run(
    n: int,
    out: str,
    *,
    seed: int = 0,
    workers: int = max(1, (os.cpu_count() or 2) - 1),
    T: int = 512,
    C: int = 5,
    anomaly_types=("spike", "level_shift", "variance_change", "drift"),
    missing_rate: float = 0.05,
    event_prob: float = 0.5,
    do_evolve: bool = False,
    resume: bool = False,
) -> int:
    """Stream samples to ``out`` as they are produced.

    Writes each sample to disk the instant its future completes, while keeping
    the output file in index order via a tiny reorder buffer (size <= workers).
    Peak memory is thus O(workers) instead of O(n): 500k samples never all sit
    in RAM, and a crash loses only the unwritten tail rather than everything.

    With ``resume=True`` (and ``--n`` as the *target* total), count existing
    lines in ``out`` and only generate the missing suffix, appending. Repeated
    crashes converge to the full target: each resume picks up where it stopped.
    """
    Path(out).parent.mkdir(parents=True, exist_ok=True)
    start_idx = 0
    mode = "w"
    if resume and os.path.exists(out):
        with open(out, "r", encoding="utf-8") as f:
            start_idx = sum(1 for _ in f)
        if start_idx >= n:
            print(f"[gen] {out} already has {start_idx} lines >= target {n}; nothing to do")
            return start_idx
        mode = "a"
        print(f"[gen] resume: {out} has {start_idx}/{n}, appending {n - start_idx}")
    work = [
        (T, C, seed + i, anomaly_types, missing_rate, event_prob, do_evolve)
        for i in range(start_idx, n)
    ]
    written = start_idx
    pending: dict[int, Any] = {}   # index -> sample, awaiting its turn
    next_i = start_idx              # next index that must be flushed
    with open(out, mode, encoding="utf-8") as f, \
         ProcessPoolExecutor(max_workers=workers) as ex:
        futures = {ex.submit(_one, w): i for i, w in zip(range(start_idx, n), work)}
        for fut in tqdm(as_completed(futures), total=len(work),
                        desc=os.path.basename(out), initial=start_idx):
            i = futures[fut]
            pending[i] = fut.result()
            # flush everything that is now contiguously ready, in order
            while next_i in pending:
                f.write(json.dumps(_to_jsonable(pending.pop(next_i)),
                                   ensure_ascii=False))
                f.write("\n")
                next_i += 1
                written += 1
            f.flush()
    return written


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description="Generate synthetic TS-QA JSONL data.")
    p.add_argument("--n", type=int, default=1000, help="number of samples")
    p.add_argument("--out", type=str, default="data/align.jsonl", help="output JSONL path")
    p.add_argument("--seed", type=int, default=0, help="base RNG seed")
    p.add_argument("--workers", type=int, default=max(1, (os.cpu_count() or 2) - 1))
    p.add_argument("--T", type=int, default=512, help="series length")
    p.add_argument("--C", type=int, default=5, help="number of channels")
    p.add_argument("--missing-rate", type=float, default=0.05)
    p.add_argument("--event-prob", type=float, default=0.5)
    p.add_argument("--evolve", action="store_true", help="apply Evol-Instruct rephrasing (for sft set)")
    p.add_argument("--resume", action="store_true",
                   help="count existing lines in --out and append only the missing suffix (crash-safe)")
    args = p.parse_args(argv)

    n = run(
        n=args.n,
        out=args.out,
        seed=args.seed,
        workers=args.workers,
        T=args.T,
        C=args.C,
        missing_rate=args.missing_rate,
        event_prob=args.event_prob,
        do_evolve=args.evolve,
        resume=args.resume,
    )
    print(f"wrote {n} samples -> {args.out}")
    return 0


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