"""tac-qlib R&D MCP server (stdio transport).

Exposes the quant research workflow as MCP tools backed by tac-qlib + upstream qlib:

    prepare data  -> tac-engine lake tools (tradeac-lake skill; not here)
    preprocess    -> rd_dataset   (TACHandler + processor pipeline inspection)
    train         -> rd_train     (LightGBM on the lake, logged to mlflow)
    test          -> rd_evaluate  (IC / ICIR / Rank IC / Rank ICIR)
    predict       -> rd_predict   (persist pred.pkl / label.pkl)
    backtest      -> rd_backtest  (TopkDropoutStrategy + cost model + risk analysis)
    one-shot      -> rd_run_workflow (execute a qrun workflow yaml — in a background
                      thread, or in a separate process via run_in_new_process=true)

Run::

    .venv/bin/python -m tac_qlib.rd_server        # stdio MCP server

All tools return JSON-safe dicts. Logging goes to stderr; stdout is reserved for
the MCP protocol.
"""

from __future__ import annotations

import datetime as _dt
import json
import os
import pickle
import re
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

import numpy as np
import pandas as pd

from mcp.server.mcpserver import MCPServer

from tac_qlib.risk_limits import parse_limits

# Load the repo-root `.env` so the server reads the same env file whether
# spawned by opencode MCP, tac-app, or a shell. Existing process env (e.g.
# Coolify or `node --env-file`) wins — load_dotenv does not override
# already-set variables, mirroring tac-engine's dotenvy load.
#
# The module may resolve from the repo source (`tac-qlib/tac_qlib/`, e.g. when
# tac-app sets PYTHONPATH=<root>/tac-qlib) or from a wheel copy in site-packages
# (e.g. the Docker image's non-editable `uv pip install ./tac-qlib`, spawned via
# opencode). Locate the repo root by walking up until a `.env` file is found —
# fall back to the cwd so the container (no .env on disk, vars injected by
# Coolify) still picks up a local file if one is mounted.
from dotenv import load_dotenv

def _load_repo_env() -> None:
    # opencode (or a wrapper) may inject `{env:VAR}` substitutions as empty
    # strings for unset variables; drop those so the repo `.env` (or a real
    # injected value) can apply below.
    for _k in ("TAC_LAKE_DIR", "DATABASE_URL", "MLRUNS_URI"):
        if os.environ.get(_k) == "":
            os.environ.pop(_k, None)
    here = Path(__file__).resolve()
    for parent in here.parents:
        if (parent / ".env").is_file():
            load_dotenv(parent / ".env", override=False)
            return
    load_dotenv(Path.cwd() / ".env", override=False)

_load_repo_env()

server = MCPServer(
    name="tac-qlib-rd",
    title="TradeAC quant R&D",
    instructions=(
        "Tools for the quant R&D loop on the TradeAC parquet lake: dataset "
        "(preprocess), train, predict, evaluate, backtest, and one-shot qrun of "
        "a workflow yaml. Data preparation (bar/feature backfill) is done with the "
        "tac-engine lake tools (tradeac-lake skill)."
    ),
    version="0.1.0",
)

_INIT_STATE: Dict[str, Any] = {}


# --------------------------------------------------------------------------- utils
def _default_lake() -> str:
    """The lake root from ``TAC_LAKE_DIR`` (mandatory — no fallback)."""
    lake = os.environ.get("TAC_LAKE_DIR")
    if not lake:
        raise RuntimeError(
            "TAC_LAKE_DIR is not set. Point it at the TradeAC lake root, e.g. "
            "export TAC_LAKE_DIR=/home/data/lake (docker) or set an absolute "
            "path in your local .env."
        )
    return lake


def _default_mlruns_uri() -> str:
    """MLflow tracking URI: Postgres when ``$DATABASE_URL`` is set (no fallback).

    With ``$DATABASE_URL`` set the tracking store lives in the shared Postgres
    (MLflow's own ``experiments``/``runs``/... tables; artifacts stay on the
    filesystem under ``<lake>/mlruns``). The trace FK references Postgres
    ``runs(run_uuid)``, so the store MUST be Postgres whenever ``DATABASE_URL``
    exists — never a silent sqlite fallback. Only when ``DATABASE_URL`` is
    unset do we fall back to the lake sqlite store (with a loud warning).
    ``MLRUNS_URI`` overrides.
    """
    from tac_qlib.data.config import resolve_lake_root

    url = (os.environ.get("DATABASE_URL") or "").strip()
    if url.startswith("postgres://"):
        return "postgresql+psycopg://" + url[len("postgres://") :]
    if url.startswith("postgresql://"):
        return "postgresql+psycopg://" + url[len("postgresql://") :]
    sqlite_uri = f"sqlite:///{resolve_lake_root(None)}/mlruns.db"
    print(
        "[rd_server] WARNING: DATABASE_URL not set — using the lake sqlite tracking store "
        f"({sqlite_uri}). Runs will NOT be in Postgres; traced experiments (FK to "
        "`runs(run_uuid)`) will fail to finalize. Set DATABASE_URL in the environment.",
        file=sys.stderr,
    )
    return sqlite_uri


def _tracking_sqlite_path(uri: str):
    """Extract the sqlite db path from an MLflow tracking uri (None for non-sqlite)."""
    uri = (uri or "").strip()
    if uri.startswith("sqlite:///"):
        return Path(uri[len("sqlite:///") :]).expanduser()
    if uri.startswith("sqlite://"):
        return Path(uri[len("sqlite://") :]).expanduser()
    return None


def _set_experiment_artifact_location(uri: str, exp_id: str, art_dir: str) -> bool:
    """Force an experiment's ``artifact_location`` to ``art_dir``.

    MLflow exposes no public API to change an experiment's artifact location
    after creation. Since runs derive their ``artifact_uri`` from it at
    ``create_run`` time (``append_to_uri_path(exp.artifact_location, run_id,
    "artifacts")``), a bare root like ``<lake>/mlruns`` makes runs land in
    ``mlruns/<run_id>/`` instead of ``mlruns/<exp_id>/<run_id>/``. For both the
    sqlite and postgres tracking stores the ``experiments`` table is
    authoritative, so patch it directly before any run is created.
    """
    if uri.startswith(("postgres://", "postgresql://", "postgresql+psycopg://")):
        try:
            import psycopg

            with psycopg.connect(_postgres_dsn(uri)) as conn, conn.cursor() as cur:
                cur.execute(
                    "UPDATE experiments SET artifact_location = %s WHERE experiment_id = %s",
                    (art_dir, str(exp_id)),
                )
            return True
        except Exception:  # noqa: BLE001 - best-effort; runs still work, just not co-located
            return False
    db = _tracking_sqlite_path(uri)
    if db is None or not db.exists():
        return False
    try:
        import sqlite3

        conn = sqlite3.connect(db)
        try:
            conn.execute(
                "UPDATE experiments SET artifact_location=? WHERE experiment_id=?",
                (art_dir, str(exp_id)),
            )
            conn.commit()
        finally:
            conn.close()
        return True
    except Exception:  # noqa: BLE001 - best-effort; runs still work, just not co-located
        return False


def _postgres_dsn(uri: str) -> str:
    """Normalise a postgres tracking uri to a psycopg/libpq conninfo string.

    MLflow URIs may be ``postgresql+psycopg://``, ``postgresql://`` or
    ``postgres://`` (the latter is not valid libpq); strip the SQLAlchemy
    driver suffix and normalise ``postgres://`` -> ``postgresql://``.
    """
    for prefix in ("postgresql+psycopg://", "postgresql://"):
        if uri.startswith(prefix):
            return "postgresql://" + uri[len(prefix) :]
    if uri.startswith("postgres://"):
        return "postgresql://" + uri[len("postgres://") :]
    return uri


