"""
JSON-config-driven **multi-sleeve** portfolio merge: unlimited named sleeves, each with its own
JSONL trade log, PnL key, and risk-reference model.

This replaces ad-hoc “slot” limits in ``portfolio_vrp_plus_vxx`` (single put / straddle / RR)
for **new** workflows. Legacy CLI can keep using ``portfolio_vrp_plus_vxx.py``; prefer this
engine when combining VRP + IV + lit4 + multiple VXX legs + future sleeves.

**Daily contribution** (same math as the legacy merge, unified):

    contrib_s(t) = (capital_s / ref_s) * raw_pnl_s(t)

where ``raw_pnl_s`` is the sum of ``pnl_key`` on each calendar day (exit date). For
``risk_model: vrp``, ``ref_s`` is ``vrp_ref_usd`` from the sleeve block.

**Config:** JSON (stdlib) under ``RenTech/strategy_stack/config/*.json`` — add sleeves to the
``sleeves`` array; set ``enabled: false`` to reserve a slot without loading files.

See ``config/portfolio_sleeves_default.json``.
"""
from __future__ import annotations

import csv
import json
import math
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Literal

import pandas as pd

_REPO_ROOT = Path(__file__).resolve().parents[2]
if str(_REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(_REPO_ROOT))

from RenTech.strategy_stack.iv_mispricing_complement import _load_jsonl, _pnl_series_from_trades
from RenTech.strategy_stack.portfolio_vrp_plus_vxx import (
    _infer_broker_risk_overlay,
    _infer_broker_risk_vxx,
    _metrics_block,
    _overlay_sleeve_risk_reference_usd,
)

RiskModel = Literal["vrp", "broker_risk_overlay", "broker_risk_vxx", "fixed_ref"]


def _resolve_path(repo_root: Path, p: str | Path) -> Path:
    path = Path(p).expanduser()
    if path.is_absolute():
        return path.resolve()
    return (repo_root / path).resolve()


@dataclass(frozen=True)
class ResolvedSleeve:
    id: str
    trades_path: Path
    pnl_key: str
    exit_key: str
    risk_model: RiskModel
    capital_usd: float
    ref_usd: float
    enabled: bool


def _infer_ref_usd(path: Path, model: RiskModel) -> float:
    if not path.is_file():
        return 1.0
    if model == "broker_risk_overlay":
        return max(_overlay_sleeve_risk_reference_usd(path, _infer_broker_risk_overlay), 1.0)
    if model == "broker_risk_vxx":
        return max(_overlay_sleeve_risk_reference_usd(path, _infer_broker_risk_vxx), 1.0)
    raise ValueError(f"_infer_ref_usd: unexpected model {model}")


def load_sleeve_config(
    path: Path,
    *,
    repo_root: Path | None = None,
    total_capital_override: float | None = None,
) -> tuple[dict[str, Any], list[ResolvedSleeve]]:
    repo_root = repo_root or _REPO_ROOT
    raw = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(raw, dict) or "sleeves" not in raw:
        raise ValueError(f"Invalid config: {path}")
    acct = raw.get("account") or {}
    total = float(
        total_capital_override
        if total_capital_override is not None
        else acct.get("total_capital_usd", 100_000.0)
    )
    exit_key = str(acct.get("exit_key", "exit_date"))

    resolved: list[ResolvedSleeve] = []
    for block in raw["sleeves"]:
        if not isinstance(block, dict):
            continue
        if not bool(block.get("enabled", True)):
            continue
        sid = str(block["id"]).strip()
        if not sid:
            raise ValueError("sleeve missing id")
        tpath = _resolve_path(repo_root, block["trades_jsonl"])
        pnl_key = str(block.get("pnl_key", "pnl_total"))
        model = str(block.get("risk_model", "broker_risk_overlay"))
        if model not in ("vrp", "broker_risk_overlay", "broker_risk_vxx", "fixed_ref"):
            raise ValueError(f"sleeve {sid}: unknown risk_model {model}")

        if model == "fixed_ref":
            ref_usd = float(block.get("ref_usd", 1.0))
        elif model == "vrp":
            ref_usd = float(block.get("vrp_ref_usd", 100_000.0))
        else:
            ref_usd = _infer_ref_usd(tpath, model)  # type: ignore[arg-type]

        if "capital_usd" in block and block["capital_usd"] is not None:
            cap = float(block["capital_usd"])
        elif "capital_pct" in block and block["capital_pct"] is not None:
            cap = total * float(block["capital_pct"])
        else:
            cap = 0.0

        resolved.append(
            ResolvedSleeve(
                id=sid,
                trades_path=tpath,
                pnl_key=pnl_key,
                exit_key=exit_key,
                risk_model=model,  # type: ignore[arg-type]
                capital_usd=float(cap),
                ref_usd=max(float(ref_usd), 1e-12),
                enabled=True,
            )
        )
    return raw, resolved


