"""Enrich snapshot with macro regime features: macro_vix, macro_hy_spread, macro_t10y2y from FRED.

Usage:
    uv run python3 scripts/enrich_macro_features.py \
        --input data/parquet/midwide-liquid-long-v1_tier2 \
        --output data/parquet/midwide-liquid-long-v1_tier2_macro
"""
from __future__ import annotations

import argparse
import json
import shutil
import sys
from pathlib import Path

import pyarrow as pa
import pyarrow.parquet as pq
import requests

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

ORACLE_URL = "http://localhost:18001"
FEATURES = ["macro_vix", "macro_hy_spread", "macro_t10y2y"]


def fetch_fred_series(series_id: str, start: str, end: str) -> dict[str, float]:
    """Fetch FRED observations and return {date: value} dict."""
    resp = requests.get(
        f"{ORACLE_URL}/api/v1/fred/proxy/series/observations",
        params={"series_id": series_id, "observation_start": start, "observation_end": end},
        timeout=30,
    )
    if resp.status_code != 200:
        print(f"  FRED error {resp.status_code} for {series_id}")
        return {}
    data = resp.json()
    obs = data.get("data", {}).get("observations", [])
    result = {}
    for o in obs:
        v = o.get("value")
        if v not in (None, ".", ""):
            try:
                result[o["date"]] = float(v)
            except (ValueError, TypeError):
                pass
    return result


def get_prior_value(series_dict: dict[str, float], date: str) -> float | None:
    """Get the most recent value on or before date (forward-fill)."""
    sorted_dates = sorted(series_dict.keys())
    result = None
    for d in sorted_dates:
        if d <= date:
            result = series_dict[d]
        else:
            break
    return result


def enrich_split(input_path: Path, output_path: Path, vix_dict: dict, hy_dict: dict, t10y2y_dict: dict):
    table = pq.read_table(input_path)
    n = len(table)
    existing_cols = set(table.column_names)

    if "event_date" not in existing_cols:
        print("  Empty or schema-less split, copying")
        pq.write_table(table, output_path)
        return

    # Check if already enriched (all 3 features present with <10% null)
    if all(f in existing_cols for f in FEATURES):
        null_counts = {f: table.column(f).null_count for f in FEATURES}
        if all(v < n * 0.1 for v in null_counts.values()):
            print(f"  Already enriched, copying as-is")
            pq.write_table(table, output_path)
            return

    event_dates = table.column("event_date").to_pylist()
    vix_vals = []
    hy_vals = []
    t10y2y_vals = []

    for ed in event_dates:
        date_str = str(ed)
        vix_vals.append(get_prior_value(vix_dict, date_str))
        hy_vals.append(get_prior_value(hy_dict, date_str))
        t10y2y_vals.append(get_prior_value(t10y2y_dict, date_str))

    enriched = sum(1 for v in vix_vals if v is not None)
    enriched_yc = sum(1 for v in t10y2y_vals if v is not None)
    print(f"  Enriched {enriched}/{n} rows with macro_vix/macro_hy_spread, {enriched_yc}/{n} with macro_t10y2y")

    for col_name, values in [("macro_vix", vix_vals), ("macro_hy_spread", hy_vals), ("macro_t10y2y", t10y2y_vals)]:
        arr = pa.array(values, type=pa.float64())
        if col_name in existing_cols:
            idx = table.column_names.index(col_name)
            table = table.set_column(idx, col_name, arr)
        else:
            table = table.append_column(col_name, arr)

    pq.write_table(table, output_path)
    print(f"  Written to {output_path}")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", required=True, help="Input snapshot dir")
    parser.add_argument("--output", required=True, help="Output snapshot dir")
    args = parser.parse_args()

    input_dir = Path(args.input)
    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Copy manifest
    manifest_src = input_dir / "manifest.json"
    if manifest_src.exists():
        manifest = json.loads(manifest_src.read_text())
        manifest["snapshot_id"] = output_dir.name
        (output_dir / "manifest.json").write_text(json.dumps(manifest, indent=2))

    # Determine date range across all splits
    all_dates = []
    for split in ["train", "valid", "test"]:
        p = input_dir / f"{split}.parquet"
        if p.exists():
            t = pq.read_table(p)
            if "event_date" not in t.column_names:
                continue
            t = t.select(["event_date"])
            all_dates.extend([str(d) for d in t.column("event_date").to_pylist()])

    if not all_dates:
        print("No dates found, copying as-is")
        for split in ["train", "valid", "test"]:
            p = input_dir / f"{split}.parquet"
            if p.exists():
                shutil.copy2(p, output_dir / f"{split}.parquet")
        return

    start = min(all_dates)
    end = max(all_dates)
    print(f"Fetching FRED data for {start} to {end}...")

    vix_dict = fetch_fred_series("VIXCLS", start, end)
    hy_dict = fetch_fred_series("BAMLH0A0HYM2", start, end)
    t10y2y_dict = fetch_fred_series("T10Y2Y", start, end)
    print(f"  VIX: {len(vix_dict)} obs, HY: {len(hy_dict)} obs, T10Y2Y: {len(t10y2y_dict)} obs")

    for split in ["train", "valid", "test"]:
        input_path = input_dir / f"{split}.parquet"
        if not input_path.exists():
            print(f"Skipping {split} (not found)")
            continue
        output_path = output_dir / f"{split}.parquet"
        print(f"\nProcessing {split}...")
        enrich_split(input_path, output_path, vix_dict, hy_dict, t10y2y_dict)

    print("\nDone.")


if __name__ == "__main__":
    main()