def _pin_artifact_location(uri: str, exp_name: str) -> None:
    """Create (or fix) the experiment with artifacts pinned under the lake ``mlruns/``.

    MLflow's file artifact store defaults to ``./mlruns`` relative to the process
    cwd, so a run written by a server launched from another directory (e.g. the
    repo root) splits the DB (lake) from its artifacts (cwd). Pre-creating the
    experiment with an explicit ``artifact_location`` under the lake keeps both
    under ``$TAC_LAKE_DIR``.

    The location must include the experiment id (``<lake>/mlruns/<exp_id>``),
    otherwise runs land in ``<lake>/mlruns/<run_id>/`` and the experiment id is
    dropped from the on-disk layout. The id is only known after creation, so we
    create with a default location and then patch the sqlite ``experiments`` row
    to the id-suffixed path. Existing experiments (e.g. created before this
    function existed, or by an older pin) get their location corrected too.
    """
    try:
        from mlflow.tracking import MlflowClient

        from tac_qlib.data.config import resolve_lake_root

        art_root = str((resolve_lake_root(None) / "mlruns").resolve())
        client = MlflowClient(tracking_uri=uri)
        exp = client.get_experiment_by_name(exp_name)
        if exp is None:
            exp_id = client.create_experiment(exp_name)
        else:
            exp_id = exp.experiment_id
        _set_experiment_artifact_location(uri, exp_id, str(Path(art_root) / exp_id))
    except Exception:  # noqa: BLE001 - best-effort; runs still work, just not co-located
        pass


def _jsonable(obj: Any) -> Any:
    """Recursively convert numpy/pandas/datetime values to JSON-safe python types."""
    if isinstance(obj, dict):
        return {k: _jsonable(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple, set)):
        return [_jsonable(v) for v in obj]
    if isinstance(obj, np.generic):
        return obj.item()
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    if isinstance(obj, (pd.Timestamp, _dt.datetime, _dt.date, _dt.time)):
        return obj.isoformat()
    if isinstance(obj, float):
        return obj if not np.isnan(obj) and not np.isinf(obj) else None
    if obj is None or isinstance(obj, (str, int, bool)):
        return obj
    return str(obj)


def _resolve_universe(universe: Optional[str], lake_root: str, market: str) -> List[str]:
    if universe in (None, "", "all", "ALL"):
        from tac_qlib.data.config import LakeConfig

        return LakeConfig(lake_root, market).load_symbols()
    return [s.strip().upper() for s in str(universe).replace(" ", "").split(",") if s]


def _default_dates(lake_root: str, market: str):
    from tac_qlib.data.config import LakeConfig

    dates = LakeConfig(lake_root, market).load_calendar_dates()
    if not dates:
        raise ValueError(
            "lake calendar is empty; seed it first with tac-engine `backfill_lake_calendar`"
        )
    return dates[0], dates[-1]


def _ensure_init(lake_root: Optional[str], market: str) -> str:
    lake_root = str(Path(lake_root or _default_lake()).expanduser().resolve())
    market = (market or "US").upper()
    if _INIT_STATE.get("lake_root") == lake_root and _INIT_STATE.get("market") == market:
        return lake_root
    from tac_qlib.qlib_init import qlib_init

    uri = os.environ.get("MLRUNS_URI") or _default_mlruns_uri()
    exp_manager = {
        "class": "MLflowExpManager",
        "module_path": "qlib.workflow.expm",
        "kwargs": {"uri": uri, "default_exp_name": "tac-rd"},
    }
    qlib_init(provider_uri=lake_root, market=market, exp_manager=exp_manager)
    _INIT_STATE.update(lake_root=lake_root, market=market)
    return lake_root


def _feature_fields(features: Optional[str], freq: str, lake_root: str, market: str) -> List[str]:
    from tac_qlib.contrib.data.handler import TACHandler

    return TACHandler._normalize_feature_fields(features or None, freq, lake_root, market)


def _build_dataset(
    lake_root: str,
    market: str,
    universe: List[str],
    start_time: str,
    end_time: str,
    fit_start_time: str,
    fit_end_time: str,
    freq: str,
    features: Optional[str],
    label: Optional[str],
    train: Optional[tuple],
    valid: Optional[tuple],
    test: Optional[tuple],
):
    from qlib.data.dataset import DatasetH
    from tac_qlib.contrib.data.handler import TACHandler

    handler = TACHandler(
        instruments=universe,
        start_time=start_time,
        end_time=end_time,
        freq=freq,
        fit_start_time=fit_start_time,
        fit_end_time=fit_end_time,
        feature_fields=features,
        label=label,
        lake_root=lake_root,
        market=market,
    )
    segments = {}
    if train is not None:
        segments["train"] = train
    if valid is not None:
        segments["valid"] = valid
    if test is not None:
        segments["test"] = test
    dataset = DatasetH(handler=handler, segments=segments) if segments else None
    return dataset, handler


def _seg_stats(dataset, seg: str) -> dict:
    from qlib.data.dataset.handler import DataHandlerLP

    df = dataset.prepare(seg, col_set="feature", data_key=DataHandlerLP.DK_I)
    dates = df.index.get_level_values(0)
    insts = sorted(df.index.get_level_values(1).dropna().astype(str).unique())
    return {
        "rows": int(len(df)),
        "date_min": str(pd.Timestamp(dates.min()).date()),
        "date_max": str(pd.Timestamp(dates.max()).date()),
        "instruments": insts,
    }


def _nan_ratios(lake_root: str, market: str, universe: List[str], fields: List[str],
                start_time: str, end_time: str, freq: str) -> dict:
    from qlib.data import D

    out = {}
    for f in fields:
        df = D.features(universe, [f], start_time, end_time, freq)
        n = int(len(df))
        nn = int(df.iloc[:, 0].notna().sum())
        out[f] = {"rows": n, "non_nan": nn, "nan_ratio": round((n - nn) / n, 4) if n else None}
    return out


def _load_model(model_path: Optional[str], run_id: Optional[str], experiment_name: str):
    if model_path:
        with open(model_path, "rb") as fh:
            return pickle.load(fh)
    if run_id:
        return _load_artifact(experiment_name, run_id, "params.pkl")
    raise ValueError("provide either model_path or run_id")


def _load_artifact(experiment_name: str, run_id: str, name: str):
    from qlib.workflow import R

    exp = R.get_exp(experiment_name=experiment_name)
    rec = exp.get_recorder(recorder_id=run_id)
    if rec is None:
        raise ValueError(f"no recorder {run_id!r} in experiment {experiment_name!r}")
    return rec.load_object(name)


def _head_rows(pred: pd.DataFrame, top: int) -> List[dict]:
    rows = []
    for (ts, inst), v in list(pred.iloc[:, 0].items())[:top]:
        rows.append({"date": str(pd.Timestamp(ts).date()), "instrument": str(inst),
                     "score": round(float(v), 6)})
    return rows


# --------------------------------------------------------------------------- tools
@server.tool(structured_output=False)
def rd_status(lake_root: str = "", market: str = "US") -> dict:
    """Inspect the TradeAC lake: root, calendar range, symbols, per-symbol coverage windows and available ta-lib feature columns. Call this first to confirm the lake is populated before running dataset / train / predict / backtest tools."""
    lake_root = lake_root or _default_lake()
    from tac_qlib.data.config import LakeConfig

    cfg = LakeConfig(lake_root, market)
    root = str(cfg.lake_root)
    cal = cfg.load_calendar_dates()
    symbols = cfg.load_symbols()

    coverage = []
    cp = cfg.coverage_path()
    if cp.exists():
        df = pd.read_parquet(cp)
        if len(df):
            cols = [c for c in ("market", "timeframe", "symbol", "feed", "first_t", "last_t") if c in df.columns]
            for _, r in df[cols].sort_values(["symbol", "timeframe"]).iterrows():
                coverage.append(_jsonable(r.to_dict()))

    features = []
    from tac_qlib.contrib.data.handler import get_common_feature_fields

    for timeframe in ("1d", "1m", "1h"):
        try:
            fl = get_common_feature_fields(lake_root, market, timeframe)
        except Exception:
            fl = []
        if fl:
            features.append({"timeframe": timeframe, "columns": fl})

    return {
        "lake_root": root,
        "market": market.upper(),
        "calendar_days": len(cal),
        "calendar_start": str(pd.Timestamp(cal[0]).date()) if cal else None,
        "calendar_end": str(pd.Timestamp(cal[-1]).date()) if cal else None,
        "symbols": symbols,
        "symbol_count": len(symbols),
        "coverage": coverage,
        "features": features,
    }


