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 = (
                