from __future__ import annotations

from typing import Annotated
from typing import Literal

from pydantic import BaseModel
from pydantic import ConfigDict
from pydantic import Discriminator
from pydantic import Field

from amsdal_ml.mcp_server.mcp_oauth.services.types import GrantTypes
from amsdal_ml.mcp_server.mcp_oauth.services.types import OAuthConstants
from amsdal_ml.mcp_server.mcp_oauth.services.types import ResponseTypes


class RegisterRequest(BaseModel):
    redirect_uris: list[str] = Field(default_factory=list)
    grant_types: list[Literal[GrantTypes.REFRESH_TOKEN, GrantTypes.AUTHORIZATION_CODE]] = Field(default_factory=list)
    response_types: list[Literal[ResponseTypes.CODE]] = Field(default_factory=list)


class RegisterResponse(BaseModel):
    client_id: str
    expires_at: int
    client_id_issued_at: int


class AuthorizationCodeFormData(BaseModel):
    model_config = ConfigDict(extra='forbid')

    grant_type: Literal[GrantTypes.AUTHORIZATION_CODE]
    code: str
    code_verifier: str
    redirect_uri: str
    client_id: str
    resource: str | None = None


class RefreshTokenFormData(BaseModel):
    model_config = ConfigDict(extra='forbid')

    grant_type: Literal[GrantTypes.REFRESH_TOKEN]
    refresh_token: str
    client_id: str
    resource: str | None = None


class AuthorizeResponse(BaseModel):
    code: str
    state: str


class TokenResponse(BaseModel):
    access_token: str
    token_type: str = OAuthConstants.TOKEN_TYPE_BEARER
    expires_in: int
    refresh_token: str
    scope: str | None = None


class JwtDecodedPayload(BaseModel):
    model_config = ConfigDict(extra='ignore')

    iss: str | None = None
    sub: str | None = None
    email: str | None = None
    aud: str | list[str] | None = None
    token_type: str | None = None
    resource: str | None = None


class InternalAuthTokenPayload(BaseModel):
    email: str
    iat: int
    exp: int


TokenFromData = Annotated[
    AuthorizationCodeFormData | RefreshTokenFormData,
    Discriminator('grant_type'),
]