@server.tool(structured_output=False)
def rd_dataset(
    lake_root: str = "",
    market: str = "US",
    universe: str = "all",
    start_time: str = "",
    end_time: str = "",
    fit_start_time: str = "",
    fit_end_time: str = "",
    freq: str = "day",
    features: str = "",
    label: str = "",
    train_start: str = "",
    train_end: str = "",
    valid_start: str = "",
    valid_end: str = "",
    test_start: str = "",
    test_end: str = "",
) -> dict:
    """Build (but do not train) the dataset for the given universe / window / segments and report the preprocessing outcome: final feature columns, per-segment row counts + date ranges, and raw NaN coverage per feature over the fit window (columns that are 100% NaN get dropped by DropAllNaN)."""
    lake_root = _ensure_init(lake_root, market)
    market = (market or "US").upper()
    universe_l = _resolve_universe(universe, lake_root, market)
    if not universe_l:
        raise ValueError("empty universe: the lake has no symbols; backfill bars first")

    if not start_time or not end_time:
        first, last = _default_dates(lake_root, market)
        start_time = start_time or str(pd.Timestamp(first).date())
        end_time = end_time or str(pd.Timestamp(last).date())
    fit_start_time = fit_start_time or start_time
    fit_end_time = fit_end_time or end_time

    fields = _feature_fields(features, freq, lake_root, market)
    label_expr = label or "Ref($close,-2)/Ref($close,-1)-1"

    train = (train_start, train_end) if train_start and train_end else None
    valid = (valid_start, valid_end) if valid_start and valid_end else None
    test = (test_start, test_end) if test_start and test_end else None

    dataset, handler = _build_dataset(
        lake_root, market, universe_l, start_time, end_time, fit_start_time, fit_end_time,
        freq, features or None, label_expr, train, valid, test,
    )

    segments = {}
    if dataset is not None:
        for name in ("train", "valid", "test"):
            if name in dataset.segments:
                segments[name] = _seg_stats(dataset, name)

    ratios = _nan_ratios(lake_root, market, universe_l, fields, fit_start_time, fit_end_time, freq)
    all_nan = [f for f, r in ratios.items() if r["non_nan"] == 0]
    nan_ratio = {f: r["nan_ratio"] for f, r in ratios.items() if r["nan_ratio"] and r["nan_ratio"] > 0}

    return {
        "lake_root": lake_root,
        "market": market,
        "freq": freq,
        "universe": universe_l,
        "universe_size": len(universe_l),
        "window": {"start": start_time, "end": end_time, "fit_start": fit_start_time, "fit_end": fit_end_time},
        "feature_fields": fields,
        "label": label_expr,
        "segments": segments,
        "all_nan_dropped": all_nan,
        "nan_ratio_nonzero": nan_ratio,
    }


@server.tool(structured_output=False)
def rd_train(
    wait: bool = True,
    lake_root: str = "",
    market: str = "US",
    universe: str = "all",
    train_start: str = "",
    train_end: str = "",
    valid_start: str = "",
    valid_end: str = "",
    test_start: str = "",
    test_end: str = "",
    freq: str = "day",
    features: str = "",
    label: str = "",
    learning_rate: float = 0.05,
    num_leaves: int = 15,
    num_boost_round: int = 200,
    early_stopping_rounds: int = 50,
    colsample_bytree: float = 0.8,
    subsample: float = 0.8,
    subsample_freq: int = 1,
    reg_alpha: float = 0.01,
    reg_lambda: float = 0.01,
    seed: int = 2026,
    experiment_name: str = "tac-rd",
    out_dir: str = "tac-rd-output",
    record_analysis: bool = True,
    benchmark: str = "",
    topk: int = 2,
    n_drop: int = 1,
    account: int = 1000000,
    risk_degree: float = 0.95,
    open_cost: float = 0.0005,
    close_cost: float = 0.0015,
    min_cost: float = 5.0,
    risk_limits: str = "",
) -> dict:
    """Fit a LightGBM (LGBModel, mse) on the lake dataset and log it to mlflow. When ``record_analysis`` is set (default), also runs SignalRecord / SigAnaRecord (ana_long_short) / PortAnaRecord on the test segment so the run records pred/label, IC/Rank IC series, long-short group returns and the portfolio backtest. The model is also pickled to <out_dir>/params.pkl so later steps can reuse it via model_path (or via the returned run_id). With ``wait=False`` the training runs in the background and the tool returns immediately — poll rd_exp_get_run / rd_exp_list for the newest run of ``experiment_name`` (status RUNNING until it finishes) and reuse its run_id. risk_limits: JSON spec recorded as run tag ``rd.risk_limits``; the run's PortAnaRecord backtest folds size/concentration caps into risk_degree (liquidity floor is applied by rd_backtest on the recorded pred)."""
    kwargs = dict(
        lake_root=lake_root, market=market, universe=universe,
        train_start=train_start, train_end=train_end,
        valid_start=valid_start, valid_end=valid_end,
        test_start=test_start, test_end=test_end,
        freq=freq, features=features, label=label,
        learning_rate=learning_rate, num_leaves=num_leaves,
        num_boost_round=num_boost_round, early_stopping_rounds=early_stopping_rounds,
        colsample_bytree=colsample_bytree, subsample=subsample,
        subsample_freq=subsample_freq, reg_alpha=reg_alpha, reg_lambda=reg_lambda,
        seed=seed, experiment_name=experiment_name, out_dir=out_dir,
        record_analysis=record_analysis, benchmark=benchmark, topk=topk, n_drop=n_drop,
        account=account, risk_degree=risk_degree, open_cost=open_cost,
        close_cost=close_cost, min_cost=min_cost, risk_limits=risk_limits,
    )
    if wait:
        return _rd_train_impl(**kwargs)
    threading.Thread(target=_rd_train_impl, kwargs=kwargs, daemon=True).start()
    return {
        "status": "started",
        "background": True,
        "experiment_name": experiment_name,
        "out_dir": out_dir,
        "note": ("training running in the background — poll rd_exp_get_run / rd_exp_list for the "
                 "newest run of this experiment (status RUNNING until it finishes), reuse its run_id, "
                 "then use the model via run_id or the params.pkl written to out_dir"),
    }


