# Code generated by sqlc. DO NOT EDIT.
# versions:
#   sqlc v1.29.0
# source: task.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_TASK = """-- name: create_task \\:one
INSERT INTO "task" (benchmark_id,
                    name,
                    category,
                    question,
                    data,
                    is_public)
VALUES (:p1, :p2, :p3, :p4, :p5, :p6)
RETURNING id, benchmark_id, name, category, question, data, is_public, created_at, updated_at
"""


@dataclasses.dataclass()
class CreateTaskParams:
    benchmark_id: int
    name: str
    category: Optional[str]
    question: Optional[str]
    data: Optional[str]
    is_public: Optional[bool]


DELETE_TASK = """-- name: delete_task \\:exec
DELETE
FROM "task"
WHERE benchmark_id = :p1 AND name = :p2
"""


GET_TASK_BY_ID = """-- name: get_task_by_id \\:one
SELECT id, benchmark_id, name, category, question, data, is_public, created_at, updated_at
FROM "task"
WHERE id = :p1
LIMIT 1
"""


GET_TASK_BY_NAME = """-- name: get_task_by_name \\:one
SELECT id, benchmark_id, name, category, question, data, is_public, created_at, updated_at
FROM "task"
WHERE benchmark_id = :p1 AND name = :p2
LIMIT 1
"""


GET_TASK_ID = """-- name: get_task_id \\:one
SELECT id
FROM "task"
WHERE benchmark_id = :p1 AND name = :p2
LIMIT 1
"""


GET_TASK_NAMES_IN_BENCHMARK = """-- name: get_task_names_in_benchmark \\:many
SELECT name
FROM "task"
WHERE benchmark_id = :p1
"""


UPDATE_TASK = """-- name: update_task \\:one
UPDATE "task"
SET category   = COALESCE(:p1, category),
    question   = COALESCE(:p2, question),
    data       = COALESCE(:p3, data),
    is_public  = COALESCE(:p4, is_public),
    updated_at = COALESCE(:p5, updated_at)
