# Code generated by sqlc. DO NOT EDIT.
# versions:
#   sqlc v1.29.0
# source: benchmark_job.sql
import dataclasses
import datetime
from typing import AsyncIterator, Iterator, Optional

import sqlalchemy
import sqlalchemy.ext.asyncio

from mcpuniverse.app.db.sqlc import models


CREATE_BENCHMARK_JOB = """-- name: create_benchmark_job \\:one
INSERT INTO "benchmark_job" (job_id,
                             owner_id,
                             benchmark_id,
                             project_id)
VALUES (:p1, :p2, :p3, :p4)
RETURNING id, job_id, owner_id, benchmark_id, project_id, status, progress, results, score, celery_id, created_at, updated_at
"""


GET_BENCHMARK_JOB_BY_ID = """-- name: get_benchmark_job_by_id \\:one
SELECT id, job_id, owner_id, benchmark_id, project_id, status, progress, results, score, celery_id, created_at, updated_at
FROM "benchmark_job"
WHERE job_id = :p1
LIMIT 1
"""


GET_BENCHMARK_JOBS = """-- name: get_benchmark_jobs \\:many
SELECT id, job_id, owner_id, benchmark_id, project_id, status, progress, results, score, celery_id, created_at, updated_at
FROM "benchmark_job"
WHERE owner_id = :p1
"""


GET_NUM_JOBS = """-- name: get_num_jobs \\:one
SELECT COUNT(*)
FROM "benchmark_job"
WHERE owner_id = :p1 AND status = :p2 AND updated_at > :p3
"""


GET_TOTAL_NUM_JOBS = """-- name: get_total_num_jobs \\:one
SELECT COUNT(*)
FROM "benchmark_job"
WHERE status = :p1 AND updated_at > :p2
"""


UPDATE_BENCHMARK_JOB = """-- name: update_benchmark_job \\:one
UPDATE "benchmark_job"
SET status     = COALESCE(:p1, status),
    progress   = COALESCE(:p2, progress),
    results    = COALESCE(:p3, results),
    score      = COALESCE(:p4, score),
    celery_id  = COALESCE(:p5, celery_id),
    updated_at = COALESCE(:p6, updated_at)
WHERE job_id = :p7
RETURNING id, job_id, owner_id, benchmark_id, project_id, status, progress, results, score, celery_id, created_at, updated_at
"""


@dataclasses.dataclass()
class UpdateBenchmarkJobParams:
    status: Optional[str]
    progress: Optional[int]
    results: Optional[str]
    score: Optional[float]
    celery_id: Optional[str]
    updated_at: Optional[datetime.datetime]
    job_id: str


class Querier:
    def __init__(self, conn: sqlalchemy.engine.Connection):
        self._conn = conn

    def create_benchmark_job(self, *, job_id: str, owner_id: int, benchmark_id: int, project_id: int) -> Optional[models.BenchmarkJob]:
        row = self._conn.execute(sqlalchemy.text(CREATE_BENCHMARK_JOB), {
            "p1": job_id,
            "p2": owner_id,
            "p3": benchmark_id,
            "p4": project_id,
        }).first()
        if row is None:
            return None
        return models.BenchmarkJob(
            id=row[0],
            job_id=row[1],
            owner_id=row[2],
            benchmark_id=row[3],
            project_id=row[4],
            status=row[5],
            progress=row[6],
            results=row[7],
            score=row[8],
            celery_id=row[9],
            created_at=row[10],
            updated_at=row[11],
        )

    def get_benchmark_job_by_id(self, *, job_id: str) -> Optional[models.BenchmarkJob]:
        row = self._conn.execute(sqlalchemy.text(GET_BENCHMARK_JOB_BY_ID), {"p1": job_id}).first()
        if row is None:
            return None
        return models.BenchmarkJob(
            id=row[0],
            job_id=row[1],
            owner_id=row[2],
            benchmark_id=row[3],
            project_id=row[4],
            status=row[5],
            progress=row[6],
            results=row[7],
            score=row[8],
            celery_id=row[9],
            created_at=row[10],
            updated_at=row[11],
        )

    def get_benchmark_jobs(self, *, owner_id: int) -> Iterator[models.BenchmarkJob]:
        result = self._conn.execute(sqlalchemy.text(GET_BENCHMARK_JOBS), {"p1": owner_id})
        for row in result:
            yield models.BenchmarkJob(
                id=row[0],
                job_id=row[1],
                owner_id=row[2],
                benchmark_id=row[3],
                project_id=row[4],
                status=row[5],
                progress=row[6],
                results=row[7],
                score=row[8],
                celery_id=row[9],
                created_at=row[10],
                updated_at=row[11],
            )

    def get_num_jobs(self, *, owner_id: int, status: Optional[str], updated_at: datetime.datetime) -> Optional[int]:
        row = self._conn.execute(sqlalchemy.text(GET_NUM_JOBS), {"p1": owner_id, "p2": status, "p3": updated_at}).first()
        if row is None:
            return None
        return row[0]

    def get_total_num_jobs(self, *, status: Optional[str], updated_at: datetime.datetime) -> Optional[int]:
        row = self._conn.execute(sqlalchemy.text(GET_TOTAL_NUM_JOBS), {"p1": status, "p2": updated_at}).first()
        if row is None:
            return None
        return row[0]

    def update_benchmark_job(self, arg: UpdateBenchmarkJobParams) -> Optional[models.BenchmarkJob]:
        row = self._conn.execute(sqlalchemy.text(UPDATE_BENCHMARK_JOB), {
            "p1": arg.status,
            "p2": arg.progress,
            "p3": arg.results,
            "p4": arg.score,
            "p5": arg.celery_id,
            "p6": arg.updated_at,
            "p7": arg.job_id,
        }).first()
        if row is None:
            return None
        return models.BenchmarkJob(
            id=row[0],
            job_id=row[1],
            owner_id=row[2],
            benchmark_id=row[3],
            project_id=row[4],
            status=row[5],
            progress=row[6],
            results=row[7],
            score=row[8],
            celery_id=row[9],
            created_at=row[10],
            updated_at=row[11],
        )


