#!/usr/bin/env python3
from __future__ import annotations

import argparse
import datetime as dt
import json
from pathlib import Path
from typing import Any

from apps.backtester.run import _build_merged_snapshot_store
from libs.backtest.allocator import _macro_regime_state
from libs.backtest.manifests import load_manifest, resolve_config
from libs.backtest.overlay import (
    build_overlay_curve,
    load_equity_curve_csv,
    load_merged_store_from_snapshot_dir,
    summarize_overlay_curve,
)
from libs.common.config import get_settings


def _parse_date(value: str | None) -> dt.date | None:
    if not value:
        return None
    return dt.date.fromisoformat(value)


def _compute_regimes(
    *,
    snapshot_dir: str | Path,
    split: str,
    config_path: str | Path,
    start_date: dt.date | None,
    end_date: dt.date | None,
) -> dict[dt.date, str]:
    del split  # overlay regimes should cover the full requested window, not a single split
    manifest = load_manifest(config_path)
    config = resolve_config(manifest, config_root=".")

    raw_snapshot_dir = Path(snapshot_dir)
    if (raw_snapshot_dir / "train.parquet").exists() or (raw_snapshot_dir / "test.parquet").exists():
        settings = get_settings()
        store = load_merged_store_from_snapshot_dir(
            raw_snapshot_dir,
            oracle_url=settings.stock_oracle_url,
            db_dsn=settings.postgres_dsn,
        )
    else:
        try:
            store = _build_merged_snapshot_store(
                manifest,
                config,
                snapshot_dir_override=str(raw_snapshot_dir),
            )
        except FileNotFoundError:
            store = _build_merged_snapshot_store(
                manifest,
                config,
                snapshot_dir_override=None,
            )

    if start_date or end_date:
        lower = start_date or dt.date.min
        upper = end_date or dt.date.max
        store = store.slice_by_date_range(lower, upper)

    regimes: dict[dt.date, str] = {}
    for date in store.all_trading_days():
        regimes[date] = _macro_regime_state(config, store.get_macro_for_date(date))
    return regimes


def _load_spec(path: str | Path) -> dict[str, Any]:
    return json.loads(Path(path).read_text())


def main() -> None:
    parser = argparse.ArgumentParser(description="Evaluate a regime-switched overlay from book equity curves")
    parser.add_argument("--spec", required=True, help="Path to overlay spec JSON")
    parser.add_argument("--output-dir", required=True, help="Directory to write overlay outputs")
    args = parser.parse_args()

    spec = _load_spec(args.spec)
    initial_equity = float(spec.get("initial_equity", 10_000.0))

    curves = {
        book["label"]: load_equity_curve_csv(book["equity_csv"])
        for book in spec["books"]
    }
    regime_source = spec["regime_source"]
    start_date = _parse_date(spec.get("start_date"))
    end_date = _parse_date(spec.get("end_date"))
    regimes = _compute_regimes(
        snapshot_dir=regime_source["snapshot_dir"],
        split=regime_source.get("split", "train"),
        config_path=regime_source["config_path"],
        start_date=start_date,
        end_date=end_date,
    )

    curve = build_overlay_curve(
        curves=curves,
        allocations=spec["allocations"],
        regimes_by_date=regimes,
        initial_equity=initial_equity,
    )
    if start_date:
        curve = curve[curve["date"] >= start_date]
    if end_date:
        curve = curve[curve["date"] <= end_date]

    summary = summarize_overlay_curve(curve, initial_equity=initial_equity)
    summary["overlay_name"] = spec.get("overlay_name", Path(args.spec).stem)
    summary["books"] = [book["label"] for book in spec["books"]]
    summary["allocations"] = spec["allocations"]

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    curve.assign(date=curve["date"].astype(str)).to_csv(out_dir / "overlay_equity.csv", index=False)
    (out_dir / "overlay_summary.json").write_text(json.dumps(summary, indent=2) + "\n")
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()