#!/usr/bin/env python
"""Phase 2c: streaming export -> session features at 1m/5m/30m/1d -> rolling windows, per symbol (multiprocess).
Raw Parquet is deleted after the feature rows are written; manifests keep hashes and dataset versions."""
from __future__ import annotations

import argparse
import datetime as dt
import json
import os
import sys
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path

import polars as pl

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from wp39.api_client import FinDataClient  # noqa: E402
from wp39.corporate_actions import load_actions  # noqa: E402
from wp39.features import compute_spy_buckets, compute_symbol_features, session_bounds_table  # noqa: E402
from wp39.ingest import export_symbol, mark_consumed  # noqa: E402
from wp39.pipeline import Paths, get_logger, load_config  # noqa: E402
from wp39.provenance import append_jsonl, now_iso, write_json  # noqa: E402
from wp39.samples import rolling_window_features  # noqa: E402

START, END = dt.date(2011, 1, 1), dt.date(2026, 10, 30)


def work(symbol: str, keep_raw: bool) -> dict:
    P = Paths.default()
    contract = load_config("data_contract")
    ds = contract["findata"]["datasets"]["bars_1m"]
    out = P.cache / "features" / f"{symbol}.parquet"
    win_out = P.cache / "windows" / f"{symbol}.parquet"
    win_out.parent.mkdir(exist_ok=True)
    if out.exists() and win_out.exists():
        return {"symbol": symbol, "status": "cached"}
    t0 = time.time()
    client = FinDataClient(log_path=P.logs / f"api_calls_worker_{os.getpid()}.jsonl")
    try:
        paths = export_symbol(client, symbol, START, END, ds, P.cache / "bars", P.manifests / "bars_manifest.jsonl", keep=keep_raw)
    finally:
        client.close()
    if not paths:
        return {"symbol": symbol, "status": "no_export"}
    bounds = pl.read_parquet(P.cache / "external" / "session_bounds.parquet")
    spy = {k: pl.read_parquet(P.cache / "external" / f"spy_buckets_{k}.parquet") for k in ("1m", "5m", "30m")}
    actions = pl.read_parquet(P.cache / "external" / "actions.parquet")
    feat = compute_symbol_features(paths, symbol, bounds, spy, actions)
    if feat.height == 0:
        return {"symbol": symbol, "status": "empty"}
    feat.write_parquet(out)
    rolling_window_features(feat).write_parquet(win_out)
    if not keep_raw:
        for p in paths:
            p.unlink(missing_ok=True)
        mark_consumed(P.manifests / "bars_manifest.jsonl", symbol, ds)
    return {"symbol": symbol, "status": "ok", "sessions": feat.height, "first": str(feat["session_date"].min()),
            "last": str(feat["session_date"].max()), "secs": round(time.time() - t0, 1)}


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--workers", type=int, default=6)
    ap.add_argument("--limit", type=int, default=0, help="pilot: only the first N eligible symbols")
    ap.add_argument("--symbols", type=str, default="", help="comma-separated explicit list")
    ap.add_argument("--keep-raw", action="store_true")
    args = ap.parse_args()
    P = Paths.default()
    log = get_logger("03_features", P.logs / "03_features.log")
    contract = load_config("data_contract")
    ext = P.cache / "external"
    # shared inputs: session bounds, SPY buckets, corporate actions
    if not (ext / "session_bounds.parquet").exists():
        session_bounds_table(START, END).write_parquet(ext / "session_bounds.parquet")
    if not (ext / "spy_buckets_1m.parquet").exists():
        client = FinDataClient(log_path=P.logs / "api_calls.jsonl")
        paths = export_symbol(client, "SPY", START, END, contract["findata"]["datasets"]["bars_1m_etf"], P.cache / "bars",
                              P.manifests / "bars_manifest.jsonl", keep=True)
        client.close()
        bounds = pl.read_parquet(ext / "session_bounds.parquet")
        for k, v in compute_spy_buckets(paths, bounds).items():
            v.write_parquet(ext / f"spy_buckets_{k}.parquet")
        log.info("SPY buckets written")
    if not (ext / "actions.parquet").exists():
        load_actions(ext / "corporate_actions_raw.csv").write_parquet(ext / "actions.parquet")
    if args.symbols:
        symbols = args.symbols.split(",")
    else:
        uni = pl.read_parquet(P.processed / "universe.parquet").filter(pl.col("eligible")).sort("ticker")
        symbols = uni["ticker"].to_list()
        if args.limit:
            symbols = symbols[: args.limit]
    log.info("symbols to process: %d (workers=%d)", len(symbols), args.workers)
    status = P.logs / "03_features_status.jsonl"
    done = {r["symbol"] for r in (json.loads(l) for l in open(status)) if r.get("status") in ("ok", "cached", "empty", "no_export")} if status.exists() else set()
    todo = [s for s in symbols if s not in done]
    t0 = time.time(); n_ok = 0
    with ProcessPoolExecutor(max_workers=args.workers) as ex:
        futs = {ex.submit(work, s, args.keep_raw): s for s in todo}
        for i, f in enumerate(as_completed(futs)):
            try:
                r = f.result()
            except Exception as exc:  # noqa: BLE001
                r = {"symbol": futs[f], "status": "error", "error": str(exc)[:300]}
            r["ts"] = now_iso()
            append_jsonl(status, r)
            n_ok += r["status"] in ("ok", "cached")
            if i % 50 == 0:
                log.info("%d/%d done (ok=%d) elapsed=%.0fs last=%s", i + 1, len(todo), n_ok, time.time() - t0, r)
    log.info("finished: %d symbols, %.0fs", len(todo), time.time() - t0)
    return 0


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