execution/pdt_tracker.py

"""
execution/pdt_tracker.py — Pattern Day Trade rule enforcement.
Uses FINRA settlement-date logic (T+1 for equities).
"""
from __future__ import annotations

import datetime
import logging

import pandas_market_calendars as mcal

import config
from data import storage

logger = logging.getLogger("trading_system.pdt_tracker")

_nyse_cal = mcal.get_calendar("NYSE")


def next_business_day(dt: datetime.date) -> datetime.date:
    """Get the next NYSE business day after dt (T+1 settlement)."""
    start = dt + datetime.timedelta(days=1)
    end = dt + datetime.timedelta(days=10)
    schedule = _nyse_cal.schedule(
        start_date=start.isoformat(),
        end_date=end.isoformat(),
    )
    if schedule.empty:
        return dt + datetime.timedelta(days=1)
    return schedule.index[0].date()


def settlement_date(trade_date: datetime.date) -> datetime.date:
    """Compute settlement date for a trade (T+1 for equities)."""
    return next_business_day(trade_date)


def get_last_n_business_days(n: int = 5) -> tuple[datetime.date, datetime.date]:
    """Get the date range covering the last N NYSE business days."""
    today = datetime.date.today()
    start = today - datetime.timedelta(days=n * 3)  # buffer for weekends/holidays
    schedule = _nyse_cal.schedule(
        start_date=start.isoformat(),
        end_date=today.isoformat(),
    )
    if len(schedule) < n:
        return start, today
    return schedule.index[-n].date(), today


def is_day_trade(
    symbol: str,
    side: str,
    open_positions: dict[str, dict],
) -> bool:
    """Check if closing a position would count as a day trade.

    A day trade = opening AND closing the same position within the same
    settlement day.

    Args:
        symbol: Symbol being traded
        side: "buy" or "sell" — the incoming order side
        open_positions: {symbol: {side, open_date, qty}} of current positions

    Returns: True if this would be a day trade
    """
    if symbol not in open_positions:
        return False

    pos = open_positions[symbol]
    pos_side = pos.get("side", "")

    # Check if this order would close the position
    # Buy closes a short, sell closes a long
    is_closing = (pos_side == "buy" and side == "sell") or \
                 (pos_side == "sell" and side == "buy")

    if not is_closing:
        return False

    # Compare settlement dates
    open_date = pos.get("open_date")
    if open_date is None:
        return False

    if isinstance(open_date, str):
        open_date = datetime.date.fromisoformat(open_date[:10])
    elif isinstance(open_date, datetime.datetime):
        open_date = open_date.date()

    close_date = datetime.date.today()
    open_settlement = settlement_date(open_date)
    close_settlement = settlement_date(close_date)

    return open_settlement == close_settlement


def count_day_trades() -> int:
    """Count day trades in the last 5 NYSE business days."""
    start_date, end_date = get_last_n_business_days(5)
    return storage.count_pdt_trades_in_window(
        start_date.isoformat(), end_date.isoformat()
    )


def record_day_trade(
    symbol: str,
    side: str,
    qty: float,
    open_date: datetime.datetime,
    close_date: datetime.datetime,
):
    """Record a completed day trade in the DB."""
    settle = settlement_date(close_date.date())
    storage.insert_pdt_trade(
        symbol=symbol,
        open_date=open_date.isoformat(),
        close_date=close_date.isoformat(),
        settlement_date=settle.isoformat(),
        side=side,
        qty=qty,
    )
    logger.info(
        "Recorded day trade: %s %s %.2f shares, settlement=%s",
        symbol, side, qty, settle,
    )


def can_day_trade(
    symbol: str,
    side: str,
    open_positions: dict[str, dict],
    alert_callback=None,
) -> tuple[bool, str | None]:
    """Check if a day trade is allowed under PDT rules.

    Returns: (allowed, reason_if_blocked)
    """
    # PDT doesn't apply if above $25k or cash account
    if config.ACCOUNT_BALANCE_ABOVE_25K:
        return True, None

    if config.ALPACA_ACCOUNT_TYPE == "cash":
        return True, None

    # Check if this order would create a day trade
    if not is_day_trade(symbol, side, open_positions):
        return True, None

    # Count existing day trades
    current_count = count_day_trades()
    if current_count >= config.PDT_MAX_DAY_TRADES:
        reason = (
            f"PDT limit reached: {current_count}/{config.PDT_MAX_DAY_TRADES} "
            f"day trades in 5-day window"
        )
        logger.warning(reason)
        if alert_callback:
            alert_callback(f"⚠️ {reason}")
        return False, reason

    return True, None