class AsyncQuerier:
    def __init__(self, conn: sqlalchemy.ext.asyncio.AsyncConnection):
        self._conn = conn

    async def create_benchmark_job(self, *, job_id: str, owner_id: int, benchmark_id: int, project_id: int) -> Optional[models.BenchmarkJob]:
        row = (await self._conn.execute(sqlalchemy.text(CREATE_BENCHMARK_JOB), {
            "p1": job_id,
            "p2": owner_id,
            "p3": benchmark_id,
            "p4": project_id,
        })).first()
        if row is None:
            return None
        return models.BenchmarkJob(
            id=row[0],
            job_id=row[1],
            owner_id=row[2],
            benchmark_id=row[3],
            project_id=row[4],
            status=row[5],
            progress=row[6],
            results=row[7],
            score=row[8],
            celery_id=row[9],
            created_at=row[10],
            updated_at=row[11],
        )

    async def get_benchmark_job_by_id(self, *, job_id: str) -> Optional[models.BenchmarkJob]:
        row = (await self._conn.execute(sqlalchemy.text(GET_BENCHMARK_JOB_BY_ID), {"p1": job_id})).first()
        if row is None:
            return None
        return models.BenchmarkJob(
            id=row[0],
            job_id=row[1],
            owner_id=row[2],
            benchmark_id=row[3],
            project_id=row[4],
            status=row[5],
            progress=row[6],
            results=row[7],
            score=row[8],
            celery_id=row[9],
            created_at=row[10],
            updated_at=row[11],
        )

    async def get_benchmark_jobs(self, *, owner_id: int) -> AsyncIterator[models.BenchmarkJob]:
        result = await self._conn.stream(sqlalchemy.text(GET_BENCHMARK_JOBS), {"p1": owner_id})
        async for row in result:
            yield models.BenchmarkJob(
                id=row[0],
                job_id=row[1],
                owner_id=row[2],
                benchmark_id=row[3],
                project_id=row[4],
                status=row[5],
                progress=row[6],
                results=row[7],
                score=row[8],
                celery_id=row[9],
                created_at=row[10],
                updated_at=row[11],
            )

    async def get_num_jobs(self, *, owner_id: int, status: Optional[str], updated_at: datetime.datetime) -> Optional[int]:
        row = (await self._conn.execute(sqlalchemy.text(GET_NUM_JOBS), {"p1": owner_id, "p2": status, "p3": updated_at})).first()
        if row is None:
            return None
        return row[0]

    async def get_total_num_jobs(self, *, status: Optional[str], updated_at: datetime.datetime) -> Optional[int]:
        row = (await self._conn.execute(sqlalchemy.text(GET_TOTAL_NUM_JOBS), {"p1": status, "p2": updated_at})).first()
        if row is None:
            return None
        return row[0]

    async def update_benchmark_job(self, arg: UpdateBenchmarkJobParams) -> Optional[models.BenchmarkJob]:
        row = (await self._conn.execute(sqlalchemy.text(UPDATE_BENCHMARK_JOB), {
            "p1": arg.status,
            "p2": arg.progress,
            "p3": arg.results,
            "p4": arg.score,
            "p5": arg.celery_id,
            "p6": arg.updated_at,
            "p7": arg.job_id,
        })).first()
        if row is None:
            return None
        return models.BenchmarkJob(
            id=row[0],
            job_id=row[1],
            owner_id=row[2],
            benchmark_id=row[3],
            project_id=row[4],
            status=row[5],
            progress=row[6],
            results=row[7],
            score=row[8],
            celery_id=row[9],
            created_at=row[10],
            updated_at=row[11],
        )
