import argparse
import asyncio
import dataclasses
import datetime as dt
import json
import logging
from pathlib import Path

import pandas as pd

from .broker import NEW_YORK, Broker
from .config import Config, load_config
from .rebalance import plan_orders
from . import backtest
from .strategy import WARMUP, select_with_buffer, sessions_until_rebalance, signal_overnight_recovery
from .universe import fetch_sp500, load_universe

# MOC orders must reach NYSE by 15:50 ET; keep a margin.
ORDER_WINDOW = (dt.time(9, 31), dt.time(15, 45))


def update_universe(cfg: Config) -> None:
    table = fetch_sp500()
    cfg.data.universe.parent.mkdir(parents=True, exist_ok=True)
    table.to_csv(cfg.data.universe, index=False)
    print(f"{len(table)} symbols -> {cfg.data.universe}")


def show_rebalance_schedule(cfg: Config, today: dt.date) -> None:
    if not cfg.data.bars.exists():
        print("rebalance schedule unavailable: no daily bars; run `make bars`")
        return
    dates = pd.read_csv(cfg.data.bars, usecols=["date"], parse_dates=["date"])["date"]
    sessions = pd.DatetimeIndex(dates.dropna().unique()).sort_values()
    sessions = sessions[sessions.date <= today]
    if sessions.empty:
        print("rebalance schedule unavailable: no sessions through today; run `make bars`")
        return
    last = sessions[-1].date()
    st = cfg.strategy
    wait = sessions_until_rebalance(sessions, st.anchor, st.rebalance_every)
    if wait is None:
        schedule = f"first trading session on or after {st.anchor}"
    elif wait == 0:
        schedule = f"{last} (0 trading sessions remaining)"
    else:
        schedule = f"in {wait} trading sessions"
    print(f"next rebalance: {schedule} (as of {last}; every {st.rebalance_every} trading sessions)")
    if last < today:
        print(f"Bars end {last}; run `make bars` to update the rebalance countdown for {today}.")


async def check(cfg: Config) -> None:
    async with Broker(cfg.ib) as broker:
        state = await broker.account_state()
        print(f"account {state.account}  NAV ${state.nav_usd:,.2f}  cash ${state.cash_usd:,.2f}")
        holdings = pd.DataFrame([dataclasses.asdict(h) for h in state.holdings])
        print(holdings.to_string(index=False) if len(holdings) else "no positions")
        trades = await broker.open_trades()
        print(f"{len(trades)} open orders")
        for t in trades:
            print(f"  {t.contract.symbol} {t.order.action} {t.order.totalQuantity} {t.order.orderType} {t.orderStatus.status}")
        show_rebalance_schedule(cfg, dt.datetime.now(NEW_YORK).date())
        contracts, _ = await broker.qualify_stocks(["SPY"])
        if "SPY" not in contracts:
            print("no SPY price: IBKR could not qualify the SPY contract")
            return
        errors: dict[int, str] = {}

        def on_error(req_id, code, message, contract):
            if contract == contracts["SPY"]:
                errors[code] = message

        broker.ib.errorEvent += on_error
        try:
            prices = await broker.snapshot_prices(contracts)
        finally:
            broker.ib.errorEvent -= on_error
        if "SPY" in prices:
            print(f"SPY snapshot {prices['SPY']:.2f} (requested market_data_type={cfg.ib.market_data_type})")
        else:
            reason = "; ".join(f"IBKR {code}: {message}" for code, message in errors.items()) or "no valid snapshot received"
            print(f"no SPY price (requested market_data_type={cfg.ib.market_data_type}): {reason}")
            if 10197 in errors:
                print(
                    "Log out of the competing live market-data session, then retry `make check`. "
                    "Delayed mode (3/4) does not force delayed quotes when live data permissions exist; "
                    "changing ib.market_data_type cannot resolve this session conflict."
                )
            elif cfg.ib.market_data_type in (1, 2):
                print("Check the market data subscription, or set ib.market_data_type = 3 to allow delayed data.")
            else:
                print("Delayed data is already requested; check market data permissions and availability in IB Gateway/TWS.")


