# Code generated by sqlc. DO NOT EDIT.
# versions:
#   sqlc v1.29.0
# source: update_executor_stats.sql
import dataclasses
import datetime
from typing import Any, Optional
import uuid

import sqlalchemy
import sqlalchemy.ext.asyncio

from . import models


UPDATE_EXECUTOR_STATS = """-- name: update_executor_stats \\:one
WITH updated AS (
    UPDATE hyrex_executor
        SET last_heartbeat = CURRENT_TIMESTAMP,
            stats = :p1
        WHERE hyrex_executor.id = :p2
            AND hyrex_executor.status = 'RUNNING'\\:\\:executor_status
        RETURNING id, status, last_heartbeat, stats),
     executor_state AS (SELECT CASE WHEN u.id IS NOT NULL THEN 'ACCEPTED' ELSE 'REJECTED' END AS result,
                               COALESCE(u.status, e.status)\\:\\:text                             AS status,
                               COALESCE(u.last_heartbeat, e.last_heartbeat)                   AS last_heartbeat,
                               COALESCE(u.stats, e.stats)                                     AS stats
                        FROM hyrex_executor e
                                 LEFT JOIN updated u ON e.id = u.id
                        WHERE e.id = :p2),
     insert_log AS (
         INSERT INTO hyrex_system_logs (id, timestamp, event_name, event_body)
             SELECT gen_random_uuid(),
                    CURRENT_TIMESTAMP,
                    'HEARTBEAT_REJECTED',
                    json_build_object('executor_id', :p2, 'current_status', es.status, 'stats', :p1)
             FROM executor_state es
             WHERE es.result = 'REJECTED'
             RETURNING NULL)
SELECT es.result, es.status, es.last_heartbeat, es.stats
FROM executor_state es
"""


@dataclasses.dataclass()
class UpdateExecutorStatsParams:
    stats: Optional[Any]
    id: uuid.UUID


@dataclasses.dataclass()
class UpdateExecutorStatsRow:
    result: str
    status: str
    last_heartbeat: Optional[datetime.datetime]
    stats: Optional[Any]


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

    def update_executor_stats(self, arg: UpdateExecutorStatsParams) -> Optional[UpdateExecutorStatsRow]:
        row = self._conn.execute(sqlalchemy.text(UPDATE_EXECUTOR_STATS), {"p1": arg.stats, "p2": arg.id}).first()
        if row is None:
            return None
        return UpdateExecutorStatsRow(
            result=row[0],
            status=row[1],
            last_heartbeat=row[2],
            stats=row[3],
        )


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

    async def update_executor_stats(self, arg: UpdateExecutorStatsParams) -> Optional[UpdateExecutorStatsRow]:
        row = (await self._conn.execute(sqlalchemy.text(UPDATE_EXECUTOR_STATS), {"p1": arg.stats, "p2": arg.id})).first()
        if row is None:
            return None
        return UpdateExecutorStatsRow(
            result=row[0],
            status=row[1],
            last_heartbeat=row[2],
            stats=row[3],
        )
