import uuid
from typing import Any
from uuid import UUID

from sqlalchemy import JSON, create_engine, text
from sqlalchemy.ext.mutable import MutableDict
from sqlmodel import Field, Session, SQLModel, select


class Foo(SQLModel, table=True):
    id: UUID = Field(primary_key=True, default_factory=uuid.uuid4)
    foo: str
    meta: dict[str, Any] = Field(sa_type=MutableDict.as_mutable(JSON))


engine = create_engine("sqlite:///:memory:", echo=True)


SQLModel.metadata.create_all(engine)


with Session(engine) as session:
    session.add(
        Foo(
            foo="Bar",
            meta={"hello": "world", "foo": {"bar": 42}},
        )
    )
    session.commit()

    # Ensure row is available
    foo = session.exec(select(Foo)).first()
    print(">>> first row:", foo)

    # Query by JSON-Content
    foo = session.exec(select(Foo).where(Foo.meta["hello"] == "world")).first()
    print(">>> select SQLModel:", foo)

    # Query like created by the expression above, doesn't work as well
    foo = session.execute(text("SELECT foo.id, foo.foo, foo.meta FROM foo "
                          "WHERE JSON_QUOTE(JSON_EXTRACT(foo.meta, '$.\"hello\"')) = 'world'")).first()
    print(">>> select text:", foo)

    # removed JSON_QUOTE - now it works!
    foo = session.execute(text("SELECT foo.id, foo.foo, foo.meta FROM foo "
                          "WHERE JSON_EXTRACT(foo.meta, '$.\"hello\"') = 'world'")).first()
    print(">>> select text without JSON_QUOTE:", foo)

session.exec(select(Foo).where(Foo.meta["hello"] == "world")).first()    # => None

session.exec(select(Foo)).first()      # => meta={'hello': 'world', 'foo': {'bar': 42}} id=UUID('3511f27f-266e-42e5-a200-e80a8abef4e9') foo='Bar'

session.exec(select(Foo).where(Foo.meta["hello"].as_string() == "world")).first()