"""asyncpg pool + tiny migrator (plain SQL files in migrations/, tracked in schema_migrations)."""

from __future__ import annotations

import json
import re
from pathlib import Path
from typing import Any

import asyncpg

from ktrade.logging import get_logger

log = get_logger("db")
MIGRATIONS_DIR = Path(__file__).resolve().parents[2] / "migrations"


def dsn(url: str) -> str:
    return re.sub(r"^postgresql\+\w+://", "postgresql://", url)


def _json_default(o: Any) -> Any:
    import datetime as _dt
    import decimal
    import uuid

    import numpy as np

    if isinstance(o, np.bool_):
        return bool(o)
    if isinstance(o, np.integer):
        return int(o)
    if isinstance(o, np.floating):
        v = float(o)
        return v if v == v and v not in (float("inf"), float("-inf")) else None
    if isinstance(o, np.ndarray):
        return o.tolist()
    if isinstance(o, decimal.Decimal | uuid.UUID):
        return str(o)
    if isinstance(o, _dt.datetime | _dt.date):
        return o.isoformat()
    raise TypeError(f"not JSON serializable: {type(o).__name__}")


def json_dumps(o: Any) -> str:
    return json.dumps(o, default=_json_default, allow_nan=False)


async def _init_conn(conn: asyncpg.Connection) -> None:
    await conn.set_type_codec("jsonb", encoder=json_dumps, decoder=json.loads, schema="pg_catalog")
    await conn.set_type_codec("json", encoder=json_dumps, decoder=json.loads, schema="pg_catalog")


async def create_pool(url: str, min_size: int = 1, max_size: int = 10) -> asyncpg.Pool:
    return await asyncpg.create_pool(dsn(url), min_size=min_size, max_size=max_size, init=_init_conn, command_timeout=60)


async def migrate(url: str) -> list[str]:
    applied: list[str] = []
    conn = await asyncpg.connect(dsn(url))
    try:
        await conn.execute("CREATE TABLE IF NOT EXISTS schema_migrations (version TEXT PRIMARY KEY, applied_at TIMESTAMPTZ NOT NULL DEFAULT now())")
        done = {r["version"] for r in await conn.fetch("SELECT version FROM schema_migrations")}
        for f in sorted(MIGRATIONS_DIR.glob("*.sql")):
            if f.stem in done:
                continue
            async with conn.transaction():
                await conn.execute(f.read_text())
                await conn.execute("INSERT INTO schema_migrations (version) VALUES ($1)", f.stem)
            applied.append(f.stem)
            log.info("db.migrated", version=f.stem)
    finally:
        await conn.close()
    return applied


async def has_timescale(pool: asyncpg.Pool) -> bool:
    return bool(await pool.fetchval("SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname='timescaledb')"))


def rows_to_dicts(rows: list[asyncpg.Record]) -> list[dict[str, Any]]:
    return [dict(r) for r in rows]