async def probe_gateway(cfg: Config, history: bool = False) -> None:
    # Broker's connect fails unless the Gateway answers positions, orders and
    # account updates, which needs a live session with IBKR's servers.
    async with Broker(dataclasses.replace(cfg.ib, client_id=cfg.ib.health_client_id)) as broker:
        status = f"account {broker.account}, server version {broker.ib.client.serverVersion()}"
        if history:
            # The historical data farm can fail on its own: on 2026-10-12 account requests
            # worked while every bar request timed out for 35 minutes, until the Gateway's
            # daily restart. Without this check `make gateway` never restarted it.
            contracts, _ = await broker.qualify_stocks(["SPY"])
            bars = (await broker.daily_bars(contracts, "5 D", timeout=30))[0] if contracts else None
            if bars is None or bars.empty:
                raise SystemExit(f"IB Gateway answers account requests ({status}) but returned no SPY bars "
                                 "within 30 s; its historical data farm is not serving")
            status += f", historical data OK (SPY through {bars['date'].max()})"
        print(f"IB Gateway OK: {status}")


async def update_bars(cfg: Config) -> None:
    symbols = load_universe(cfg.data.universe)
    async with Broker(cfg.ib) as broker:
        contracts, unknown = await broker.qualify_stocks(symbols)
        # Fail fast when the historical farm is down instead of timing out every symbol.
        probe = dict([next(iter(contracts.items()))])
        if (await broker.daily_bars(probe, "5 D", timeout=30))[0].empty:
            raise SystemExit(f"no historical data for {next(iter(probe))}; check the Gateway's HMDS farm status")
        bars, failed = await broker.daily_bars(contracts, cfg.data.bars_duration)
    if bars.empty:
        raise SystemExit(f"no bars for any of {len(contracts)} symbols; kept existing {cfg.data.bars}")
    cfg.data.bars.parent.mkdir(parents=True, exist_ok=True)
    bars.to_csv(cfg.data.bars, index=False)
    print(f"{bars['symbol'].nunique()} symbols, {bars['date'].min()} .. {bars['date'].max()} -> {cfg.data.bars}")
    if unknown:
        print(f"unqualified contracts ({len(unknown)}): {unknown}")
    if failed:
        print(f"no bars ({len(failed)}): {failed}")


def load_bars(path: Path, today: dt.date) -> tuple[pd.DataFrame, pd.DataFrame]:
    bars = pd.read_csv(path, parse_dates=["date"])
    opens, closes = (bars.pivot(index="date", columns="symbol", values=v).sort_index() for v in ("open", "close"))
    if opens.index[-1].date() != today:
        raise SystemExit(f"bars end {opens.index[-1].date()}, need today's open ({today}); run `make bars` after 09:30 ET")
    return opens, closes