def build_daily_matrix(
    sleeves: list[ResolvedSleeve],
    *,
    repo_root: Path | None = None,
) -> tuple[pd.DatetimeIndex, dict[str, pd.Series], dict[str, float]]:
    repo_root = repo_root or _REPO_ROOT
    series_by_id: dict[str, pd.Series] = {}
    scale_by_id: dict[str, float] = {}
    nonempty: list[pd.Series] = []

    for s in sleeves:
        if not s.enabled:
            continue
        if not s.trades_path.is_file():
            series_by_id[s.id] = pd.Series(dtype=float)
            scale_by_id[s.id] = float(s.capital_usd) / float(s.ref_usd)
            print(f"WARNING: missing trades file for sleeve {s.id}: {s.trades_path}", file=sys.stderr)
            continue
        rows = _load_jsonl(s.trades_path)
        ser = _pnl_series_from_trades(rows, exit_key=s.exit_key, pnl_key=s.pnl_key)
        ser = ser.sort_index()
        series_by_id[s.id] = ser
        nonempty.append(ser)
        scale_by_id[s.id] = float(s.capital_usd) / float(s.ref_usd)

    if not nonempty:
        raise SystemExit("No enabled sleeve had a readable trades JSONL.")

    all_idx = nonempty[0].index
    for ser in nonempty[1:]:
        all_idx = all_idx.union(ser.index)

    all_idx = all_idx.sort_values()
    t_pad = all_idx[0] - pd.Timedelta(days=1)
    if t_pad not in all_idx:
        all_idx = all_idx.insert(0, t_pad)

    for s in sleeves:
        if not s.enabled:
            continue
        ser = series_by_id.get(s.id, pd.Series(dtype=float)).reindex(all_idx, fill_value=0.0)
        series_by_id[s.id] = ser
        scale_by_id[s.id] = float(s.capital_usd) / float(s.ref_usd)

    return all_idx, series_by_id, scale_by_id


def merged_equity_frame(
    sleeves: list[ResolvedSleeve],
    *,
    total_capital: float,
    repo_root: Path | None = None,
) -> pd.DataFrame:
    idx, raw_by_id, scale_by_id = build_daily_matrix(sleeves, repo_root=repo_root)
    eq0 = float(total_capital)
    parts: dict[str, pd.Series] = {}
    total = pd.Series(0.0, index=idx)
    for s in sleeves:
        if not s.enabled:
            continue
        r = raw_by_id[s.id] * scale_by_id[s.id]
        parts[f"pnl_{s.id}"] = r
        total = total + r
    parts["pnl_total"] = total
    parts["eq_full"] = eq0 + total.cumsum()
    return pd.DataFrame(parts, index=idx)


