from typing import Annotated

from advanced_alchemy.extensions.litestar.plugins.init.config.engine import EngineConfig
from litestar import Litestar, get
from litestar.contrib.sqlalchemy.dto import SQLAlchemyDTO
from litestar.contrib.sqlalchemy.plugins import (
    SQLAlchemyPlugin,
    SQLAlchemySyncConfig,
)
from litestar.dto import DTOConfig
from litestar.testing import TestClient
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column


class Base(DeclarativeBase):
    pass


class Address(Base):
    __tablename__ = "addresses"

    id: Mapped[int] = mapped_column(primary_key=True)
    street: Mapped[str]
    city: Mapped[str]
    state: Mapped[str]
    zip: Mapped[str]


db_config = SQLAlchemySyncConfig(
    connection_string="sqlite://",
    engine_config=EngineConfig(echo=False),
    create_all=True,
)


AddressDTO = SQLAlchemyDTO[
    Annotated[
        Address,
        DTOConfig(exclude={"zip"}),
    ]
]


@get("/address", dto=AddressDTO, sync_to_thread=False)
def get_address() -> Address:
    return Address(id=1, street="123 Main St", city="Anytown", state="NY", zip="12345")


app_with_plugin = Litestar(
    route_handlers=[get_address],
    plugins=[SQLAlchemyPlugin(config=db_config)],
)

app_without_plugin = Litestar(
    route_handlers=[get_address],
)

with TestClient(app=app_without_plugin) as client:
    if "zip" not in client.get("/address").json():
        print("'zip' is properly excluded without the plugin")

with TestClient(app=app_with_plugin) as client:
    if "zip" in client.get("/address").json():
        print("'zip' is not excluded with the plugin")