WHERE benchmark_id = :p6 AND name = :p7
RETURNING id, benchmark_id, name, category, question, data, is_public, created_at, updated_at
"""


@dataclasses.dataclass()
class UpdateTaskParams:
    category: Optional[str]
    question: Optional[str]
    data: Optional[str]
    is_public: Optional[bool]
    updated_at: Optional[datetime.datetime]
    benchmark_id: int
    name: str


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

    def create_task(self, arg: CreateTaskParams) -> Optional[models.Task]:
        row = self._conn.execute(sqlalchemy.text(CREATE_TASK), {
            "p1": arg.benchmark_id,
            "p2": arg.name,
            "p3": arg.category,
            "p4": arg.question,
            "p5": arg.data,
            "p6": arg.is_public,
        }).first()
        if row is None:
            return None
        return models.Task(
            id=row[0],
            benchmark_id=row[1],
            name=row[2],
            category=row[3],
            question=row[4],
            data=row[5],
            is_public=row[6],
            created_at=row[7],
            updated_at=row[8],
        )

    def delete_task(self, *, benchmark_id: int, name: str) -> None:
        self._conn.execute(sqlalchemy.text(DELETE_TASK), {"p1": benchmark_id, "p2": name})

    def get_task_by_id(self, *, id: int) -> Optional[models.Task]:
        row = self._conn.execute(sqlalchemy.text(GET_TASK_BY_ID), {"p1": id}).first()
        if row is None:
            return None
        return models.Task(
            id=row[0],
            benchmark_id=row[1],
            name=row[2],
            category=row[3],
            question=row[4],
            data=row[5],
            is_public=row[6],
            created_at=row[7],
            updated_at=row[8],
        )

    def get_task_by_name(self, *, benchmark_id: int, name: str) -> Optional[models.Task]:
        row = self._conn.execute(sqlalchemy.text(GET_TASK_BY_NAME), {"p1": benchmark_id, "p2": name}).first()
        if row is None:
            return None
        return models.Task(
            id=row[0],
            benchmark_id=row[1],
            name=row[2],
            category=row[3],
            question=row[4],
            data=row[5],
            is_public=row[6],
            created_at=row[7],
            updated_at=row[8],
        )

    def get_task_id(self, *, benchmark_id: int, name: str) -> Optional[int]:
        row = self._conn.execute(sqlalchemy.text(GET_TASK_ID), {"p1": benchmark_id, "p2": name}).first()
        if row is None:
            return None
        return row[0]

    def get_task_names_in_benchmark(self, *, benchmark_id: int) -> Iterator[str]:
        result = self._conn.execute(sqlalchemy.text(GET_TASK_NAMES_IN_BENCHMARK), {"p1": benchmark_id})
        for row in result:
            yield row[0]

    def update_task(self, arg: UpdateTaskParams) -> Optional[models.Task]:
        row = self._conn.execute(sqlalchemy.text(UPDATE_TASK), {
            "p1": arg.category,
            "p2": arg.question,
            "p3": arg.data,
            "p4": arg.is_public,
            "p5": arg.updated_at,
            "p6": arg.benchmark_id,
            "p7": arg.name,
        }).first()
        if row is None:
            return None
        return models.Task(
            id=row[0],
            benchmark_id=row[1],
            name=row[2],
            category=row[3],
            question=row[4],
            data=row[5],
            is_public=row[6],
            created_at=row[7],
            updated_at=row[8],
        )


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

    async def create_task(self, arg: CreateTaskParams) -> Optional[models.Task]:
        row = (await self._conn.execute(sqlalchemy.text(CREATE_TASK), {
            "p1": arg.benchmark_id,
            "p2": arg.name,
            "p3": arg.category,
            "p4": arg.question,
            "p5": arg.data,
            "p6": arg.is_public,
        })).first()
        if row is None:
            return None
        return models.Task(
            id=row[0],
            benchmark_id=row[1],
            name=row[2],
            category=row[3],
            question=row[4],
            data=row[5],
            is_public=row[6],
            created_at=row[7],
            updated_at=row[8],
        )

    async def delete_task(self, *, benchmark_id: int, name: str) -> None:
        await self._conn.execute(sqlalchemy.text(DELETE_TASK), {"p1": benchmark_id, "p2": name})

    async def get_task_by_id(self, *, id: int) -> Optional[models.Task]:
        row = (await self._conn.execute(sqlalchemy.text(GET_TASK_BY_ID), {"p1": id})).first()
        if row is None:
            return None
        return models.Task(
            id=row[0],
            benchmark_id=row[1],
            name=row[2],
            category=row[3],
            question=row[4],
            data=row[5],
            is_public=row[6],
            created_at=row[7],
            updated_at=row[8],
        )

    async def get_task_by_name(self, *, benchmark_id: int, name: str) -> Optional[models.Task]:
        row = (await self._conn.execute(sqlalchemy.text(GET_TASK_BY_NAME), {"p1": benchmark_id, "p2": name})).first()
        if row is None:
            return None
        return models.Task(
            id=row[0],
            benchmark_id=row[1],
            name=row[2],
            category=row[3],
            question=row[4],
            data=row[5],
            is_public=row[6],
            created_at=row[7],
            updated_at=row[8],
        )

    async def get_task_id(self, *, benchmark_id: int, name: str) -> Optional[int]:
        row = (await self._conn.execute(sqlalchemy.text(GET_TASK_ID), {"p1": benchmark_id, "p2": name})).first()
        if row is None:
            return None
        return row[0]

    async def get_task_names_in_benchmark(self, *, benchmark_id: int) -> AsyncIterator[str]:
        result = await self._conn.stream(sqlalchemy.text(GET_TASK_NAMES_IN_BENCHMARK), {"p1": benchmark_id})
        async for row in result:
            yield row[0]

    async def update_task(self, arg: UpdateTaskParams) -> Optional[models.Task]:
        row = (await self._conn.execute(sqlalchemy.text(UPDATE_TASK), {
            "p1": arg.category,
            "p2": arg.question,
            "p3": arg.data,
            "p4": arg.is_public,
            "p5": arg.updated_at,
            "p6": arg.benchmark_id,
            "p7": arg.name,
        })).first()
        if row is None:
            return None
        return models.Task(
            id=row[0],
            benchmark_id=row[1],
            name=row[2],
            category=row[3],
            question=row[4],
            data=row[5],
            is_public=row[6],
            created_at=row[7],
            updated_at=row[8],
        )