_PROMOTED_FROM_SRC: frozenset[str] = frozenset({
    # size / leg identity — now promoted to legs_display / contracts_qty_logged
    "contracts", "qty", "legs_json", "entry_legs_json",
    # margin / risk — promoted to broker_risk_usd_logged
    "broker_risk_usd", "broker_risk_per_contract_usd",
    "max_margin", "max_loss_total_usd", "max_loss_one_usd",
    # instrument / structure description — absorbed into legs_display
    "structure", "mode", "strategy", "underlying",
    "expiration", "expiration_near", "expiration_far",
    "strike", "put_strike", "call_strike", "short_strike", "long_strike",
    # lit4 metadata now in legs_display
    "sid", "description", "source",
    # portfolio-level noise (not trade-specific)
    "nav_at_entry_usd", "risk_pct_of_portfolio", "contracts_requested",
})


def _fmt_expiry(s: str) -> str:
    """'2016-04-15' → \"Apr'16\"."""
    if not s:
        return ""
    from datetime import datetime
    try:
        return datetime.strptime(str(s)[:10], "%Y-%m-%d").strftime("%b'%y")
    except (ValueError, TypeError):
        return str(s)[:7]


def _strike_str(v: Any) -> str:
    try:
        f = float(v)
        if math.isfinite(f):
            return f"{f:.2f}".rstrip("0").rstrip(".")
    except (TypeError, ValueError):
        pass
    return str(v)


def _build_legs_display(r: dict[str, Any]) -> str:
    """
    Brokerage-statement style trade description showing actual contracts per leg.

    Rules (in priority order):
    1. VRP rows: ``legs_json`` JSON array + ``qty``  →  per-leg strikes with qty.
    2. IV/VXX rows: ``contracts`` field + structure/mode/strategy fields.
    3. Lit4 and any other rows with no count: "1 lot <structure>" + description.
    """
    def _cstr(v: Any) -> str:
        try:
            f = float(v)
            if math.isfinite(f):
                return str(int(f)) if abs(f - round(f)) < 1e-9 else str(f)
        except (TypeError, ValueError):
            pass
        return str(v)

    # ── 1. VRP: legs_json is a JSON array ─────────────────────────────────────
    raw_legs = r.get("legs_json", "")
    if raw_legs and isinstance(raw_legs, str) and raw_legs.strip().startswith("["):
        try:
            legs: list[dict] = json.loads(raw_legs)
        except json.JSONDecodeError:
            legs = []
        qty_raw = r.get("qty", "")
        qty = _cstr(qty_raw) if qty_raw != "" else "?"
        if legs:
            expiries = [str(lg.get("expiry", "")) for lg in legs]
            rights   = [str(lg.get("right",  "")).strip() for lg in legs]
            strikes  = [_strike_str(lg.get("strike", "")) for lg in legs]
            same_exp   = len(set(expiries)) == 1
            same_right = len(set(rights))   == 1
            if same_exp and same_right and len(legs) == 2:
                # simple spread: e.g. "2× P 172/170 Apr'16"
                exp = _fmt_expiry(expiries[0])
                return f"{qty}× {rights[0]} {strikes[0]}/{strikes[1]} {exp}"
            # diagonal or complex: list every leg
            parts = [f"{r_}{k} {_fmt_expiry(e)}" for r_, k, e in zip(rights, strikes, expiries)]
            return f"{qty}× " + " / ".join(parts)
        return f"{qty}× spread"

    # ── 2. IV / VXX: contracts field present ──────────────────────────────────
    raw_c = r.get("contracts")
    if raw_c is not None:
        c = _cstr(raw_c)
        und = str(r.get("underlying", "")).strip() or "SPY"
        structure = (
            str(r.get("structure", "") or r.get("mode", "") or r.get("strategy", ""))
            .strip()
            .lower()
        )
        if structure == "otm_put":
            k = _strike_str(r.get("strike", ""))
            exp = _fmt_expiry(str(r.get("expiration", "")).strip())
            return f"{c}× {und} put K{k} {exp}"
        if structure == "straddle":
            k = _strike_str(r.get("strike", ""))
            exp = _fmt_expiry(str(r.get("expiration", "")).strip())
            return f"{c}× {und} straddle K{k} {exp}"
        if structure == "risk_reversal":
            ps  = _strike_str(r.get("put_strike",  ""))
            cs  = _strike_str(r.get("call_strike",  ""))
            exp = _fmt_expiry(str(r.get("expiration_near", "") or r.get("expiration_far", "")).strip())
            return f"{c}× {und} RR put{ps}/call{cs} {exp}"
        if structure == "bear_call":
            sk  = _strike_str(r.get("short_strike", ""))
            lk  = _strike_str(r.get("long_strike",  ""))
            exp = _fmt_expiry(str(r.get("expiration", "")).strip())
            return f"{c}× {und} bear-call C{sk}/{lk} {exp}"
        if structure == "long_call":
            ck  = _strike_str(r.get("call_strike", ""))
            exp = _fmt_expiry(str(r.get("expiration", "")).strip())
            return f"{c}× {und} long-call C{ck} {exp}"
        return f"{c}× {und} {structure}".strip()

    # ── 3. Lit4 / unknown: no count in log ────────────────────────────────────
    structure = str(r.get("structure", "")).strip()
    desc = str(r.get("description", "")).strip()
    label = structure or "position"
    if desc:
        return f"1 lot {label} — {desc}"
    return f"1 lot {label}"


