"""Shared live trading-policy state and availability checks."""

from __future__ import annotations

from collections.abc import Iterable
from datetime import date, datetime, timezone
from zoneinfo import ZoneInfo

from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.paper_trade import PaperTrade

# Gate-reset "day" boundary matches US cash equities session calendar, not UTC.
_REENTRY_DAY_TZ = ZoneInfo("America/New_York")


def _ny_trading_date(moment: datetime) -> date:
    """Calendar date in America/New_York for a gate-reset observation."""
    if moment.tzinfo is None:
        moment = moment.replace(tzinfo=timezone.utc)
    return moment.astimezone(_REENTRY_DAY_TZ).date()


MANUAL_BOOK = "manual"
SHADOW_BOOK = "shadow"


async def _latest_initial_stop_trades(
    db: AsyncSession,
    *,
    closed_before: datetime | None = None,
    book: str = MANUAL_BOOK,
) -> dict[int, PaperTrade]:
    """Return a ticker's latest closed trade only when it was an initial stop.

    Scoped to one ``book``: the discretionary and shadow books diverge as soon
    as their entries differ, so each must see only its own stop history when
    deciding whether a ticker is locked out of re-entry.
    """
    ranked_stmt = (
        select(
            PaperTrade.id.label("trade_id"),
            func.row_number()
            .over(
                partition_by=PaperTrade.ticker_id,
                order_by=(PaperTrade.closed_at.desc(), PaperTrade.id.desc()),
            )
            .label("recency"),
        )
        .where(
            PaperTrade.status == "closed",
            PaperTrade.closed_at.is_not(None),
            PaperTrade.book == book,
        )
    )
    if closed_before is not None:
        ranked_stmt = ranked_stmt.where(PaperTrade.closed_at <= closed_before)
    ranked = ranked_stmt.subquery()
    stmt = (
        select(PaperTrade)
        .join(ranked, ranked.c.trade_id == PaperTrade.id)
        .where(
            ranked.c.recency == 1,
            PaperTrade.close_reason == "stop",
        )
    )
    result = await db.execute(stmt)
    return {trade.ticker_id: trade for trade in result.scalars()}


async def get_reentry_gate_locks(
    db: AsyncSession, *, book: str = MANUAL_BOOK
) -> dict[int, datetime]:
    """Return tickers still waiting for a post-stop gate failure.

    A later qualified setup is actionable only after the daily scanner has
    observed an unqualified evaluation after the latest initial-stop exit and
    then a fresh qualification. The returned timestamp is the stop time and is
    useful for diagnostics; callers normally only need the keys.
    """
    latest = await _latest_initial_stop_trades(db, book=book)
    return {
        ticker_id: trade.closed_at
        for ticker_id, trade in latest.items()
        if trade.reentry_gate_requalified_at is None and trade.closed_at is not None
    }


async def observe_reentry_gate_transitions(
    db: AsyncSession,
    *,
    evaluated_ticker_ids: Iterable[int],
    qualified_ticker_ids: Iterable[int],
    observed_at: datetime | None = None,
    book: str = MANUAL_BOOK,
) -> set[int]:
    """Persist gate-failure and later requalification observations.

    Only tickers whose scan completed successfully belong in
    ``evaluated_ticker_ids``. This prevents a scanner exception from being
    mistaken for a real gate exit. The caller owns the transaction; this helper
    flushes so the new state is immediately visible in that transaction.
    """
    evaluated = {int(ticker_id) for ticker_id in evaluated_ticker_ids}
    if not evaluated:
        return set()
    qualified = {int(ticker_id) for ticker_id in qualified_ticker_ids}
    timestamp = observed_at or datetime.now(timezone.utc)
    latest = await _latest_initial_stop_trades(db, closed_before=timestamp, book=book)
    updated: set[int] = set()
    for ticker_id in evaluated:
        trade = latest.get(ticker_id)
        if trade is None or trade.reentry_gate_requalified_at is not None:
            continue
        if trade.reentry_gate_failed_at is None:
            if ticker_id not in qualified:
                trade.reentry_gate_failed_at = timestamp
                updated.add(ticker_id)
        elif ticker_id in qualified:
            # Study semantics: requalify only on a *subsequent* daily observation.
            # Same America/New_York calendar day as the failure does not unlock,
            # even if multiple full-universe scans run (manual + near-close).
            if _ny_trading_date(trade.reentry_gate_failed_at) < _ny_trading_date(
                timestamp
            ):
                trade.reentry_gate_requalified_at = timestamp
                updated.add(ticker_id)

    if updated:
        await db.flush()
    return updated