"""Stored results for summaries of very large sources: a dashboard total or
group-by has to read every matching record, which at crores of records takes
longer than a web request may run. On a source past record_sql.COUNT_CAP such a
query's rows are kept for summary_cache_seconds and refreshed in the background
after that, so a dashboard answers at once, from figures at most that old plus
the time a refresh takes. A smaller source always reads live.

The key is the query's SQL with its values written in, so it holds the tenant,
source, filters and access rules: people who may see different records run
different SQL and never share a result. That SQL comes only from this app's own
query builders.
"""

import hashlib
import json
import uuid
from datetime import date, datetime, timedelta, timezone
from decimal import Decimal
from typing import Any

from fastapi import HTTPException
from sqlalchemy import text
from sqlalchemy.dialects.postgresql.asyncpg import dialect as asyncpg_dialect
from sqlalchemy.exc import CompileError, DBAPIError
from sqlalchemy.ext.asyncio import AsyncSession

from app.core import jobs
from app.core.config import settings
from app.crud import record_sql

JOB_KIND = "summary_cache.refresh"
# How long a request waits for a summary with no stored result yet, before leaving it to the background.
IN_REQUEST_MS = 8_000
# Results nobody has asked for lately are removed by the next refresh.
KEEP_FOR = timedelta(days=7)
_DIALECT = asyncpg_dialect()


def _encode(value: Any) -> Any:
    if isinstance(value, Decimal):
        return {"$": "decimal", "v": str(value)}
    if isinstance(value, datetime):
        return {"$": "datetime", "v": value.isoformat()}
    if isinstance(value, date):
        return {"$": "date", "v": value.isoformat()}
    if isinstance(value, uuid.UUID):
        return {"$": "uuid", "v": str(value)}
    return value


_DECODERS = {"decimal": Decimal, "datetime": datetime.fromisoformat, "date": date.fromisoformat, "uuid": uuid.UUID}


def _decode(value: Any) -> Any:
    # Any other object is a stored field value, grouped on as it is.
    if isinstance(value, dict) and value.keys() == {"$", "v"} and value["$"] in _DECODERS:
        return _DECODERS[value["$"]](value["v"])
    return value


def _as_sql(stmt: Any) -> str | None:
    """The query with its values written in, or None for a value type that can't be."""
    try:
        return str(stmt.compile(dialect=_DIALECT, compile_kwargs={"literal_binds": True}))
    except (CompileError, NotImplementedError, TypeError, ValueError):
        return None


async def _run_raw(db: AsyncSession, sql: str) -> list[list]:
    conn = await db.connection()
    return [list(row) for row in (await conn.exec_driver_sql(sql)).all()]


async def _store(db: AsyncSession, tenant_id: uuid.UUID, key: str, raw_rows: list[list]) -> None:
    await db.execute(
        text(
            "INSERT INTO summary_cache (tenant_id, key, rows) VALUES (:t, :k, CAST(:r AS jsonb)) "
            "ON CONFLICT (tenant_id, key) DO UPDATE SET rows = EXCLUDED.rows, computed_at = now()"
        ),
        {"t": tenant_id, "k": key, "r": json.dumps(raw_rows, default=_encode)},
    )


class _TooSlow(Exception):
    pass


def _timed_out(exc: DBAPIError) -> bool:
    return getattr(getattr(exc.orig, "__cause__", None), "sqlstate", None) == "57014"


async def _attempt(db: AsyncSession, sql: str) -> list[list]:
    """Runs it under a shorter limit, in a savepoint so a timeout leaves the request usable."""
    await db.execute(text("SELECT set_config('statement_timeout', :v, true)"), {"v": str(IN_REQUEST_MS)})
    try:
        async with db.begin_nested():
            return await _run_raw(db, sql)
    except DBAPIError as exc:
        if _timed_out(exc):
            raise _TooSlow from exc
        raise
    finally:
        await db.execute(text("SELECT set_config('statement_timeout', :v, true)"), {"v": str(settings.request_statement_timeout_ms)})


async def _queue_refresh(tenant_id: uuid.UUID, key: str, sql: str) -> None:
    """In its own transaction, so it stays queued if this request ends in an error."""
    async with jobs.tenant_session(tenant_id) as session:
        await jobs.enqueue(session, JOB_KIND, {"key": key, "sql": sql}, tenant_id=tenant_id, dedupe_key=f"{JOB_KIND}:{tenant_id}:{key}")


async def read(db: AsyncSession, tenant_id: uuid.UUID, src: record_sql.RecordSource, stmt: Any) -> list:
    """The rows of a summary query: live on a smaller source, stored on a very large one."""
    if settings.summary_cache_seconds <= 0:
        return list((await db.execute(stmt)).all())
    counts = await record_sql.capped_source_counts(db, src.kind, tenant_id, [src.source_id])
    sql = _as_sql(stmt) if counts.get(src.source_id, (0, False))[1] else None
    if sql is None:
        return list((await db.execute(stmt)).all())

    key = hashlib.sha256(sql.encode()).hexdigest()
    stored = (await db.execute(text("SELECT rows, computed_at FROM summary_cache WHERE tenant_id = :t AND key = :k"), {"t": tenant_id, "k": key})).first()
    if stored is not None:
        if datetime.now(timezone.utc) - stored.computed_at > timedelta(seconds=settings.summary_cache_seconds):
            await _queue_refresh(tenant_id, key, sql)
        return [tuple(_decode(v) for v in row) for row in stored.rows]

    # Outside a web request nothing limits the time, and a background job can wait.
    if (await db.execute(text("SHOW statement_timeout"))).scalar_one() == "0":
        raw_rows = await _run_raw(db, sql)
    else:
        try:
            raw_rows = await _attempt(db, sql)
        except _TooSlow:
            await _queue_refresh(tenant_id, key, sql)
            raise HTTPException(status_code=503, detail="This summary covers a lot of records and is being prepared — try again in a minute or two.") from None
    await _store(db, tenant_id, key, raw_rows)
    return [tuple(row) for row in raw_rows]


@jobs.handler(JOB_KIND, max_attempts=3)
async def _refresh_job(job: jobs.Job) -> str:
    async with jobs.tenant_session(job.tenant_id) as db:
        raw_rows = await _run_raw(db, job.payload["sql"])
        await _store(db, job.tenant_id, job.payload["key"], raw_rows)
        await db.execute(
            text("DELETE FROM summary_cache WHERE tenant_id = :t AND computed_at < :before"),
            {"t": job.tenant_id, "before": datetime.now(timezone.utc) - KEEP_FOR},
        )
    return f"{len(raw_rows):,} row(s)"