def _broker_risk_logged(r: dict[str, Any]) -> float | None:
    """Max capital at risk exactly as logged: broker_risk_usd > max_loss_total_usd > max_margin."""
    for key in ("broker_risk_usd", "max_loss_total_usd", "max_margin"):
        v = r.get(key)
        if v is not None:
            try:
                fv = float(v)
                if math.isfinite(fv) and fv > 0:
                    return round(fv, 2)
            except (TypeError, ValueError):
                pass
    return None


def _logged_contracts_qty_str(r: dict[str, Any]) -> str:
    """Simple numeric count for CSV column; prefer ``contracts`` then ``qty``."""
    for key in ("contracts", "qty"):
        v = r.get(key)
        if v is not None:
            try:
                f = float(v)
                if math.isfinite(f):
                    return str(int(f)) if abs(f - round(f)) < 1e-9 else str(f)
            except (TypeError, ValueError):
                pass
    return ""


def scaled_trade_rows(s: ResolvedSleeve) -> list[dict[str, Any]]:
    if not s.trades_path.is_file():
        return []
    rows = _load_jsonl(s.trades_path)
    scale = float(s.capital_usd) / float(max(s.ref_usd, 1e-12))
    out: list[dict[str, Any]] = []
    for r in rows:
        raw = float(r.get(s.pnl_key, 0.0) or 0.0)
        ex = r.get(s.exit_key)
        en = r.get("entry_date", "")
        # Keep only non-promoted extra fields in src_json (engineers' detail column)
        extras = {
            k: v
            for k, v in r.items()
            if k not in (s.exit_key, s.pnl_key, "entry_date") and k not in _PROMOTED_FROM_SRC
        }
        out.append(
            {
                "sleeve_id": s.id,
                "entry_date": str(en)[:10] if en is not None else "",
                "exit_date": str(ex)[:10] if ex is not None else "",
                "legs_display": _build_legs_display(r),
                "broker_risk_usd_logged": _broker_risk_logged(r),
                "contracts_qty_logged": _logged_contracts_qty_str(r),
                "pnl_usd_raw_in_log": round(raw, 6),
                "pnl_usd_scaled": round(raw * scale, 6),
                "sleeve_scale_factor": round(scale, 10),
                "sleeve_risk_ref_usd": round(s.ref_usd, 2),
                "sleeve_capital_usd": round(s.capital_usd, 2),
                "risk_model": s.risk_model,
                "source_jsonl": str(s.trades_path),
                "src_json": json.dumps(extras, default=str)[:8000],
            }
        )
    return out