async def rebalance(cfg: Config, submit: bool) -> None:
    ex, st = cfg.execution, cfg.strategy
    now = dt.datetime.now(NEW_YORK)
    universe = set(load_universe(cfg.data.universe))
    opens, closes = load_bars(cfg.data.bars, now.date())
    wait = sessions_until_rebalance(opens.index, st.anchor, st.rebalance_every)
    if wait != 0:
        after = f"in {wait} sessions" if wait else f"on the first session from {st.anchor}"
        print(f"not a rebalance session; holdings drift until the next one {after}")
        return
    # Today's bar is still forming, so the signal uses completed sessions only.
    history = opens.index[:-1]
    if len(history) <= WARMUP:
        raise SystemExit(f"need more than {WARMUP} completed sessions of bars, have {len(history)}")
    members = opens.columns.intersection(sorted(universe))
    score = signal_overnight_recovery(opens.loc[history, members], closes.loc[history, members]).iloc[-1]
    run_dir = cfg.root / "runs" / now.strftime("%Y%m%d-%H%M%S")

    async with Broker(cfg.ib) as broker:
        state = await broker.account_state()
        # The strategy owns the whole account: every position it does not select is sold.
        # Orders go out as SMART/USD stocks, so any other kind of position stops the run.
        if foreign := [h.symbol for h in state.holdings if (h.sec_type, h.currency) != ("STK", "USD")]:
            raise SystemExit(f"positions that are not USD stocks: {foreign}; close them before rebalancing")
        investable = state.nav_usd * (1 - ex.cash_buffer)
        weights = select_with_buffer(score, {h.symbol for h in state.holdings}, st.top_n, st.buffer)
        # Sized at today's open, the latest price in the bars; filled at the close. A holding
        # outside the index has no bars and takes the account's price; its exit size does not use it.
        orders = plan_orders(
            weights,
            {h.symbol: h.quantity for h in state.holdings},
            {h.symbol: h.market_price for h in state.holdings} | opens.iloc[-1].dropna().to_dict(),
            investable,
            max(ex.min_trade_value, ex.min_trade_fraction * state.nav_usd),
            ex.resize_band,
        )

        run_dir.mkdir(parents=True)
        plan = pd.DataFrame([dataclasses.asdict(o) for o in orders], columns=["symbol", "current", "target", "delta", "price"])
        plan["weight"] = plan["symbol"].map(weights).fillna(0.0)
        plan.to_csv(run_dir / "plan.csv", index=False)
        (run_dir / "account.json").write_text(json.dumps(dataclasses.asdict(state), indent=2))
        print(f"account {state.account}  NAV ${state.nav_usd:,.2f}  investable ${investable:,.2f}  signal on sessions through {history[-1].date()}")
        print(plan.to_string(index=False) if len(plan) else "no orders")
        if not submit or not orders:
            print(f"dry run -> {run_dir}")
            return
        if not ORDER_WINDOW[0] <= now.time() <= ORDER_WINDOW[1]:
            raise SystemExit(f"{now:%H:%M} ET is outside the order window {ORDER_WINDOW[0]}-{ORDER_WINDOW[1]}")

        if open_trades := await broker.open_trades():
            # A run whose confirmation timed out exits non-zero and the trade service retries it;
            # if IBKR holds every order that run sent, the session's orders are in.
            refs = {t.order.orderRef for t in open_trades}
            sent = [cfg.root / "runs" / ref.removeprefix("ibkr-quant ") / "orders.csv" for ref in refs]
            if all(ref.startswith(f"ibkr-quant {now:%Y%m%d}-") for ref in refs) and all(p.exists() for p in sent):
                sent_ids = set(pd.concat(pd.read_csv(p) for p in sent)["order_id"])
                if missing := sent_ids - {t.order.orderId for t in open_trades}:
                    raise SystemExit(f"today's orders {sorted(missing)} are no longer open at IBKR "
                                     "(rejected, cancelled or filled); see `make check`")
                print(f"{len(open_trades)} orders from {', '.join(sorted(refs))} are open at IBKR; not submitting again")
                return
            raise SystemExit(f"{len(open_trades)} open orders already exist; cancel or let them finish first")
        contracts, unknown = await broker.qualify_stocks([o.symbol for o in orders])
        if unknown:
            raise SystemExit(f"unqualified contracts: {unknown}")
        limits: dict[str, float] = {}
        if ex.order_type == "LMT":
            snapshot = await broker.snapshot_prices(contracts)
            if missing := [o.symbol for o in orders if o.symbol not in snapshot]:
                raise SystemExit(f"no snapshot price for {missing}")
            offset = ex.limit_offset_bps / 1e4
            limits = {
                o.symbol: round(snapshot[o.symbol] * (1 + offset if o.delta > 0 else 1 - offset), 2)
                for o in orders
            }

        trades = [
            broker.place(contracts[o.symbol], o.delta, ex.order_type, limits.get(o.symbol), f"ibkr-quant {run_dir.name}")
            for o in orders
        ]
        confirmation_error = None
        try:
            await broker.confirm_orders(trades)
        except RuntimeError as exc:
            confirmation_error = exc
        submitted = pd.DataFrame(
            [
                {
                    "symbol": t.contract.symbol,
                    "action": t.order.action,
                    "quantity": t.order.totalQuantity,
                    "order_type": ex.order_type,
                    "limit": limits.get(t.contract.symbol),
                    "order_id": t.order.orderId,
                    "perm_id": t.order.permId,
                    "status": t.orderStatus.status,
                }
                for t in trades
            ]
        )
        submitted.to_csv(run_dir / "orders.csv", index=False)
        print(submitted.to_string(index=False))
        if confirmation_error:
            raise SystemExit(f"{confirmation_error}; recorded status -> {run_dir}")
        print(f"accepted {len(trades)} orders -> {run_dir}")


