import uuid
from typing import Literal, Any, Type, Callable

from pydantic import BaseModel
from sqlalchemy.dialects.postgresql import JSONB
from sqlmodel import Field, Column

from common.database.base_models import BaseSQLModel


class KVStoreMixin:
    key: str
    value: str | dict | list | None

    @classmethod
    async def key_exists(cls, key: str) -> bool:
        count = await cls.count(key=key)  # noqa
        return count > 0

    @classmethod
    async def get_value(
        cls,
        key: str,
        if_not_found: Literal["raise", "none"] = "none",
        default: Any = None,
    ):
        model = await cls.get(key=key)  # noqa
        if model is None and if_not_found == "raise":
            raise ValueError(f"Key '{key}' not found in KV store.")
        if model is None:
            return default
        return model.value

    @classmethod
    async def set_value(cls, key: str, value: Any):
        """Set or update a value in the KV store."""
        existing = await cls.get(key=key)  # noqa
        if existing is not None:
            # Update existing entry
            await existing.update(value=value)  # noqa
            return existing
        else:
            # Create new entry
            return await cls.create(key=key, value=value)  # noqa


class SimpleKVStoreSQLModel(KVStoreMixin, BaseSQLModel):
    __tablename__ = "kv_store"
    id: str = Field(default_factory=lambda: f"skv_{uuid.uuid4().hex}", primary_key=True)
    key: str = Field(index=True, unique=True)
    value: str | None = Field(default=None)


class JSONKVStoreSQLModel(KVStoreMixin, BaseSQLModel):
    __tablename__ = "kv_store"
    id: str = Field(default_factory=lambda: f"jkv_{uuid.uuid4().hex}", primary_key=True)
    key: str = Field(index=True, unique=True)
    value: dict | list | None = Field(sa_column=Column(JSONB), default=None)

    @classmethod
    async def validate_value(
        cls,
        key: str,
        model: Type[BaseModel] | None = None,
        if_not_found: Literal["raise", "none"] = "none",
        default_factory: Callable = lambda: dict(),
    ):
        value = await cls.get_value(key=key, if_not_found=if_not_found, default=None)
        if value is None and if_not_found == "raise":
            raise ValueError(f"Key '{key}' not found in KV store.")
        if value is None:
            value = default_factory()
        if model is not None:
            return model.model_validate(value)
        return value

    @classmethod
    async def set_value(cls, key: str, value: dict | list | BaseModel | None):
        """Set or update a value in the KV store. Accepts dict, list, or Pydantic models."""
        # If value is a Pydantic model, serialize it to dict
        if isinstance(value, BaseModel):
            value = value.model_dump()
        
        # Call parent implementation
        return await super().set_value(key=key, value=value)