from typing import TYPE_CHECKING, Any, Dict, Optional
from uuid import UUID

from sqlalchemy.schema import Column, ForeignKey
from sqlmodel import Field, Relationship, select
from sqlmodel.sql.sqltypes import GUID

from joj.horse.models.base import BaseORMModel
from joj.horse.services.db import db_session
from joj.horse.services.oauth import OAuth2Profile

if TYPE_CHECKING:
    from joj.horse.models import User


class UserOAuthAccount(BaseORMModel, table=True):  # type: ignore[call-arg]
    __tablename__ = "user_oauth_accounts"

    oauth_name: str = Field()
    access_token: str = Field()
    refresh_token: Optional[str] = Field(None, nullable=True)
    expires_at: Optional[int] = Field(None, nullable=True)
    account_id: str = Field(index=True)
    account_name: Optional[str] = Field(None, index=True, nullable=True)
    account_email: str = Field(index=True)

    user_id: Optional[UUID] = Field(
        sa_column=Column(GUID, ForeignKey("users.id", ondelete="CASCADE"))
    )
    user: Optional["User"] = Relationship(back_populates="oauth_accounts")

    @staticmethod
    async def create_or_update(
        oauth_name: str, token: Dict[str, Any], profile: OAuth2Profile
    ) -> "UserOAuthAccount":
        access_token = token["access_token"]
        refresh_token = token.get("refresh_token", None)
        expires_at = token.get("expires_at", None)

        async with db_session() as session:
            statement = (
                select(UserOAuthAccount)
                .where(UserOAuthAccount.oauth_name == oauth_name)
                .where(UserOAuthAccount.account_id == profile.account_id)
            )
            results = await session.exec(statement)
            oauth_account: Optional[UserOAuthAccount] = results.one_or_none()
            if oauth_account:
                oauth_account.access_token = access_token
                oauth_account.refresh_token = refresh_token
                oauth_account.expires_at = expires_at
                oauth_account.account_name = profile.account_name
            else:
                oauth_account = UserOAuthAccount(
                    oauth_name=oauth_name,
                    access_token=access_token,
                    refresh_token=refresh_token,
                    expires_at=expires_at,
                    account_id=profile.account_id,
                    account_name=profile.account_name,
                    account_email=profile.account_email,
                )
            session.sync_session.add(oauth_account)
            await session.commit()
            await session.refresh(oauth_account)

        return oauth_account

joj.horse.services.db.db_session sqlalchemy.schema.ForeignKey sqlmodel.Field sqlmodel.select sqlmodel.Relationship
[((558, 565), 'sqlmodel.Field', 'Field', ([], {}), '()\n', (563, 565), False, 'from sqlmodel import Field, Relationship, select\n'), ((590, 597), 'sqlmodel.Field', 'Field', ([], {}), '()\n', (595, 597), False, 'from sqlmodel import Field, Relationship, select\n'), ((633, 659), 'sqlmodel.Field', 'Field', (['None'], {'nullable': '(True)'}), '(None, nullable=True)\n', (638, 659), False, 'from sqlmodel import Field, Relationship, select\n'), ((692, 718), 'sqlmodel.Field', 'Field', (['None'], {'nullable': '(True)'}), '(None, nullable=True)\n', (697, 718), False, 'from sqlmodel import Field, Relationship, select\n'), ((741, 758), 'sqlmodel.Field', 'Field', ([], {'index': '(True)'}), '(index=True)\n', (746, 758), False, 'from sqlmodel import Field, Relationship, select\n'), ((793, 831), 'sqlmodel.Field', 'Field', (['None'], {'index': '(True)', 'nullable': '(True)'}), '(None, index=True, nullable=True)\n', (798, 831), False, 'from sqlmodel import Field, Relationship, select\n'), ((857, 874), 'sqlmodel.Field', 'Field', ([], {'index': '(True)'}), '(index=True)\n', (862, 874), False, 'from sqlmodel import Field, Relationship, select\n'), ((1023, 1068), 'sqlmodel.Relationship', 'Relationship', ([], {'back_populates': '"""oauth_accounts"""'}), "(back_populates='oauth_accounts')\n", (1035, 1068), False, 'from sqlmodel import Field, Relationship, select\n'), ((1393, 1405), 'joj.horse.services.db.db_session', 'db_session', ([], {}), '()\n', (1403, 1405), False, 'from joj.horse.services.db import db_session\n'), ((944, 986), 'sqlalchemy.schema.ForeignKey', 'ForeignKey', (['"""users.id"""'], {'ondelete': '"""CASCADE"""'}), "('users.id', ondelete='CASCADE')\n", (954, 986), False, 'from sqlalchemy.schema import Column, ForeignKey\n'), ((1460, 1484), 'sqlmodel.select', 'select', (['UserOAuthAccount'], {}), '(UserOAuthAccount)\n', (1466, 1484), False, 'from sqlmodel import Field, Relationship, select\n')]