def run_backtest(cfg: Config, top_ns: list[int], everies: list[int], capitals: list[float], fee: str) -> None:
    st = cfg.strategy
    bars = pd.read_csv(cfg.data.bars, parse_dates=["date"])
    opens = bars.pivot(index="date", columns="symbol", values="open").sort_index()
    closes = bars.pivot(index="date", columns="symbol", values="close").sort_index()
    # As live: pick among the index members of each day, scoring them on their
    # full price history. Only today's members are known, so names removed from
    # the index are missing (survivorship bias remains).
    added = pd.read_csv(cfg.data.universe, parse_dates=["added"]).set_index("symbol")["added"]
    member = pd.DataFrame(opens.index.values[:, None] >= added.reindex(opens.columns).values[None, :],
                          index=opens.index, columns=opens.columns)
    # Row d holds the signal from sessions through d - 1, known at d's open.
    signal = signal_overnight_recovery(opens, closes, member).shift(1)
    fee_model = backtest.ibkr_fixed_fee if fee == "ibkr-fixed" else backtest.flat_fee(2.0)
    rows = []
    for capital in capitals:
        for every in everies:
            for top_n in top_ns:
                run_cfg = backtest.RunConfig(top_n, every, st.buffer, capital, fee_model)
                result = backtest.run(signal, opens, closes, run_cfg)
                rows.append({"capital": capital, "every": every, "top_n": top_n} | backtest.metrics(result, capital))
    days = signal.index[backtest.WARMUP:]
    # Same-universe baseline: equal weight, re-weighted at every close, no costs.
    ew = closes.ffill().pct_change().shift(-1).loc[days[:-1]].mean(axis=1)
    table = pd.DataFrame(rows)
    path = cfg.root / "runs" / f"backtest-{dt.datetime.now(NEW_YORK):%Y%m%d-%H%M%S}.csv"
    path.parent.mkdir(exist_ok=True)
    table.to_csv(path, index=False)
    pct = ["total", "h1", "h2", "vol", "max_dd", "turnover_day", "fees_pct"]
    print(f"{days[0].date()} .. {days[-1].date()} ({len(days)} sessions), fee={fee}")
    print(f"equal-weight universe: total {((1 + ew).prod() - 1) * 100:.1f}%, sharpe {ew.mean() / ew.std() * 252 ** 0.5:.2f}")
    print(table.assign(**{c: (table[c] * 100).round(1) for c in pct}).round({"sharpe": 2}).to_string(index=False))
    print(f"-> {path}")


async def fills(cfg: Config) -> None:
    async with Broker(cfg.ib) as broker:
        executions = await broker.executions()
    if executions.empty:
        print("no executions")
        return
    path = cfg.root / "runs" / f"fills-{dt.datetime.now(NEW_YORK):%Y%m%d}.csv"
    path.parent.mkdir(exist_ok=True)
    executions.to_csv(path, index=False)
    print(executions.to_string(index=False))
    print(f"-> {path}")


def main() -> None:
    parser = argparse.ArgumentParser(prog="ibkr-quant")
    parser.add_argument("--config", type=Path, default=Path("config.toml"))
    commands = parser.add_subparsers(dest="command", required=True)
    commands.add_parser("universe", help="download S&P 500 constituents")
    commands.add_parser("check", help="connect and show account, positions, open orders, SPY price")
    gateway_parser = commands.add_parser("gateway", help="exit non-zero unless IB Gateway answers account requests")
    gateway_parser.add_argument("--history", action="store_true", help="also require SPY daily bars from the historical data farm")
    commands.add_parser("bars", help="download adjusted daily bars for the universe")
    rebalance_parser = commands.add_parser("rebalance", help="plan orders (dry run unless --submit)")
    rebalance_parser.add_argument("--submit", action="store_true")
    commands.add_parser("fills", help="show today's executions")
    backtest_parser = commands.add_parser("backtest", help="replay the strategy on the downloaded bars")
    backtest_parser.add_argument("--top-n", type=int, nargs="+", help="default: strategy.top_n")
    backtest_parser.add_argument("--every", type=int, nargs="+", help="rebalance interval in sessions; default: strategy.rebalance_every")
    backtest_parser.add_argument("--capital", type=float, nargs="+", default=[1_000_000])
    backtest_parser.add_argument("--fee", choices=["ibkr-fixed", "flat2"], default="ibkr-fixed")
    args = parser.parse_args()

    logging.basicConfig(level=logging.WARNING, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
    cfg = load_config(args.config)
    match args.command:
        case "universe":
            update_universe(cfg)
        case "check":
            asyncio.run(check(cfg))
        case "gateway":
            asyncio.run(probe_gateway(cfg, args.history))
        case "bars":
            asyncio.run(update_bars(cfg))
        case "rebalance":
            asyncio.run(rebalance(cfg, args.submit))
        case "fills":
            asyncio.run(fills(cfg))
        case "backtest":
            run_backtest(cfg, args.top_n or [cfg.strategy.top_n], args.every or [cfg.strategy.rebalance_every], args.capital, args.fee)