#!/usr/bin/env python3
"""Rank live strategy state files by current marked-to-market return."""

import argparse
import glob
import json
import os
from dataclasses import dataclass
from pathlib import Path

import data_manager
from trader import get_prices_for_date, portfolio_value
from universe import UNIVERSES


@dataclass
class RankedStrategy:
    strategy: str
    return_pct: float
    value: float
    cash: float
    n_positions: int
    buys: int
    sells: int
    holdings: str
    state_date: str


def _state_name(path: str, market: str) -> str:
    base = os.path.basename(path)
    prefix = f"trader_{market}_"
    return base[len(prefix):-len(".json")]


def _load_close_data(market: str, update: bool):
    if update:
        universe = UNIVERSES[market]
        tickers = universe["fetch"]()
        benchmark = universe["benchmark"]
        all_tickers = sorted(set(tickers + [benchmark]))
        return data_manager.update(market, all_tickers)

    data = data_manager.load(market)
    if data is None:
        raise RuntimeError(
            f"No cached data found for market '{market}'. Run without --no-update first."
        )
    return data


def _format_holdings(holdings: dict) -> str:
    return ", ".join(f"{ticker}:{shares:g}" for ticker, shares in sorted(holdings.items()))


def rank_market(market: str, update: bool, include_sim: bool) -> tuple[str, list[RankedStrategy]]:
    close = _load_close_data(market, update)
    latest_date = str(close.index[-1].date())

    rows: list[RankedStrategy] = []
    for path in sorted(glob.glob(f"data/trader_{market}_*.json")):
        strategy = _state_name(path, market)
        if strategy.startswith("sim_") and not include_sim:
            continue

        state = json.loads(Path(path).read_text())
        if "initial_capital" not in state:
            continue

        holdings = state.get("holdings", {}) or {}
        cash = float(state.get("cash", 0.0) or 0.0)
        prices = get_prices_for_date(list(holdings), close.index[-1], close)
        value = portfolio_value(holdings, prices, cash)
        initial = float(state["initial_capital"])
        return_pct = (value / initial - 1.0) * 100.0 if initial else 0.0

        trades = state.get("trade_log", []) or []
        buys = sum(1 for trade in trades if trade.get("action") == "BUY")
        sells = sum(1 for trade in trades if trade.get("action") == "SELL")
        equity = state.get("daily_equity", {}) or {}
        state_date = max(equity.keys()) if equity else ""

        rows.append(
            RankedStrategy(
                strategy=strategy,
                return_pct=return_pct,
                value=value,
                cash=cash,
                n_positions=len(holdings),
                buys=buys,
                sells=sells,
                holdings=_format_holdings(holdings),
                state_date=state_date,
            )
        )

    rows.sort(key=lambda row: row.return_pct, reverse=True)
    return latest_date, rows


def print_market_table(market: str, latest_date: str, rows: list[RankedStrategy], top: int) -> None:
    print(f"\n{market.upper()} Top{top}  latest_price_date={latest_date}")
    print(
        f"{'#':>2}  {'Strategy':<42} {'Return':>8} {'Value':>10} "
        f"{'Cash':>9} {'Pos':>3} {'Buy/Sell':>8} {'StateDate':>10}  Holdings"
    )
    print("-" * 132)

    for idx, row in enumerate(rows[:top], 1):
        print(
            f"{idx:>2}  {row.strategy:<42} {row.return_pct:>7.2f}% "
            f"{row.value:>10.2f} {row.cash:>9.2f} {row.n_positions:>3} "
            f"{row.buys:>3}/{row.sells:<4} {row.state_date:>10}  {row.holdings}"
        )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Update market data and rank live strategy state files."
    )
    parser.add_argument(
        "--market",
        choices=["all", "us", "cn"],
        default="all",
        help="Market to rank. Default: all.",
    )
    parser.add_argument(
        "--top",
        type=int,
        default=10,
        help="Number of strategies to print per market. Default: 10.",
    )
    parser.add_argument(
        "--no-update",
        action="store_true",
        help="Use cached data only; do not download new prices.",
    )
    parser.add_argument(
        "--include-sim",
        action="store_true",
        help="Include state files whose strategy name starts with sim_.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    markets = ["us", "cn"] if args.market == "all" else [args.market]

    for market in markets:
        latest_date, rows = rank_market(
            market=market,
            update=not args.no_update,
            include_sim=args.include_sim,
        )
        print_market_table(market, latest_date, rows, args.top)


if __name__ == "__main__":
    main()