from __future__ import annotations

from typing import TYPE_CHECKING, ClassVar

from sqlalchemy.ext.asyncio import AsyncSession

if TYPE_CHECKING:
    pass

import sqlalchemy as sa
from advanced_alchemy.extensions.litestar import (
    AsyncSessionConfig,
    SQLAlchemyAsyncConfig,
    SQLAlchemyDTO,
    SQLAlchemyDTOConfig,
    SQLAlchemyPlugin,
    base,
)
from litestar.app import Litestar
from litestar.handlers import get
from sqlalchemy.orm import mapped_column, relationship
from sqlalchemy.orm.base import Mapped

# ------------------------------------------------------------
# DB Config
# ------------------------------------------------------------

session_config = AsyncSessionConfig(expire_on_commit=False)
alchemy_config = SQLAlchemyAsyncConfig(
    connection_string="sqlite+aiosqlite:///:memory:",
    before_send_handler="autocommit",
    session_config=session_config,
    create_all=True,
)
alchemy = SQLAlchemyPlugin(config=alchemy_config)

# ------------------------------------------------------------
# User and Group models
# ------------------------------------------------------------

class User(base.BigIntBase):
    __tablename__ = "user"
    name: Mapped[str] = mapped_column(sa.String())
    group_id: Mapped[int | None] = mapped_column(sa.ForeignKey("group.id"), nullable=True)
    group: Mapped["GroupMembership | None"] = relationship(back_populates="user")

class GroupMembership(base.BigIntBase):
    __tablename__ = "group"
    group_name: Mapped[str] = mapped_column(sa.String())
    user: Mapped[User | None] = relationship(back_populates="group")

# ------------------------------------------------------------
# User DTO and read route
# ------------------------------------------------------------

class UserRead(SQLAlchemyDTO[User]):
  config: ClassVar[SQLAlchemyDTOConfig] = SQLAlchemyDTOConfig()

@get(path="/user", return_dto=UserRead)
async def get_user(db_session: AsyncSession) -> User:
    user = await db_session.get(User, 1)
    return user

# ------------------------------------------------------------
# App
# ------------------------------------------------------------

async def on_startup(app: Litestar) -> None:
    async with alchemy_config.get_session() as db_session:
        db_session.add(User(name="User 1"))
        await db_session.commit()


app = Litestar([get_user], plugins=[alchemy], on_startup=[on_startup], debug=True)

from __future__ import annotations

from typing import TYPE_CHECKING, Any, ClassVar

if TYPE_CHECKING:
    from litestar.types import Dependencies

import sqlalchemy as sa
from advanced_alchemy.extensions.litestar import (
    AsyncSessionConfig,
    SQLAlchemyAsyncConfig,
    SQLAlchemyDTO,
    SQLAlchemyDTOConfig,
    SQLAlchemyPlugin,
    base,
)
from advanced_alchemy.extensions.litestar.providers import create_service_provider
from advanced_alchemy.repository import SQLAlchemyAsyncRepository
from advanced_alchemy.service import SQLAlchemyAsyncRepositoryService
from litestar import Controller, Request
from litestar.app import Litestar
from litestar.connection import ASGIConnection
from litestar.di import Provide
from litestar.handlers import get
from litestar.response import Response
from litestar.security.jwt import JWTCookieAuth, Token
from sqlalchemy.orm import joinedload, mapped_column, relationship
from sqlalchemy.orm.base import Mapped

session_config = AsyncSessionConfig(expire_on_commit=False)
alchemy_config = SQLAlchemyAsyncConfig(
    connection_string="sqlite+aiosqlite:///:memory:",
    before_send_handler="autocommit",
    session_config=session_config,
    create_all=True,
)
alchemy = SQLAlchemyPlugin(config=alchemy_config)

# ------------------------------------------------------------
# User and Group Model and Service
# ------------------------------------------------------------

class User(base.BigIntBase):
    __tablename__ = "user"
    name: Mapped[str] = mapped_column(sa.String())
    group: Mapped["GroupMembership | None"] = relationship(back_populates="user")

class GroupMembership(base.BigIntBase):
    __tablename__ = "group"
    group_name: Mapped[str] = mapped_column(sa.String())
    user_id: Mapped[int] = mapped_column(sa.ForeignKey("user.id"))
    user: Mapped[User | None] = relationship(back_populates="group")

class UserService(SQLAlchemyAsyncRepositoryService[User]):
    class UserRepository(SQLAlchemyAsyncRepository[User]):
        model_type: type[User] = User

    repository_type: type[UserRepository] = UserRepository

# ------------------------------------------------------------
# Dependencies
# ------------------------------------------------------------

provide_user_service = create_service_provider(
    UserService,
    load=[joinedload(User.group)]
)

async def current_user_dependency(request: Request[User, Any, Any]) -> User:
    return request.user

# ------------------------------------------------------------
# DTO and User Controller
# ------------------------------------------------------------

class UserRead(SQLAlchemyDTO[User]):
  config: ClassVar[SQLAlchemyDTOConfig] = SQLAlchemyDTOConfig()

class UserController(Controller):
    dependencies: Dependencies | None = {
        "user_service": Provide(provide_user_service),
        "current_user": Provide(current_user_dependency),
    }
    path: str = "/user"

    @get(path="/me", return_dto=UserRead)
    async def get_user(self, current_user: User) -> User:
        return current_user

class LoginController(Controller):
    dependencies: Dependencies | None = {"user_service": Provide(provide_user_service)}
    path = "/login"

    @get(path="/", exclude_from_auth=True)
    async def login(self, user_service: UserService) -> Response[dict[str, str]]:
        user = await user_service.get_one(id=1)
        return auth.login(identifier=str(user.id), response_body={"message": "logged in"})

# ------------------------------------------------------------
# Guards
# -----------------------------------------------------------
async def get_current_user_from_token(token: Token, connection: ASGIConnection) -> User | None:
    service_provider = provide_user_service(
        alchemy_config.provide_session(connection.app.state, connection.scope)
    )
    async for service in service_provider:
        user = await service.get_one_or_none(id=int(token.sub))
        return user if user else None
    return None

auth = JWTCookieAuth[User](
    retrieve_user_handler=get_current_user_from_token,
    token_secret="secret",
    exclude=["/login"],
)

# ------------------------------------------------------------
# App
# -----------------------------------------------------------

async def on_startup(app: Litestar) -> None:
    async with alchemy_config.get_session() as db_session:
        db_session.add(User(name="User 1"))
        await db_session.commit()


app = Litestar(
    [UserController, LoginController],
    plugins=[alchemy],
    on_startup=[on_startup],
    debug=True,
    on_app_init=[auth.on_app_init],
)

from litestar.testing import TestClient

from app import app

with TestClient(app) as client:
    login_response = client.get("/login")
    print(f"Login: {login_response.status_code}")
    
    me_response = client.get("/user/me")
    print(f"/user/me body: {me_response.text}")