def _rd_train_impl(
    lake_root: str = "",
    market: str = "US",
    universe: str = "all",
    train_start: str = "",
    train_end: str = "",
    valid_start: str = "",
    valid_end: str = "",
    test_start: str = "",
    test_end: str = "",
    freq: str = "day",
    features: str = "",
    label: str = "",
    learning_rate: float = 0.05,
    num_leaves: int = 15,
    num_boost_round: int = 200,
    early_stopping_rounds: int = 50,
    colsample_bytree: float = 0.8,
    subsample: float = 0.8,
    subsample_freq: int = 1,
    reg_alpha: float = 0.01,
    reg_lambda: float = 0.01,
    seed: int = 2026,
    experiment_name: str = "tac-rd",
    out_dir: str = "tac-rd-output",
    record_analysis: bool = True,
    benchmark: str = "",
    topk: int = 2,
    n_drop: int = 1,
    account: int = 1000000,
    risk_degree: float = 0.95,
    open_cost: float = 0.0005,
    close_cost: float = 0.0015,
    min_cost: float = 5.0,
    risk_limits: str = "",
) -> dict:
    lake_root = _ensure_init(lake_root, market)
    market = (market or "US").upper()
    universe_l = _resolve_universe(universe, lake_root, market)
    if not (train_start and train_end):
        raise ValueError("train_start and train_end are required")
    test_start = test_start or train_end
    test_end = test_end or str(pd.Timestamp(pd.Timestamp(train_end) + pd.Timedelta(days=30)).date())

    first, _last = _default_dates(lake_root, market)
    start_time = min(str(pd.Timestamp(first).date()), train_start)
    end_time = test_end
    fit_start_time, fit_end_time = train_start, train_end

    label_expr = label or "Ref($close,-2)/Ref($close,-1)-1"
    train = (train_start, train_end)
    valid = (valid_start, valid_end) if valid_start and valid_end else None
    test = (test_start, test_end)

    dataset, handler = _build_dataset(
        lake_root, market, universe_l, start_time, end_time, fit_start_time, fit_end_time,
        freq, features or None, label_expr, train, valid, test,
    )

    from qlib.contrib.model.gbdt import LGBModel
    from qlib.workflow import R

    model = LGBModel(
        loss="mse",
        learning_rate=learning_rate,
        num_leaves=num_leaves,
        colsample_bytree=colsample_bytree,
        subsample=subsample,
        subsample_freq=subsample_freq,
        reg_alpha=reg_alpha,
        reg_lambda=reg_lambda,
        seed=seed,
    )

    feature_fields = _feature_fields(features, freq, lake_root, market)
    segments_cfg = {}
    for name, seg in (("train", train), ("valid", valid), ("test", test)):
        if seg:
            segments_cfg[name] = list(seg)
    config = {
        "qlib_init": {
            "provider_uri": lake_root,
            "region": market.lower(),
            "exp_manager": {"class": "MLflowExpManager", "module_path": "qlib.workflow.expm",
                            "kwargs": {"uri": os.environ.get("MLRUNS_URI") or _default_mlruns_uri(),
                                       "default_exp_name": experiment_name}},
        },
        "task": {
            "model": {"class": "LGBModel", "module_path": "qlib.contrib.model.gbdt", "kwargs": {
                "loss": "mse", "learning_rate": learning_rate, "num_leaves": num_leaves,
                "n_estimators": num_boost_round, "colsample_bytree": colsample_bytree,
                "subsample": subsample, "subsample_freq": subsample_freq,
                "reg_alpha": reg_alpha, "reg_lambda": reg_lambda, "seed": seed}},
            "dataset": {"class": "DatasetH", "module_path": "qlib.data.dataset", "kwargs": {
                "handler": {"class": "TACHandler", "module_path": "tac_qlib.contrib.data.handler", "kwargs": {
                    "instruments": universe_l, "start_time": start_time, "end_time": end_time,
                    "fit_start_time": fit_start_time, "fit_end_time": fit_end_time, "freq": freq,
                    "feature_fields": feature_fields, "label": label_expr,
                    "lake_root": lake_root, "market": market}},
                "segments": segments_cfg}},
            "record": [],
        },
    }

    t0 = time.time()
    record_warns: List[str] = []
    _pin_artifact_location(os.environ.get("MLRUNS_URI") or _default_mlruns_uri(), experiment_name)
    limits = parse_limits(risk_limits)
    eff_degree = risk_degree
    if limits:
        if limits.get("size_cap_pct"):
            eff_degree = min(eff_degree, limits["size_cap_pct"] * topk / 100.0)
        if limits.get("concentration_cap_pct"):
            eff_degree = min(eff_degree, limits["concentration_cap_pct"] / 100.0)
    with R.start(experiment_name=experiment_name):
        model.fit(dataset, num_boost_round=num_boost_round,
                  early_stopping_rounds=early_stopping_rounds, verbose_eval=0)
        R.save_objects(**{"params.pkl": model})
        R.save_objects(config=config)
        rec = R.get_recorder()
        try:
            R.set_tags(rd_experiment_id=str(rec.info.get("experiment_id", "")),
                       rd_run_id=str(rec.info.get("id", "")))
            if limits:
                R.set_tags(**{"rd.risk_limits": json.dumps(limits, sort_keys=True),
                              "rd.risk_degree_effective": str(eff_degree)})
        except Exception:  # noqa: BLE001 - tags are best-effort
            pass
        if record_analysis:
            try:
                from qlib.workflow.record_temp import PortAnaRecord, SigAnaRecord, SignalRecord

                SignalRecord(model=model, dataset=dataset, recorder=rec).generate()
                SigAnaRecord(recorder=rec, ana_long_short=True, ann_scaler=252).generate()
                port_config = {
                    "strategy": {
                        "class": "TopkDropoutStrategy",
                        "module_path": "qlib.contrib.strategy",
                        "kwargs": {"signal": "<PRED>", "topk": topk, "n_drop": n_drop,
                                   "only_tradable": True, "risk_degree": eff_degree},
                    },
                    "backtest": {
                        "start_time": test_start,
                        "end_time": test_end,
                        "account": account,
                        "exchange_kwargs": {"codes": universe_l, "deal_price": "$close", "freq": freq,
                                            "open_cost": open_cost, "close_cost": close_cost, "min_cost": min_cost},
                    },
                }
                if benchmark:
                    port_config["backtest"]["benchmark"] = benchmark
                PortAnaRecord(recorder=rec, config=port_config, risk_analysis_freq="1d").generate()
            except Exception as exc:  # noqa: BLE001 - analysis is best-effort; training already succeeded
                record_warns.append(f"analysis records failed: {exc}")
    fit_sec = round(time.time() - t0, 2)
    info = dict(rec.info)

    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    with open(out / "params.pkl", "wb") as fh:
        pickle.dump(model, fh)

    return {
        "experiment_name": experiment_name,
        "run_id": info.get("id"),
        "status": info.get("status"),
        "fit_seconds": fit_sec,
        "feature_fields": feature_fields,
        "label": label_expr,
        "segments": {s: _seg_stats(dataset, s) for s in ("train", "valid", "test") if s in dataset.segments},
        "model_path": str(out / "params.pkl"),
        "out_dir": str(out),
        "record_warnings": record_warns,
    }


@server.tool(structured_output=False)
def rd_predict(
    model_path: str = "",
    run_id: str = "",
    experiment_name: str = "tac-rd",
    lake_root: str = "",
    market: str = "US",
    universe: str = "all",
    train_start: str = "",
    train_end: str = "",
    valid_start: str = "",
    valid_end: str = "",
    test_start: str = "",
    test_end: str = "",
    freq: str = "day",
    features: str = "",
    label: str = "",
    out_dir: str = "tac-rd-output",
    top: int = 5,
) -> dict:
    """Predict on the test segment with a fitted model (from model_path pickle or an mlflow run_id) and persist pred.pkl / label.pkl to out_dir. Returns the score head and summary stats."""
    lake_root = _ensure_init(lake_root, market)
    market = (market or "US").upper()
    universe_l = _resolve_universe(universe, lake_root, market)
    if not (test_start and test_end):
        raise ValueError("test_start and test_end are required")
    if not (train_start and train_end):
        first, _ = _default_dates(lake_root, market)
        train_start = str(pd.Timestamp(first).date())
        train_end = str((pd.Timestamp(test_start) - pd.Timedelta(days=1)).date())

    first, _last = _default_dates(lake_root, market)
    start_time = min(str(pd.Timestamp(first).date()), train_start)
    end_time = test_end
    fit_start_time, fit_end_time = train_start, train_end

    label_expr = label or "Ref($close,-2)/Ref($close,-1)-1"
    train = (train_start, train_end)
    valid = (valid_start, valid_end) if valid_start and valid_end else None
    test = (test_start, test_end)

    dataset, _handler = _build_dataset(
        lake_root, market, universe_l, start_time, end_time, fit_start_time, fit_end_time,
        freq, features or None, label_expr, train, valid, test,
    )

    model = _load_model(model_path or None, run_id or None, experiment_name)
    pred = model.predict(dataset, segment="test")
    pred_df = pred.to_frame("score")

    from qlib.data.dataset.handler import DataHandlerLP

    label_df = dataset.prepare("test", col_set="label", data_key=DataHandlerLP.DK_I)

    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    pred_df.to_pickle(out / "pred.pkl")
    label_df.to_pickle(out / "label.pkl")

    s = pred_df.iloc[:, 0]
    return {
        "pred_path": str(out / "pred.pkl"),
        "label_path": str(out / "label.pkl"),
        "count": int(len(s)),
        "date_min": str(pd.Timestamp(s.index.get_level_values(0).min()).date()),
        "date_max": str(pd.Timestamp(s.index.get_level_values(0).max()).date()),
        "instruments": sorted(s.index.get_level_values(1).dropna().astype(str).unique()),
        "score_mean": round(float(s.mean()), 6),
        "score_std": round(float(s.std()), 6),
        "score_min": round(float(s.min()), 6),
        "score_max": round(float(s.max()), 6),
        "head": _head_rows(pred_df, int(top)),
    }


