import asyncio
import logging
from datetime import timedelta

import asyncpg

from utils.files import media_file_name
from utils.read.media import get_media
from utils.read.models import TrashOutcome
from utils.storage import ContentAddressedStorage

logger = logging.getLogger(__name__)

RETENTION = timedelta(days=7)
PURGE_INTERVAL_SECONDS = 3600
BATCH = 200

_SHARED_ELSEWHERE = """
SELECT EXISTS (SELECT 1 FROM stories WHERE storage_key = $1)
    OR EXISTS (SELECT 1 FROM avatars WHERE storage_key = $1)
    OR EXISTS (SELECT 1 FROM custom_emoji WHERE storage_key = $1)
"""

_EVICT = """
UPDATE media SET storage_key = NULL, downloaded = false,
    evicted_at = now(), evicted_key = $2
WHERE storage_key = $1
"""

_INSERT = """
INSERT INTO media_trash
    (storage_key, account_id, chat_id, message_id, kind, mime, file_name, file_size)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
ON CONFLICT (storage_key) DO UPDATE SET trashed_at = now()
"""

_RESTORE = """
UPDATE media SET storage_key = $1, downloaded = true,
    evicted_at = NULL, evicted_key = NULL
WHERE evicted_key = $1 OR (
    evicted_at IS NOT NULL AND unique_id IN (
        SELECT unique_id FROM media
        WHERE evicted_key = $1 AND unique_id IS NOT NULL
    )
)
"""

_FORGET_KEYS = """
UPDATE media SET evicted_key = NULL WHERE evicted_key = ANY($1::text[])
"""

_EVICT_MANY = """
UPDATE media SET evicted_key = storage_key, storage_key = NULL,
    downloaded = false, evicted_at = now()
WHERE storage_key = ANY($1::text[])
"""


async def move_to_trash(
    pool: asyncpg.Pool,
    storage: ContentAddressedStorage,
    account_id: int,
    media_ids: list[int],
) -> TrashOutcome:
    outcome = TrashOutcome()
    for media_id in media_ids:
        media = await get_media(pool, media_id)
        if media is None or media.account_id != account_id or not media.storage_key:
            continue
        key = media.storage_key
        on_disk = storage.exists(key)
        async with pool.acquire() as conn, conn.transaction():
            if await conn.fetchval(_SHARED_ELSEWHERE, key):
                outcome.blocked.append(media_id)
                continue
            await conn.execute(_EVICT, key, key if on_disk else None)
            if not on_disk:
                continue
            await conn.execute(
                _INSERT,
                key,
                account_id,
                media.chat_id,
                media.message_id,
                media.kind,
                media.mime,
                media.file_name,
                media.file_size,
            )
            storage.trash(key)
        outcome.keys.append(key)
        outcome.files += 1
        outcome.bytes += media.file_size or 0
    return outcome


def _move_batch(storage: ContentAddressedStorage, keys: list[str]) -> None:
    moved: list[str] = []
    try:
        for key in keys:
            if storage.exists(key):
                storage.trash(key)
                moved.append(key)
    except OSError:
        for key in moved:
            storage.restore(key)
        raise


async def trash_blobs(
    pool: asyncpg.Pool,
    storage: ContentAddressedStorage,
    account_id: int,
    blobs: list[asyncpg.Record],
) -> TrashOutcome:
    outcome = TrashOutcome()
    for start in range(0, len(blobs), BATCH):
        batch = blobs[start : start + BATCH]
        keys = [blob["storage_key"] for blob in batch]
        async with pool.acquire() as conn, conn.transaction():
            await conn.execute(_EVICT_MANY, keys)
            await conn.executemany(
                _INSERT,
                [
                    (
                        blob["storage_key"],
                        account_id,
                        blob["chat_id"],
                        blob["message_id"],
                        blob["kind"],
                        blob["mime"],
                        media_file_name(blob["kind"], blob["mime"], blob["message_id"]),
                        blob["file_size"],
                    )
                    for blob in batch
                ],
            )
            await asyncio.to_thread(_move_batch, storage, keys)
        outcome.keys.extend(keys)
        outcome.files += len(batch)
        outcome.bytes += sum(blob["file_size"] or 0 for blob in batch)
    return outcome


async def restore(
    pool: asyncpg.Pool,
    storage: ContentAddressedStorage,
    account_id: int,
    keys: list[str] | None,
) -> TrashOutcome:
    if keys is None:
        rows = await pool.fetch(
            "SELECT storage_key FROM media_trash WHERE account_id = $1", account_id
        )
        keys = [row["storage_key"] for row in rows]
    outcome = TrashOutcome()
    for key in keys:
        async with pool.acquire() as conn, conn.transaction():
            row = await conn.fetchrow(
                "DELETE FROM media_trash WHERE storage_key = $1 AND account_id = $2 "
                "RETURNING file_size",
                key,
                account_id,
            )
            if row is None:
                continue
            if not storage.restore(key):
                await conn.execute(_FORGET_KEYS, [key])
                continue
            await conn.execute(_RESTORE, key)
        outcome.files += 1
        outcome.bytes += row["file_size"] or 0
    return outcome


async def _purge(
    pool: asyncpg.Pool, storage: ContentAddressedStorage, rows: list[asyncpg.Record]
) -> TrashOutcome:
    keys = [row["storage_key"] for row in rows]
    for key in keys:
        storage.purge(key)
    async with pool.acquire() as conn, conn.transaction():
        await conn.execute(
            "DELETE FROM media_trash WHERE storage_key = ANY($1::text[])", keys
        )
        await conn.execute(_FORGET_KEYS, keys)
    return TrashOutcome(
        files=len(rows), bytes=sum(row["file_size"] or 0 for row in rows)
    )


async def purge(
    pool: asyncpg.Pool,
    storage: ContentAddressedStorage,
    account_id: int,
    keys: list[str] | None,
) -> TrashOutcome:
    rows = await pool.fetch(
        "SELECT storage_key, file_size FROM media_trash WHERE account_id = $1 "
        "AND ($2::text[] IS NULL OR storage_key = ANY($2))",
        account_id,
        keys,
    )
    return await _purge(pool, storage, rows)


async def purge_expired(
    pool: asyncpg.Pool, storage: ContentAddressedStorage
) -> TrashOutcome:
    rows = await pool.fetch(
        "SELECT storage_key, file_size FROM media_trash "
        "WHERE trashed_at < now() - $1::interval",
        RETENTION,
    )
    return await _purge(pool, storage, rows)


async def purge_loop(pool: asyncpg.Pool, storage: ContentAddressedStorage) -> None:
    while True:
        try:
            outcome = await purge_expired(pool, storage)
            if outcome.files:
                logger.info(
                    "Trash purged: %s files, %s bytes", outcome.files, outcome.bytes
                )
        except Exception:
            logger.exception("Trash purge failed")
        await asyncio.sleep(PURGE_INTERVAL_SECONDS)