import json
from typing import Any

import asyncpg

from utils.db.models import Job, JobStatus


def _row_to_job(row: asyncpg.Record) -> Job:
    data = dict(row)
    data["params"] = json.loads(data["params"])
    data["progress"] = json.loads(data["progress"])
    data["cursor"] = json.loads(data["cursor"]) if data["cursor"] is not None else None
    return Job(**data)


async def enqueue(
    pool: asyncpg.Pool, account_id: int, kind: str, params: dict[str, Any]
) -> int:
    return await pool.fetchval(
        "INSERT INTO jobs (account_id, kind, params) "
        "VALUES ($1, $2, $3::jsonb) RETURNING id",
        account_id,
        kind,
        json.dumps(params),
    )


async def claim_next(pool: asyncpg.Pool, account_id: int) -> Job | None:
    row = await pool.fetchrow(
        "UPDATE jobs SET status = 'running', started_at = now(), "
        "updated_at = now(), attempts = attempts + 1 "
        "WHERE id = ("
        "    SELECT id FROM jobs WHERE account_id = $1 AND status = 'pending' "
        "    AND (run_after IS NULL OR run_after <= now()) "
        "    ORDER BY id FOR UPDATE SKIP LOCKED LIMIT 1"
        ") RETURNING *",
        account_id,
    )
    return _row_to_job(row) if row else None


async def requeue_running(pool: asyncpg.Pool, account_id: int) -> int:
    result = await pool.execute(
        "UPDATE jobs SET status = 'pending', updated_at = now() "
        "WHERE account_id = $1 AND status = 'running'",
        account_id,
    )
    return int(result.split()[-1])


async def prune_finished(pool: asyncpg.Pool, account_id: int) -> int:
    result = await pool.execute(
        "DELETE FROM jobs WHERE account_id = $1 "
        "AND status IN ('done', 'canceled', 'failed') "
        "AND finished_at < now() - interval '30 days'",
        account_id,
    )
    return int(result.split()[-1])


async def save_cursor(pool: asyncpg.Pool, job_id: int, cursor: dict[str, Any]) -> None:
    await pool.execute(
        "UPDATE jobs SET cursor = $2::jsonb, updated_at = now() WHERE id = $1",
        job_id,
        json.dumps(cursor),
    )


async def report_progress(
    pool: asyncpg.Pool, job_id: int, progress: dict[str, Any]
) -> None:
    await pool.execute(
        "UPDATE jobs SET progress = $2::jsonb, updated_at = now() WHERE id = $1",
        job_id,
        json.dumps(progress),
    )


async def defer_flood_wait(pool: asyncpg.Pool, job_id: int, seconds: int) -> None:
    await pool.execute(
        "UPDATE jobs SET status = 'pending', flood_waits = flood_waits + 1, "
        "run_after = now() + make_interval(secs => $2), updated_at = now() "
        "WHERE id = $1 AND status = 'running'",
        job_id,
        seconds,
    )


async def seconds_to_next_due(pool: asyncpg.Pool, account_id: int) -> float | None:
    return await pool.fetchval(
        "SELECT extract(epoch FROM min(run_after) - now())::float FROM jobs "
        "WHERE account_id = $1 AND status = 'pending' AND run_after > now()",
        account_id,
    )


async def finish(
    pool: asyncpg.Pool, job_id: int, status: JobStatus, error: str | None = None
) -> None:
    await pool.execute(
        "UPDATE jobs SET status = $2, error = $3, finished_at = now(), "
        "updated_at = now() WHERE id = $1 AND status = 'running'",
        job_id,
        status.value,
        error,
    )


async def is_canceled(pool: asyncpg.Pool, job_id: int) -> bool:
    status = await pool.fetchval("SELECT status FROM jobs WHERE id = $1", job_id)
    return status == JobStatus.CANCELED.value


async def get_job(pool: asyncpg.Pool, job_id: int) -> Job | None:
    row = await pool.fetchrow("SELECT * FROM jobs WHERE id = $1", job_id)
    return _row_to_job(row) if row else None