@server.tool(structured_output=False)
def rd_evaluate(
    pred_path: str = "",
    label_path: str = "",
    run_id: str = "",
    experiment_name: str = "tac-rd",
    ann_scaler: int = 252,
) -> dict:
    """Evaluate a saved prediction against the label with per-day IC / Rank IC (qlib calc_ic): mean, std, annual vol, IR, skew, kurtosis, max drawdown, plus the headline IC / ICIR / Rank IC / Rank ICIR metrics that SigAnaRecord logs. Loads pred.pkl / label.pkl from paths or from an mlflow run."""
    if pred_path:
        pred = pd.read_pickle(pred_path)
    elif run_id:
        pred = _load_artifact(experiment_name, run_id, "pred.pkl")
    else:
        raise ValueError("provide either pred_path or run_id")

    if label_path:
        label = pd.read_pickle(label_path)
    elif run_id:
        label = _load_artifact(experiment_name, run_id, "label.pkl")
    else:
        raise ValueError("a label is required to compute IC; pass label_path or run_id")

    from qlib.contrib.eva.alpha import calc_ic

    ic, ric = calc_ic(pred.iloc[:, 0], label.iloc[:, 0], dropna=True)

    def _stats(s: pd.Series) -> dict:
        s = s.dropna()
        mean = float(s.mean())
        std = float(s.std(ddof=1))
        return {
            "days": int(len(s)),
            "mean": round(mean, 6),
            "std": round(std, 6),
            "ann_vol": round(std * float(np.sqrt(ann_scaler)), 6),
            "ir": round(mean / std * float(np.sqrt(ann_scaler)), 6) if std > 0 else None,
            "skew": round(float(s.skew()), 6),
            "kurt": round(float(s.kurt()), 6),
            "maxdd": round(float((s.cumsum() - s.cumsum().cummax()).min()), 6),
        }

    ic_std, ric_std = float(ic.std(ddof=1)), float(ric.std(ddof=1))
    return {
        "IC": _stats(ic),
        "RankIC": _stats(ric),
        "headline": {
            "IC": round(float(ic.mean()), 6),
            "ICIR": round(float(ic.mean()) / ic_std, 6) if ic_std > 0 else None,
            "Rank IC": round(float(ric.mean()), 6),
            "Rank ICIR": round(float(ric.mean()) / ric_std, 6) if ric_std > 0 else None,
        },
    }


@server.tool(structured_output=False)
def rd_backtest(
    pred_path: str,
    start_time: str,
    end_time: str,
    lake_root: str = "",
    market: str = "US",
    topk: int = 2,
    n_drop: int = 1,
    risk_degree: float = 0.95,
    benchmark: str = "",
    account: float = 1000000.0,
    universe: str = "",
    open_cost: float = 0.0005,
    close_cost: float = 0.0015,
    min_cost: float = 5.0,
    out_dir: str = "tac-rd-output",
    risk_limits: str = "",
) -> dict:
    """Run a daily backtest of the saved prediction with TopkDropoutStrategy and a cost model, then compute risk analysis (mean, std, annualized return, information ratio, max drawdown). Writes report/positions/risk to out_dir. benchmark: any symbol the lake holds, or empty for none. risk_limits: JSON (liquidity_floor_adv, size_cap_pct, concentration_cap_pct) applied as a signal filter so the same spec that gates live also gates backtest."""
    lake_root = _ensure_init(lake_root, market)
    pred = pd.read_pickle(pred_path)
    if isinstance(pred, pd.DataFrame):
        col = "score" if "score" in pred.columns else pred.columns[0]
        pred = pred[col]
    pred = pred.astype(float).sort_index()
    insts = sorted(pred.index.get_level_values(1).dropna().astype(str).unique())
    codes = [s.strip().upper() for s in str(universe).replace(" ", "").split(",") if s] or insts

    from qlib.contrib.evaluate import backtest_daily, risk_analysis
    from qlib.contrib.strategy.signal_strategy import TopkDropoutStrategy

    signal = pred
    limits_applied = {}
    limits = parse_limits(risk_limits)
    if limits:
        from tac_qlib.risk_limits import dollar_adv

        end_ts = pd.Timestamp(end_time)
        adv = dollar_adv(codes, lake_root, market, asof=end_ts.isoformat())
        floor = limits.get("liquidity_floor_adv")
        if floor:
            illiquid = sorted(s for s in insts if adv.get(str(s).upper(), 0.0) < floor)
            if illiquid:
                bad = {str(s).upper() for s in illiquid}
                inst_lev = signal.index.get_level_values(1).astype(str).str.upper()
                keep = ~inst_lev.isin(bad)
                signal = signal[keep]
                limits_applied["dropped_liquidity"] = sorted(bad)
                limits_applied["note"] = f"liquidity floor ${floor:,.0f} ADV dropped {len(bad)} symbols"

        # size / concentration caps fold into effective risk_degree (equal-weight TopkDropout)
        eff_degree = risk_degree
        if limits.get("size_cap_pct"):
            eff_degree = min(eff_degree, limits["size_cap_pct"] * topk / 100.0)
            limits_applied["size_cap_risk_degree"] = round(eff_degree, 4)
        if limits.get("concentration_cap_pct"):
            eff_degree = min(eff_degree, limits["concentration_cap_pct"] / 100.0)
            limits_applied["concentration_cap_risk_degree"] = round(eff_degree, 4)
        if abs(eff_degree - risk_degree) > 1e-9:
            limits_applied["risk_degree_from"] = risk_degree
            limits_applied["risk_degree"] = eff_degree

    strategy = TopkDropoutStrategy(
        signal=signal, topk=topk, n_drop=n_drop, only_tradable=True,
        risk_degree=eff_degree if limits else risk_degree,
    )
    report, positions = backtest_daily(
        start_time=start_time,
        end_time=end_time,
        strategy=strategy,
        account=account,
        benchmark=benchmark or None,
        exchange_kwargs={
            "codes": codes,
            "deal_price": "$close",
            "freq": "day",
            "open_cost": open_cost,
            "close_cost": close_cost,
            "min_cost": min_cost,
        },
    )

    risk = risk_analysis(report["return"], freq="day")

    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    report.to_csv(out / "report_normal.csv")
    pd.DataFrame({ts: pos.get_stock_amount_dict() for ts, pos in positions.items()}).T.to_csv(
        out / "positions_normal.csv"
    )
    risk.to_csv(out / "risk.csv")

    return {
        "start_time": start_time,
        "end_time": end_time,
        "trading_days": int(len(report)),
        "risk": {str(k): round(float(v), 6) for k, v in risk.iloc[:, 0].items()},
        "benchmark": benchmark or None,
        "topk": topk,
        "n_drop": n_drop,
        "risk_limits_applied": limits_applied,
        "artifacts": [str(out / n) for n in ("report_normal.csv", "positions_normal.csv", "risk.csv")],
    }


