"""Shared live/backtest trading-policy constants and availability checks."""

from __future__ import annotations

from collections import defaultdict
from datetime import date, datetime, timezone

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

from app.models.benchmark_price import BenchmarkPrice
from app.models.ohlcv import OHLCVRecord
from app.models.paper_trade import PaperTrade
from app.services.benchmark_service import BENCHMARK_SYMBOL

# A ticker stopped at its initial stop may qualify again immediately, but the
# July 2026 event study showed that waiting five market sessions materially
# improved the production book. The stop session is wait_session=0; the first
# permitted re-entry is wait_session=5, provided the normal gate still passes.
REENTRY_LOCKDOWN_SESSIONS = 5


async def get_reentry_lockdowns(
    db: AsyncSession,
    *,
    as_of: date | None = None,
    sessions: int = REENTRY_LOCKDOWN_SESSIONS,
) -> dict[int, int]:
    """Return ``{ticker_id: remaining_sessions}`` for active lockdowns.

    SPY is the canonical calendar for the platform's US-equity universe. When
    the stored benchmark history does not reach an older stop, only that
    ticker's own OHLCV dates are used as a conservative fallback. Unrelated
    ticker dates can therefore never shorten a lockdown.
    """
    sessions = max(0, int(sessions))
    if sessions == 0:
        return {}

    session_cutoff = as_of or datetime.now(timezone.utc).date()
    stop_result = await db.execute(
        select(
            PaperTrade.ticker_id,
            func.max(PaperTrade.closed_at).label("last_stop_at"),
        )
        .where(
            PaperTrade.status == "closed",
            PaperTrade.close_reason == "stop",
            PaperTrade.closed_at.is_not(None),
        )
        .group_by(PaperTrade.ticker_id)
    )
    stop_dates = {
        ticker_id: stopped_at.date()
        for ticker_id, stopped_at in stop_result.all()
        if stopped_at is not None and stopped_at.date() <= session_cutoff
    }
    if not stop_dates:
        return {}

    benchmark_result = await db.execute(
        select(BenchmarkPrice.date)
        .where(
            BenchmarkPrice.symbol == BENCHMARK_SYMBOL,
            BenchmarkPrice.date <= session_cutoff,
        )
        .order_by(BenchmarkPrice.date.asc())
    )
    benchmark_dates = [row[0] for row in benchmark_result.all()]

    lockdowns: dict[int, int] = {}
    fallback_stops: dict[int, date] = {}
    first_benchmark_date = benchmark_dates[0] if benchmark_dates else None
    for ticker_id, stop_date in stop_dates.items():
        completed = sum(day > stop_date for day in benchmark_dates)
        if completed >= sessions:
            continue
        if first_benchmark_date is not None and first_benchmark_date <= stop_date:
            lockdowns[ticker_id] = sessions - completed
        else:
            # The benchmark table starts after this stop (or is empty), so it
            # cannot prove how many sessions elapsed. Resolve only this ticker
            # against its own bars instead of using universe-wide dates.
            fallback_stops[ticker_id] = stop_date

    if fallback_stops:
        own_session_result = await db.execute(
            select(OHLCVRecord.ticker_id, OHLCVRecord.date)
            .where(
                OHLCVRecord.ticker_id.in_(fallback_stops),
                OHLCVRecord.date > min(fallback_stops.values()),
                OHLCVRecord.date <= session_cutoff,
            )
            .distinct()
        )
        own_dates: dict[int, set[date]] = defaultdict(set)
        for ticker_id, market_date in own_session_result.all():
            own_dates[ticker_id].add(market_date)
        for ticker_id, stop_date in fallback_stops.items():
            completed = sum(day > stop_date for day in own_dates[ticker_id])
            if completed < sessions:
                lockdowns[ticker_id] = sessions - completed

    return lockdowns


async def get_reentry_lockdown_ticker_ids(
    db: AsyncSession,
    *,
    as_of: date | None = None,
    sessions: int = REENTRY_LOCKDOWN_SESSIONS,
) -> set[int]:
    """Compatibility wrapper for callers that only need blocked ticker ids."""
    return set(
        await get_reentry_lockdowns(
            db,
            as_of=as_of,
            sessions=sessions,
        )
    )