import json
from operator import attrgetter

from asyncpg import Connection
from smb.common.pgswim import SwimEngine

from maps_adv.geosmb.scenarist.server.lib.enums import (
    MessageType,
    ScenarioName,
    SubscriptionStatus,
)


class DB(SwimEngine):
    @staticmethod
    async def _set_codecs(con: Connection):
        await con.set_type_codec(
            "jsonb", schema="pg_catalog", encoder=json.dumps, decoder=json.loads
        )
        await con.set_type_codec(
            "scenario_name",
            encoder=attrgetter("name"),
            decoder=lambda x: ScenarioName[x],
        )
        await con.set_type_codec(
            "subscription_status",
            encoder=attrgetter("name"),
            decoder=lambda x: SubscriptionStatus[x],
        )
        await con.set_type_codec(
            "messagetype", encoder=attrgetter("name"), decoder=lambda x: MessageType[x]
        )
