from __future__ import annotations
import logging

logger = logging.getLogger(__name__)


class LLMConfigStore:
    def __init__(self, pg_pool):
        self._pool = pg_pool

    async def get_backend(self, caller):
        async with self._pool.acquire(timeout=10) as conn:
            row = await conn.fetchrow(
                'SELECT backend FROM ab_llm_config WHERE caller = $1', caller)
            return row['backend'] if row else None

    async def set_backend(self, caller, backend, set_by, note=None):
        async with self._pool.acquire(timeout=10) as conn:
            await conn.execute(
                """INSERT INTO ab_llm_config (caller, backend, set_by, note)
                   VALUES ($1, $2, $3, $4)
                   ON CONFLICT (caller) DO UPDATE
                     SET backend=$2, set_by=$3, set_at=NOW(), note=$4""",
                caller, backend, set_by, note)
        logger.info('LLMConfigStore.set_backend caller=%s backend=%s set_by=%s', caller, backend, set_by)

    async def get_all(self):
        async with self._pool.acquire(timeout=10) as conn:
            rows = await conn.fetch('SELECT * FROM ab_llm_config ORDER BY caller')
            return [dict(r) for r in rows]

    async def reset(self, caller):
        async with self._pool.acquire(timeout=10) as conn:
            await conn.execute('DELETE FROM ab_llm_config WHERE caller = $1', caller)
        logger.info('LLMConfigStore.reset caller=%s', caller)