"""Monthly partition maintenance for `events` (PARTITION BY RANGE(received_at)).

Idempotent: creates current + N future month partitions and drops old ones beyond
the retention window. Names are derived from year/month, so they cannot inject SQL.
"""

from __future__ import annotations

import asyncio
from datetime import UTC, date, datetime

from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker

from ..logging_config import get_logger
from ..settings import get_settings

log = get_logger("monlet.partitioning")


def _month_start(d: date) -> date:
    return date(d.year, d.month, 1)


def _next_month(d: date) -> date:
    year = d.year + (1 if d.month == 12 else 0)
    month = 1 if d.month == 12 else d.month + 1
    return date(year, month, 1)


def _shift_months(d: date, n: int) -> date:
    y, m = d.year, d.month
    total = y * 12 + (m - 1) + n
    return date(total // 12, (total % 12) + 1, 1)


def _partition_name(d: date) -> str:
    return f"events_p{d.year:04d}_{d.month:02d}"


async def ensure_event_partitions(engine_or_sm: AsyncEngine | async_sessionmaker) -> None:
    """Create current + N future month partitions for events. Idempotent."""
    s = get_settings()
    today = _month_start(datetime.now(UTC).date())
    months = [_shift_months(today, i) for i in range(0, max(s.events_future_partitions, 0) + 1)]

    statements: list[str] = []
    for start in months:
        end = _next_month(start)
        # Explicit UTC qualifier — Postgres otherwise interprets bare date literals via
        # the session TimeZone, which would shift partition bounds on non-UTC servers.
        statements.append(
            f"CREATE TABLE IF NOT EXISTS {_partition_name(start)} "
            f"PARTITION OF events FOR VALUES FROM ('{start.isoformat()} UTC') "
            f"TO ('{end.isoformat()} UTC');"
        )

    async with _connect(engine_or_sm) as conn:
        for stmt in statements:
            await conn.execute(text(stmt))


async def drop_old_event_partitions(engine_or_sm: AsyncEngine | async_sessionmaker) -> None:
    """Drop month partitions older than events_retention_months.

    Each DROP runs in its own transaction so the ACCESS EXCLUSIVE lock on parent
    `events` is held only for one partition at a time and ingest is not blocked
    by a long maintenance batch.
    """
    s = get_settings()
    if s.events_retention_months <= 0:
        return
    cutoff = _shift_months(_month_start(datetime.now(UTC).date()), -s.events_retention_months)

    engine = _resolve_engine(engine_or_sm)
    async with engine.connect() as conn:
        rows = (
            await conn.execute(
                text(
                    """
                    SELECT c.relname
                    FROM pg_inherits i
                    JOIN pg_class c ON c.oid = i.inhrelid
                    JOIN pg_class p ON p.oid = i.inhparent
                    WHERE p.relname = 'events'
                    """
                )
            )
        ).all()

    to_drop: list[str] = []
    for (name,) in rows:
        if not name.startswith("events_p"):
            continue
        try:
            y = int(name[8:12])
            m = int(name[13:15])
        except ValueError:
            continue
        if date(y, m, 1) < cutoff:
            to_drop.append(name)

    for name in to_drop:
        async with engine.begin() as conn:
            await conn.execute(text(f"DROP TABLE IF EXISTS {name};"))
        log.info("event_partition_dropped", name=name)


def _resolve_engine(engine_or_sm: AsyncEngine | async_sessionmaker) -> AsyncEngine:
    if isinstance(engine_or_sm, async_sessionmaker):
        return engine_or_sm.kw["bind"]
    return engine_or_sm


def _connect(engine_or_sm):
    """Return an async ctx for a transactional connection from either engine or sessionmaker."""
    return _resolve_engine(engine_or_sm).begin()


async def run_partition_maintenance(sm: async_sessionmaker, stop_event: asyncio.Event) -> None:
    """Independent worker — runs even when MONLET_ENABLE_DETECTOR=false."""
    interval = max(get_settings().partition_maintenance_interval_sec, 60)
    while not stop_event.is_set():
        try:
            await ensure_event_partitions(sm)
            await drop_old_event_partitions(sm)
        except Exception as exc:
            log.warning("partition_maintenance_failed", error=str(exc))
        try:
            await asyncio.wait_for(stop_event.wait(), timeout=interval)
        except TimeoutError:
            continue


# Re-export helpers for tests.
__all__ = [
    "run_partition_maintenance",
    "ensure_event_partitions",
    "drop_old_event_partitions",
    "_partition_name",
    "_shift_months",
    "_month_start",
]