def _trade_log_fieldnames(rows: list[dict[str, Any]]) -> list[str]:
    if not rows:
        return []
    keys_union: set[str] = {k for row in rows for k in row}
    preferred = [
        "sleeve_id",
        "entry_date",
        "exit_date",
        "legs_display",
        "broker_risk_usd_logged",
        "contracts_qty_logged",
        "pnl_usd_raw_in_log",
        "pnl_usd_scaled",
        "sleeve_scale_factor",
        "sleeve_risk_ref_usd",
        "sleeve_capital_usd",
        "risk_model",
        "source_jsonl",
        "src_json",
    ]
    ordered = [k for k in preferred if k in keys_union]
    ordered.extend(sorted(keys_union - set(ordered)))
    return ordered


def write_all_sleeves_trade_log(path: Path, sleeves: list[ResolvedSleeve]) -> int:
    rows: list[dict[str, Any]] = []
    for s in sleeves:
        if s.enabled:
            rows.extend(scaled_trade_rows(s))

    def _key(r: dict) -> tuple:
        return (r.get("exit_date") or "", str(r.get("sleeve_id") or ""))

    rows.sort(key=_key)
    if not rows:
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text("", encoding="utf-8")
        return 0
    keys = _trade_log_fieldnames(rows)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=keys, extrasaction="ignore")
        w.writeheader()
        for r in rows:
            w.writerow(r)
    return len(rows)


def write_manifest(
    path: Path,
    *,
    config_path: Path,
    sleeves: list[ResolvedSleeve],
    metrics: dict,
    outputs: dict[str, str],
) -> None:
    def _ser(x: Any) -> Any:
        if hasattr(x, "item"):
            try:
                return float(x)
            except Exception:
                return str(x)
        return x

    m2 = {k: _ser(v) for k, v in metrics.items()}
    lines = [
        f"config: {config_path}",
        "",
        "enabled sleeves:",
    ]
    for s in sleeves:
        if not s.enabled:
            continue
        lines.append(
            f"  - {s.id}: capital=${s.capital_usd:,.2f}  ref=${s.ref_usd:,.2f}  "
            f"scale={s.capital_usd/s.ref_usd:.8f}  model={s.risk_model}  file={s.trades_path.name}"
        )
    lines.extend(["", "metrics (eq_full):", json.dumps(m2, indent=2)])
    lines.extend(["", "outputs:"])
    for k, v in outputs.items():
        lines.append(f"  {k}: {v}")
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def run_from_config(
    config_path: Path,
    *,
    repo_root: Path | None = None,
    total_capital_override: float | None = None,
    out_equity_csv: Path,
    out_trade_log_csv: Path | None = None,
    out_manifest: Path | None = None,
    print_metrics: bool = True,
) -> pd.DataFrame:
    repo_root = repo_root or _REPO_ROOT
    raw, sleeves = load_sleeve_config(
        config_path,
        repo_root=repo_root,
        total_capital_override=total_capital_override,
    )
    total = float(
        total_capital_override
        if total_capital_override is not None
        else (raw.get("account") or {}).get("total_capital_usd", 100_000.0)
    )
    frame = merged_equity_frame(sleeves, total_capital=total, repo_root=repo_root)
    out_equity_csv.parent.mkdir(parents=True, exist_ok=True)
    frame.to_csv(out_equity_csv)

    n_trades = 0
    if out_trade_log_csv:
        n_trades = write_all_sleeves_trade_log(out_trade_log_csv, sleeves)

    m = _metrics_block(frame["eq_full"], "multi_sleeve")
    if print_metrics:
        print(f"  return%={m['return_pct']:.2f}  maxDD%={m['max_dd_pct']:.2f}  "
              f"CAGR%={m['cagr_pct']:.2f}  Sharpe={m['sharpe']:.3f}  end=${m['end_equity']:,.0f}")
        print(f"  equity_csv: {out_equity_csv}")
        if out_trade_log_csv:
            print(f"  trade_log_rows: {n_trades}  path: {out_trade_log_csv}")

    if out_manifest:
        write_manifest(
            out_manifest,
            config_path=config_path,
            sleeves=sleeves,
            metrics=m,
            outputs={"equity": str(out_equity_csv), "trades": str(out_trade_log_csv or "")},
        )
        if print_metrics:
            print(f"  manifest: {out_manifest}")

    return frame