@server.tool(structured_output=False)
def rd_risk_calibrate(
    pred_path: str,
    start_time: str,
    end_time: str,
    lake_root: str = "",
    market: str = "US",
    topk: int = 2,
    n_drop: int = 1,
    risk_degree: float = 0.95,
    benchmark: str = "",
    account: float = 1000000.0,
    universe: str = "",
    open_cost: float = 0.0005,
    close_cost: float = 0.0015,
    min_cost: float = 5.0,
    risk_limits: str = "",
    out_dir: str = "tac-rd-calibrate",
) -> dict:
    """Prove the risk-limit spec: run limit-vs-no-limit A/B plus a sensitivity grid over each limit dimension (size_cap_pct, concentration_cap_pct, liquidity_floor_adv), and compare Sharpe / Calmar / max drawdown / ann. return. Pass the round's `risk_limits` spec as the A/B candidate. Returns rows ordered by Sharpe and the raw per-run risk dicts; use it (B2) to pick the spec that keeps live ≈ backtest."""
    lake_root = _ensure_init(lake_root, market)
    pred = pd.read_pickle(pred_path)
    if isinstance(pred, pd.DataFrame):
        col = "score" if "score" in pred.columns else pred.columns[0]
        pred = pred[col]
    pred = pred.astype(float).sort_index()
    insts = sorted(pred.index.get_level_values(1).dropna().astype(str).unique())
    codes = [s.strip().upper() for s in str(universe).replace(" ", "").split(",") if s] or insts

    from qlib.contrib.evaluate import backtest_daily, risk_analysis
    from qlib.contrib.strategy.signal_strategy import TopkDropoutStrategy

    from tac_qlib.risk_limits import dollar_adv

    adv = dollar_adv(codes, lake_root, market, asof=str(pd.Timestamp(end_time).date()))

    def _run(spec: dict) -> dict:
        signal = pred
        applied = {}
        floor = spec.get("liquidity_floor_adv")
        if floor:
            bad = {str(s).upper() for s in insts if adv.get(str(s).upper(), 0.0) < floor}
            if bad:
                inst_lev = signal.index.get_level_values(1).astype(str).str.upper()
                signal = signal[~inst_lev.isin(bad)]
                applied["dropped_liquidity"] = sorted(bad)
        eff = risk_degree
        if spec.get("size_cap_pct"):
            eff = min(eff, spec["size_cap_pct"] * topk / 100.0)
        if spec.get("concentration_cap_pct"):
            eff = min(eff, spec["concentration_cap_pct"] / 100.0)
        strategy = TopkDropoutStrategy(
            signal=signal, topk=topk, n_drop=n_drop, only_tradable=True, risk_degree=eff,
        )
        report, _positions = backtest_daily(
            start_time=start_time, end_time=end_time, strategy=strategy, account=account,
            benchmark=benchmark or None,
            exchange_kwargs={
                "codes": codes, "deal_price": "$close", "freq": "day",
                "open_cost": open_cost, "close_cost": close_cost, "min_cost": min_cost,
            },
        )
        risk = risk_analysis(report["return"], freq="day")
        r = {str(k): float(v) for k, v in risk.iloc[:, 0].items()}
        return {"risk": r, "applied": applied}

    def _sharpe(r: dict) -> float:
        return r.get("information_ratio") or r.get("sharpe_of_daily_return") or 0.0

    cand = parse_limits(risk_limits)
    rows = []
    runs = {}

    base = _run({})
    runs["baseline"] = base
    rows.append({"label": "baseline (no limits)", **{k: round(v, 6) for k, v in base["risk"].items()}})

    if cand:
        ab = _run(cand)
        runs["candidate"] = ab
        rows.append({"label": f"candidate {json.dumps(cand)}", **{k: round(v, 6) for k, v in ab["risk"].items()}})

        grid = []
        if cand.get("size_cap_pct"):
            for pct in (5, 10, 15, 20):
                spec = dict(cand, size_cap_pct=float(pct))
                grid.append((f"size_cap {pct}%", spec))
        if cand.get("concentration_cap_pct"):
            for pct in (20, 30, 40, 50):
                spec = dict(cand, concentration_cap_pct=float(pct))
                grid.append((f"conc {pct}%", spec))
        if cand.get("liquidity_floor_adv"):
            for floor in (1e6, 2.5e6, 5e6, 1e7):
                spec = dict(cand, liquidity_floor_adv=float(floor))
                grid.append((f"liquidity ${floor:,.0f}", spec))
        for label, spec in grid:
            r = _run(spec)
            runs[label] = r
            rows.append({"label": label, **{k: round(v, 6) for k, v in r["risk"].items()}})

    rows.sort(key=lambda row: -_sharpe(row))
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    with open(out / "risk_calibration.json", "w") as fh:
        json.dump({"rows": rows, "runs": runs, "candidate": cand}, fh, indent=2, default=str)

    headline = {k: _sharpe(r["risk"]) for k, r in runs.items()}
    best = max(headline, key=headline.get) if headline else None
    return {
        "start_time": start_time,
        "end_time": end_time,
        "trading_days": None,
        "best_by_sharpe": best,
        "sharpes": {k: round(v, 4) for k, v in headline.items()},
        "comparison": rows,
        "artifacts": [str(out / "risk_calibration.json")],
    }


