from __future__ import annotations
import json, logging

logger = logging.getLogger(__name__)


class KnowledgeStore:
    def __init__(self, pool):
        self._pool = pool

    async def upsert(self, entity_type, entity_key, facts):
        async with self._pool.acquire(timeout=10) as conn:
            await conn.execute(
                """INSERT INTO ab_knowledge_store (entity_type, entity_key, facts)
                   VALUES ($1, $2, $3)
                   ON CONFLICT (entity_type, entity_key)
                   DO UPDATE SET facts = $3, updated_at = NOW()""",
                entity_type, entity_key, json.dumps(facts))

    async def get(self, entity_type, entity_key):
        async with self._pool.acquire(timeout=10) as conn:
            row = await conn.fetchrow(
                'SELECT facts FROM ab_knowledge_store WHERE entity_type=$1 AND entity_key=$2',
                entity_type, entity_key)
        if not row:
            return {}
        f = row['facts']
        return json.loads(f) if isinstance(f, str) else f

    async def get_client_profile(self, partner_id):
        return await self.get('client', f'partner_{partner_id}')