@server.tool(structured_output=False)
def rd_strategy_targets(
    pred_path: str,
    signal_date: str = "",
    account: float = 1000000.0,
    risk_degree: float = 0.95,
    topk: int = 2,
    n_drop: int = 1,
    prices: str = "",
    risk_limits: str = "",
    equity: float = 0.0,
    peak_equity: float = 0.0,
) -> dict:
    """Turn a saved prediction into a deterministic target order list for one signal date, applying the same TopkDropout selection the backtest uses: rank the cross-sectional scores, drop the top n_drop, take the next topk as buys, sized at account * risk_degree / topk per name. With `prices` (JSON {symbol: price}) also floors to whole shares (qty) and reports expected_price. risk_limits: JSON (liquidity_floor_adv, size_cap_pct, concentration_cap_pct, drawdown_pause_pct) applied as pre-gates — same spec rd_backtest uses. equity/peak_equity enable the drawdown_pause gate (equity <= (1-pause_pct) * peak -> no buys). Chains straight into tac-engine place_order — no manual strategy replication. signal_date defaults to the last day in the prediction."""
    pred = pd.read_pickle(pred_path)
    if isinstance(pred, pd.DataFrame):
        col = "score" if "score" in pred.columns else pred.columns[0]
        scores_all = pred[col]
    else:
        scores_all = pred
    scores_all = scores_all.sort_index()
    dates = sorted(pd.Index(scores_all.index.get_level_values(0)).unique())
    if not dates:
        raise ValueError("prediction is empty — no signal dates")

    target_date = pd.Timestamp(signal_date) if signal_date else dates[-1]
    if target_date in dates:
        resolved = target_date
        note = None
    else:
        past = [d for d in dates if d <= target_date]
        resolved = past[-1] if past else dates[0]
        note = f"requested {target_date.date()} not in prediction; using {pd.Timestamp(resolved).date()}"

    day = scores_all.loc[resolved]
    day = day[day.notna()]

    ranking = day.sort_values(ascending=False)

    # ---- risk limits: drawdown pause gate, liquidity filter, sizing caps ----
    limits = parse_limits(risk_limits)
    limits_applied = {}
    if limits:
        from tac_qlib.risk_limits import apply_to_ranking, dollar_adv, drawdown_pause

        paused, pause_reason = drawdown_pause(equity, peak_equity, limits)
        if paused:
            return {
                "signal_date": str(pd.Timestamp(resolved).date()),
                "strategy": "TopkDropout",
                "reason": f"drawdown_pause: {pause_reason} — no targets",
                "risk_limits_applied": {"drawdown_pause_pct": limits["drawdown_pause_pct"], "paused": True},
                "targets": [],
                "ranking": [
                    {"symbol": s, "rank": i + 1, "score": round(float(v), 6)}
                    for i, (s, v) in enumerate(ranking.head(20).items())
                ],
            }

        floor = limits.get("liquidity_floor_adv")
        if floor:
            insts = [str(s) for s in ranking.index]
            adv = dollar_adv(insts, _default_lake(), "US", asof=str(pd.Timestamp(resolved).date()), lookback=20)
            ranking, applied = apply_to_ranking(ranking, adv, limits, account, risk_degree, topk)
        else:
            ranking, applied = apply_to_ranking(ranking, {}, limits, account, risk_degree, topk)
        limits_applied.update(applied)

    order = ranking.nlargest(topk + n_drop)
    if len(order) < topk + n_drop:
        return {
            "signal_date": str(pd.Timestamp(resolved).date()),
            "strategy": "TopkDropout",
            "reason": f"only {len(order)} names with a score on {pd.Timestamp(resolved).date()} (need {topk + n_drop}) — no targets",
            "targets": [],
            "risk_limits_applied": limits_applied or None,
            "ranking": [
                {"symbol": s, "rank": i + 1, "score": round(float(v), 6)}
                for i, (s, v) in enumerate(ranking.head(20).items())
            ],
        }
    targets = order.iloc[n_drop:]

    per_name_notional = round(limits_applied.get("per_name_notional", account * risk_degree / topk), 2)
    price_map = {}
    if prices:
        import json

        price_map = json.loads(prices)
    out_targets = []
    for symbol, score in targets.items():
        item = {
            "symbol": str(symbol),
            "rank": int(ranking.index.get_loc(symbol)) + 1,
            "score": round(float(score), 6),
            "side": "buy",
            "notional": per_name_notional,
        }
        price = price_map.get(str(symbol))
        if price:
            qty = int(per_name_notional // float(price))
            item["qty"] = qty
            item["expected_price"] = float(price)
            item["invested"] = round(qty * float(price), 2)
        out_targets.append(item)

    return {
        "signal_date": str(pd.Timestamp(resolved).date()),
        "date_note": note,
        "strategy": "TopkDropout",
        "topk": topk,
        "n_drop": n_drop,
        "risk_degree": risk_degree,
        "account": round(account, 2),
        "per_name_notional": per_name_notional,
        "risk_limits_applied": limits_applied or None,
        "count": len(out_targets),
        "targets": out_targets,
        "ranking": [
            {"symbol": s, "rank": i + 1, "score": round(float(v), 6)}
            for i, (s, v) in enumerate(ranking.head(20).items())
        ],
    }


def _spawn_subprocess_workflow(config_path: str, experiment_name: str, uri_folder: str = "") -> dict:
    """Run a workflow in a fresh process, redirecting its stdout/stderr to a log file.

    ``qlib.init`` sets process-global ``C``/``R`` singletons, so concurrent
    in-process background workflow threads would race. A subprocess gets a clean
    qlib, survives crashes, and releases memory when it exits. Because this
    server speaks MCP over stdio, the child MUST NOT inherit this process's
    stdout — everything it prints is diverted to ``<lake>/logs/rd-workflow-<exp>-<ts>.log``.

    The child writes to the same mlflow tracking store (Postgres or the lake
    sqlite), so the parent's ``rd_exp_*`` pollers see the run as it progresses.
    """
    from tac_qlib.data.config import resolve_lake_root

    cp = Path(config_path).expanduser().resolve()
    if not cp.exists():
        raise ValueError(f"config not found: {cp}")

    lake = resolve_lake_root(None)
    log_dir = lake / "logs"
    log_dir.mkdir(parents=True, exist_ok=True)
    safe = re.sub(r"[^A-Za-z0-9_.-]+", "_", experiment_name or "workflow")[:64]
    log_path = log_dir / f"rd-workflow-{safe}-{_dt.datetime.now():%Y%m%d-%H%M%S}.log"

    cmd = [sys.executable, "-m", "tac_qlib.rd_server", "--child-workflow", str(cp), experiment_name]
    if uri_folder:
        cmd.append(uri_folder)

    with open(log_path, "ab") as logf:
        proc = subprocess.Popen(
            cmd,
            stdout=logf,
            stderr=subprocess.STDOUT,
            env=os.environ.copy(),
        )
    return {
        "status": "started",
        "background": True,
        "process_mode": "subprocess",
        "pid": proc.pid,
        "config_path": str(cp),
        "experiment_name": experiment_name,
        "log_path": str(log_path),
        "note": ("workflow running in a separate process — poll rd_exp_get_run / rd_exp_list for the "
                 "newest run of this experiment (status RUNNING until it finishes), then reuse its run_id. "
                 f"Child logs: {log_path}"),
    }


@server.tool(structured_output=False)
def rd_run_workflow(config_path: str, experiment_name: str = "workflow", uri_folder: str = "", wait: bool = False, run_in_new_process: bool = False) -> dict:
    """Execute a full qrun workflow yaml in one shot (render jinja -> init qlib from its qlib_init -> train model -> run all records: SignalRecord, SigAnaRecord, PortAnaRecord). Returns run info, mlflow metrics and saved artifact list. Equivalent to: qrun <config> --experiment_name <name>.

    A 4-year training run can outlive the MCP call timeout, so `wait` defaults to False: the workflow runs in a background thread and the tool returns immediately with `status: started`. Poll rd_exp_get_run / rd_exp_list for the newest run of `experiment_name` (status RUNNING until it finishes), then reuse its run_id. Set wait=True only for small windows that finish within the call timeout.

    `run_in_new_process=true` runs the workflow in a separate OS process (stdout/stderr redirected to <lake>/logs/rd-workflow-<exp>-<ts>.log, returned as `log_path`). qlib init sets process-global state, so use this for concurrent or long workflows — it isolates crashes and releases memory, and the same rd_exp_* polling works because the child writes to the same mlflow store."""
    if run_in_new_process:
        return _spawn_subprocess_workflow(config_path, experiment_name, uri_folder)
    kwargs = {"config_path": config_path, "experiment_name": experiment_name, "uri_folder": uri_folder}
    if wait:
        return _rd_run_workflow_impl(**kwargs)
    threading.Thread(target=_rd_run_workflow_impl, kwargs=kwargs, daemon=True).start()
    return {
        "status": "started",
        "background": True,
        "config_path": config_path,
        "experiment_name": experiment_name,
        "note": ("workflow running in the background — poll rd_exp_get_run / rd_exp_list for the "
                 "newest run of this experiment (status RUNNING until it finishes), then reuse its run_id"),
    }


def _rd_run_workflow_impl(config_path: str, experiment_name: str = "workflow", uri_folder: str = "") -> dict:
    cp = Path(config_path).expanduser()
    if not cp.exists():
        raise ValueError(f"config not found: {cp}")

    import qlib
    from qlib.cli.run import render_template, sys_config
    from qlib.config import C
    from qlib.model.trainer import task_train
    from ruamel.yaml import YAML

    # Make the resolved lake root visible to workflow YAML jinja (`{{ LAKE }}`,
    # set from `TAC_LAKE_DIR`): qlib's render_template only substitutes env vars,
    # so pin TAC_LAKE_DIR to the same lake the server resolves, regardless of
    # the process cwd.
    from tac_qlib.data.config import resolve_lake_root

    os.environ.setdefault("TAC_LAKE_DIR", str(resolve_lake_root(None)))

    rendered = render_template(str(cp))
    cfg = YAML(typ="safe", pure=True).load(rendered)
    sys_config(cfg, str(cp))

    qlib_init_cfg = cfg.get("qlib_init", {}) or {}
    from tac_qlib.data.config import resolve_lake_root

    lake_root = resolve_lake_root(None)
    unified = _default_mlruns_uri()
    use_postgres = unified.startswith("postgresql+psycopg://")
    if "exp_manager" in qlib_init_cfg:
        # DATABASE_URL set -> the unified store is Postgres and the trace FK
        # references its `runs` table: never let a workflow pin sqlite/file.
        # Without DATABASE_URL, normalize a relative `sqlite:///mlruns.db` (and
        # the lake's own `sqlite:///{{ LAKE }}/mlruns.db`, rendered to
        # `sqlite:////<lake>/mlruns.db`) to the unified store; other absolute /
        # `file:` / `http(s)` uris are left untouched.
        em_kwargs = qlib_init_cfg["exp_manager"].setdefault("kwargs", {})
        if use_postgres:
            em_kwargs["uri"] = unified
        else:
            em_uri = (em_kwargs.get("uri") or "").strip()
            lake_mlruns_sqlite = f"sqlite:///{lake_root}/mlruns.db"
            if em_uri.startswith("sqlite:///") and (
                not em_uri.startswith("sqlite:////") or em_uri.startswith(lake_mlruns_sqlite)
            ):
                em_kwargs["uri"] = unified
        qlib.init(**qlib_init_cfg)
        exp_manager = qlib_init_cfg["exp_manager"]
    else:
        exp_manager = C["exp_manager"]
        if use_postgres:
            exp_manager["kwargs"]["uri"] = unified
        else:
            uri_dir = Path(uri_folder).expanduser() if uri_folder else lake_root / "mlruns"
            exp_manager["kwargs"]["uri"] = "file:" + str(uri_dir.resolve())
        qlib.init(**qlib_init_cfg, exp_manager=exp_manager)

    task = cfg.get("task", {}) or {}
    _pin_artifact_location(exp_manager["kwargs"]["uri"], experiment_name)
    recorder = task_train(task, experiment_name=experiment_name)
    recorder.save_objects(config=cfg)
    try:
        recorder.set_tags(**{
            "rd.experiment_id": str(getattr(recorder, "experiment_id", "")),
            "rd.run_id": str(recorder.id),
        })
    except Exception:  # noqa: BLE001 - tags are best-effort
        pass

    info = dict(recorder.info)
    metrics = {k: (float(v) if not isinstance(v, dict) else None) for k, v in recorder.list_metrics().items()}

    artifacts = []
    try:
        artifacts = [fi.path for fi in recorder.client.list_artifacts(recorder.id)]
    except Exception:
        pass

    return {
        "config_path": str(cp),
        "experiment_name": experiment_name,
        "run_id": info.get("id"),
        "status": info.get("status"),
        "qlib_init": _jsonable(qlib_init_cfg),
        "model": _jsonable(task.get("model")),
        "dataset": _jsonable(task.get("dataset")),
        "records": _jsonable(task.get("record")),
        "metrics": metrics,
        "artifacts": artifacts,
    }


@server.tool(structured_output=False)
def rd_factor_attribution(
    start_time: str,
    end_time: str,
    portfolio_equity: str = "",
    benchmark: str = "SPY",
    lake_root: str = "",
    market: str = "US",
    risk_free: float = 0.0,
    ann_scaler: int = 252,
    realized_slippage_bps: float = 0.0,
    realized_cost_bps: float = 0.0,
    expected_slippage_bps: float = 0.0,
    expected_cost_bps: float = 0.0,
    drift_warn_bps: float = 50.0,
) -> dict:
    """Factor-based P&L attribution for a live window: decompose the account's return into beta (benchmark exposure) and alpha (residual edge), then flag drift when realized execution cost (from round_metrics) exceeds the backtest assumption (from the calibration).

    Pass the account equity curve from tac-engine `get_portfolio_history` as `portfolio_equity` — a JSON array of `{"t":<iso date>,"e":<equity>}` (or `{"date":..., "equity":...}`). Benchmark daily bars are read from the lake. Realized execution drag comes from `round_metrics` (slippage_bps + cost_pct_of_gross); the expected values are the backtest/calibration assumption. `drift_alarm` fires when realized − expected > drift_warn_bps."""
    lake_root = _ensure_init(lake_root, market)

    df = None
    if portfolio_equity:
        rows = json.loads(portfolio_equity)
        rec = []
        for r in rows:
            t = r.get("t") or r.get("date")
            e = r.get("e") or r.get("equity")
            if t is not None and e is not None:
                rec.append({"t": pd.Timestamp(t), "e": float(e)})
        if rec:
            df = pd.DataFrame(rec).assign(t=lambda x: pd.to_datetime(x["t"]).dt.normalize()).set_index("t").sort_index()
            df = df[~df.index.duplicated(keep="last")]
            df = df[(df.index >= pd.Timestamp(start_time).normalize()) & (df.index <= pd.Timestamp(end_time).normalize())]

    from tac_qlib.data.config import LakeConfig

    cfg = LakeConfig(lake_root, market)
    bench_sym = str(benchmark or "SPY").upper()
    bench_p = cfg.bar_path("1d", bench_sym)
    if not bench_p.exists():
        for candidate in ("IVV", "QQQ", "VOO", "SPY"):
            p = cfg.bar_path("1d", candidate)
            if p.exists():
                bench_sym = candidate
                bench_p = p
                break
    bench = pd.read_parquet(bench_p)
    bcol = bench["t"] if "t" in bench.columns else bench["date"]
    bench = bench.assign(_t=pd.to_datetime(bcol).dt.normalize()).set_index("_t").sort_index()
    bench = bench[(bench.index >= pd.Timestamp(start_time).normalize()) & (bench.index <= pd.Timestamp(end_time).normalize())]
    bench_ret = bench["c"].astype(float).pct_change().dropna()

    if df is None or len(df) < 2:
        return {
            "start_time": start_time,
            "end_time": end_time,
            "benchmark": bench_sym,
            "benchmark_note": f"requested {benchmark or 'SPY'}, used {bench_sym} (lake)",
            "error": "portfolio_equity missing or has <2 points — pass the get_portfolio_history equity curve as portfolio_equity",
            "benchmark_return_pct": round(float(bench_ret.sum() * 100), 4) if len(bench_ret) else None,
        }

    port_ret = df["e"].pct_change().dropna()
    joined = pd.concat([port_ret.rename("port"), bench_ret.rename("bench")], axis=1, join="inner").dropna()
    if len(joined) < 3:
        return {
            "start_time": start_time,
            "end_time": end_time,
            "benchmark": bench_sym,
            "benchmark_note": f"requested {benchmark or 'SPY'}, used {bench_sym} (lake)",
            "error": f"only {len(joined)} overlapping return days with benchmark — need >=3",
            "benchmark_return_pct": round(float(bench_ret.sum() * 100), 4) if len(bench_ret) else None,
        }

    p = joined["port"]
    b = joined["bench"]
    beta = float(np.cov(p, b)[0, 1] / np.var(b))
    ann_port = (1 + p.mean()) ** ann_scaler - 1
    ann_bench = (1 + b.mean()) ** ann_scaler - 1
    alpha_ann = ann_port - risk_free - beta * (ann_bench - risk_free)

    equity0 = float(df["e"].iloc[0])
    equity1 = float(df["e"].iloc[-1])
    pnl = equity1 - equity0
    pnl_beta = beta * (equity1 / equity0 - 1) * equity0
    pnl_alpha = pnl - pnl_beta

    realized_bps = None
    drift_alarm = None
    realized = realized_slippage_bps + realized_cost_bps
    expected = expected_slippage_bps + expected_cost_bps
    if realized > 0 or expected > 0:
        realized_bps = round(realized, 2)
        if realized > expected + drift_warn_bps:
            drift_alarm = (
                f"drift: realized cost {realized:.1f} bps exceeds backtest assumption "
                f"{expected:.1f} bps by {realized - expected:.1f} bps (> {drift_warn_bps:g} bps warn) — "
                f"re-check sizing / calibration (rd_risk_calibrate)"
            )
        elif realized > expected:
            drift_alarm = (
                f"watch: realized cost {realized:.1f} bps vs backtest assumption "
                f"{expected:.1f} bps (within {drift_warn_bps:g} bps warn threshold)"
            )

    return {
        "start_time": start_time,
        "end_time": end_time,
        "benchmark": bench_sym,
        "benchmark_note": f"requested {benchmark or 'SPY'}, used {bench_sym} (lake)",
        "days": len(joined),
        "equity_start": round(equity0, 2),
        "equity_end": round(equity1, 2),
        "pnl_total": round(pnl, 2),
        "return_pct": round(float((equity1 / equity0 - 1) * 100), 4),
        "benchmark_return_pct": round(float(bench_ret.sum() * 100), 4),
        "beta": round(beta, 4),
        "alpha_annualized_pct": round(float(alpha_ann * 100), 4),
        "attribution": {
            "pnl_beta": round(pnl_beta, 2),
            "pnl_alpha": round(pnl_alpha, 2),
            "alpha_share_pct": round(float(pnl_alpha / pnl * 100), 2) if abs(pnl) > 1e-9 else None,
        },
        "realized_cost_bps": realized_bps,
        "drift_alarm": drift_alarm,
        "correlation": round(float(np.corrcoef(p, b)[0, 1]), 4),
    }


# --------------------------------------------------------------------------- entry
def main() -> None:
    import argparse
    import logging

    from tac_qlib.rd_explain import register_tools
    from tac_qlib.trace import register_trace_tools

    parser = argparse.ArgumentParser(prog="tac_qlib.rd_server", description="TradeAC quant R&D MCP server (stdio)")
    parser.add_argument(
        "--child-workflow",
        nargs="+",
        default=None,
        help=argparse.SUPPRESS,  # internal: run one workflow and exit (spawned by rd_run_workflow run_in_new_process)
    )
    args, _ = parser.parse_known_args()

    if args.child_workflow:
        logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
        logging.getLogger("lightgbm").setLevel(logging.WARNING)
        logging.getLogger("mlflow").setLevel(logging.WARNING)
        os.environ.setdefault("MLFLOW_ALLOW_FILE_STORE", "true")
        cfg = args.child_workflow[0]
        name = args.child_workflow[1] if len(args.child_workflow) > 1 else "workflow"
        uri = args.child_workflow[2] if len(args.child_workflow) > 2 else ""
        print(f"[rd_server] child-workflow: {cfg} experiment={name} uri_folder={uri!r}", flush=True)
        try:
            _rd_run_workflow_impl(config_path=cfg, experiment_name=name, uri_folder=uri)
        except Exception:  # noqa: BLE001 - log the traceback (stderr -> the parent's log file) and fail loudly
            import traceback

            traceback.print_exc()
            sys.exit(1)
        return

    register_tools(server)
    register_trace_tools(server)
    logging.getLogger("qlib").setLevel(logging.WARNING)
    logging.getLogger("lightgbm").setLevel(logging.WARNING)
    logging.getLogger("mlflow").setLevel(logging.WARNING)
    os.environ.setdefault("MLFLOW_ALLOW_FILE_STORE", "true")
    server.run(transport="stdio")


if __name__ == "__main__":
    main()