import logging
import urllib.request
import json
from typing import Literal, Optional, Type, Union, Annotated
from pydantic import (
    BaseModel,
    Field,
    ConfigDict,
    field_validator,
    AliasChoices,
    Discriminator,
    FilePath,
    PrivateAttr,
)
import datetime
from typing import Dict, Any
from joserfc.jwk import KeySet, RSAKey, OctKey, ECKey
from joserfc import jwt
from joserfc.jwk import GuestProtocol
from authentikate.errors import JwksError, MalformedJwtTokenError


logger = logging.getLogger(__name__)


class JWTToken(BaseModel):
    """A JWT token

    This is a pydantic model that represents a JWT token.
    It is used to validate the token and to extract information from it.
    The token is decoded using the `decode_token` function.

    """

    model_config = ConfigDict(extra="ignore")

    sub: str
    """A unique identifier for the user (is unique for the issuer)"""
    iss: str
    """The issuer of the token"""

    exp: datetime.datetime
    """The expiration time of the token"""

    active_org: str | None = None
    """The active organization of the user, if any"""

    client_id: str
    """The client_id of the app that requested the token"""
    preferred_username: str
    """The username of the user"""
    roles: list[str]
    """The roles of the user"""
    scope: str
    """The scope of the token"""

    iat: datetime.datetime
    """The issued at time of the token"""

    aud: list[str] | None = None
    """The audience of the token"""

    jti: str | None = None
    """The unique identifier for the token"""

    raw: str
    """ The raw original token string """

    client_app: str | None = None
    """ The client app name """

    client_release: str | None = None
    """ The client release version """

    client_device: str | None = None
    """ The client device identifier """

    @field_validator("aud", mode="before")
    def aud_to_list(
        cls: Type["JWTToken"], v: str | list[str] | None
    ) -> list[str] | None:
        """Convert the aud to a list"""
        if not v:
            return None
        if isinstance(v, str):
            return [v]

        return v

    @field_validator("sub", mode="before")
    def sub_to_username(cls: Type["JWTToken"], v: str) -> str:
        """Convert the sub to a username compatible string"""
        if isinstance(v, int):
            return str(v)
        return v

    @field_validator("iat", mode="before")
    def iat_to_datetime(cls: Type["JWTToken"], v: int) -> datetime.datetime:
        """Convert the iat to a datetime object"""
        if v is None:
            return None
        if isinstance(v, int):
            return datetime.datetime.fromtimestamp(v)
        return v

    @field_validator("exp", mode="before")
    def exp_to_datetime(cls: Type["JWTToken"], v: int) -> datetime.datetime:
        """Convert the exp to a datetime object"""
        if isinstance(v, int):
            return datetime.datetime.fromtimestamp(v)
        return v

    @property
    def changed_hash(self) -> str:
        """A hash that changes when the user changes"""
        return str(
            hash(
                self.sub
                + self.preferred_username
                + " ".join(self.roles)
                + (self.active_org or "")
            )
        )

    @property
    def scopes(self) -> list[str]:
        """The scopes of the token. Each scope is a string separated by a space"""
        return self.scope.split(" ")

    def has_scopes(self, scopes: list[str]) -> bool:
        """Check if the user has the given scope"""
        if not scopes:
            return True

        return all(scope in self.scopes for scope in scopes)

    def has_any_role(self, roles: list[str]) -> bool:
        """Check if the user has any of the given roles"""
        if not roles:
            return True

        return any(role in self.roles for role in roles)

    def has_roles(self, roles: list[str]) -> bool:
        """Check if the user has the given role"""
        if not roles:
            return True

        return all(role in self.roles for role in roles)

    def has_any_scope(self, scopes: list[str]) -> bool:
        """Check if the user has any of the given scopes"""
        if not scopes:
            return True

        return any(scope in self.scopes for scope in scopes)


class StaticToken(JWTToken):
    """A static JWT token"""

    sub: str
    iss: str = "static_issuer"
    """The issuer of the token"""
    iat: datetime.datetime = Field(default_factory=lambda: datetime.datetime.now())
    exp: datetime.datetime = Field(
        default_factory=lambda: datetime.datetime.now() + datetime.timedelta(days=1)
    )
    client_id: str = "static"
    client_app: str = "static_app"
    client_release: str = "v1.0.0"
    client_device: str = "static_device"
    active_org: str = "static_org"
    preferred_username: str = "static_user"
    scope: str = "openid profile email"
    roles: list[str] = Field(default_factory=lambda: ["admin"])
    raw: str = Field(default_factory=lambda: "static_token")


class ImitationRequest(BaseModel):
    """An imitation request"""

    sub: str
    iss: str


class Issuer(BaseModel):
    """An issuer

    This is a pydantic model that represents an issuer.
    It is used to validate the issuer and to extract information from it.
    """

    model_config = ConfigDict(extra="forbid")
    kind: str
    iss: str = Field(
        validation_alias=AliasChoices("iss", "issuer", "issuer_url", "ISSUER")
    )
    """The issuer of the token"""

    def get_as_jwks(self) -> list[Dict[str, Any]]:
        """Get the jwks of the issuer"""
        raise NotImplementedError(
            "get_jwks not implemented. Must be implemented in subclass"
        )

    def refresh(self) -> None:
        """Refresh the issuer jwks if applicable"""
        pass


class JWKIssuer(Issuer):
    """An issuer

    This is a pydantic model that represents an issuer.
    It is used to validate the issuer and to extract information from it.
    """

    kind: Literal["jwks_dict"] = Field(
        default="jwks_dict",
    )

    iss: str = Field(
        validation_alias=AliasChoices("iss", "issuer", "issuer_url", "ISSUER")
    )
    """The issuer of the token"""

    jwks: Dict[str, Any] = Field(
        validation_alias=AliasChoices("jwks", "JWKS", "JWKS_DICT")
    )
    """The jwks of the issuer"""

    @field_validator("jwks", mode="before")
    def validate_jwks_dict(cls: Type["JWKIssuer"], v: Dict[str, Any]) -> Dict[str, Any]:
        """Validate the jwks dict"""
        if not isinstance(v, dict):
            raise ValueError("jwks_dict must be a dict")
        if "keys" not in v:
            raise ValueError("jwks_dict must contain a keys field")
        if not isinstance(v["keys"], list):
            raise ValueError("jwks_dict keys must be a list")
        return v

    def get_as_jwks(self) -> list[Dict[str, Any]]:
        """Get the jwks of the issuer"""
        return self.jwks["keys"]


class RSAKeyIssuer(Issuer):
    """An issuer

    This is a pydantic model that represents an issuer.
    It is used to validate the issuer and to extract information from it.
    """

    model_config = ConfigDict(extra="forbid")
    kind: Literal["rsa"] = Field(
        default="rsa",
    )

    iss: str = Field(
        validation_alias=AliasChoices("iss", "issuer", "issuer_url", "ISSUER")
    )
    key_id: str = Field(
        default="1", validation_alias=AliasChoices("key_id", "kid", "KID")
    )
    """The issuer of the token"""
    public_key: str = Field(validation_alias=AliasChoices("public_key", "PUBLIC_KEY"))

    def get_as_jwks(self) -> list[Dict[str, Any]]:
        """Get the jwks of the issuer"""
        t = RSAKey.import_key(self.public_key)
        return [t.as_dict(kid=self.key_id)]


class RSAKeyFileIssuer(Issuer):
    """An issuer

    This is a pydantic model that represents an issuer.
    It is used to validate the issuer and to extract information from it.
    """

    model_config = ConfigDict(extra="forbid")
    kind: Literal["rsa_file"] = Field(
        default="rsa_file",
    )

    iss: str = Field(
        validation_alias=AliasChoices("iss", "issuer", "issuer_url", "ISSUER")
    )
    key_id: str = Field(
        default="1", validation_alias=AliasChoices("key_id", "kid", "KID")
    )
    """The issuer of the token"""
    public_key_pem_file: FilePath = Field(
        validation_alias=AliasChoices("public_key_pem_file", "PUBLIC_KEY_PEM_FILE")
    )

    def get_as_jwks(self) -> list[Dict[str, Any]]:
        """Get the jwks of the issuer"""

        with open(self.public_key_pem_file, "rb") as f:
            public_key = f.read()

        t = RSAKey.import_key(public_key)
        return [t.as_dict(kid=self.key_id)]


class JWKSUriIssuer(Issuer):
    """An issuer

    This is a pydantic model that represents an issuer.
    It is used to validate the issuer and to extract information from it.
    """

    model_config = ConfigDict(extra="forbid")
    kind: Literal["jwks_uri"] = Field(
        default="jwks_uri",
    )

    iss: str = Field(
        validation_alias=AliasChoices("iss", "issuer", "issuer_url", "ISSUER")
    )
    """The issuer of the token"""
    jwks_uri: str = Field(validation_alias=AliasChoices("jwks_uri", "JWKS_URI"))
    _cache: list[Dict[str, Any]] | None = PrivateAttr(default=None)

    def get_as_jwks(self) -> list[Dict[str, Any]]:
        """Get the jwks of the issuer"""

        if self._cache is None:
            self.refresh()

        return self._cache  # type: ignore

    def refresh(self) -> None:
        """Refresh the jwks from the uri"""
        try:
            with urllib.request.urlopen(self.jwks_uri) as response:
                data = json.loads(response.read())
                self._cache = data["keys"]
        except Exception as e:
            raise JwksError(f"Error fetching jwks from {self.jwks_uri}") from e


IssuerUnion = Annotated[
    Union[JWKIssuer, RSAKeyIssuer, RSAKeyFileIssuer, JWKSUriIssuer], Discriminator("kind")
]



class AuthentikateSettings(BaseModel):
    """The settings for authentikate

    This is a pydantic model that represents the settings for authentikate.
    It is used to configure the library.
    """

    model_config = ConfigDict(extra="forbid")

    issuers: list[IssuerUnion] = Field(
        validation_alias=AliasChoices(
            "issuers",
            "iss",
            "issuer",
            "issuer_url",
            "ISSUERS",
        )
    )
    authorization_headers: list[str] = Field(
        default_factory=lambda: [
            "Authorization",
            "X-Authorization",
            "AUTHORIZATION",
            "authorization",
        ],
        validation_alias=AliasChoices(
            "authorization_headers", "AUTHORIZATION_HEADERS", "AUTHORIZATION_HEADERS"
        ),
    )
    static_tokens: dict[str, StaticToken] = Field(
        default_factory=dict,
        validation_alias=AliasChoices(
            "static_tokens", "STATIC_TOKENS", "STATIC_TOKENS"
        ),
    )
    """A map of static tokens to their decoded values. Should only be used in tests."""

    def get_jwks(self) -> list[Dict[str, Any]]:
        """Get the jwks of the issuer"""

        merged_jwks = {}

        for issuer in self.issuers:
            keys = issuer.get_as_jwks()

            if not isinstance(keys, list):
                raise JwksError("keys must be a list")

            for key in keys:
                if key.get("kid") is None:
                    raise JwksError("key must contain a kid field")

                if key["kid"] in merged_jwks:
                    raise JwksError(f"Duplicate kid found: {key['kid']}")

                merged_jwks[key["kid"]] = key

        if not merged_jwks:
            raise JwksError("No keys found in jwks")

        validated_keys = []

        for key in merged_jwks.values():
            validated_keys.append(key)

        return validated_keys

    def load_key(self, obj: GuestProtocol) -> KeySet:
        """Resolve the key from the header"""
        kid = obj.headers().get("kid")
        if not kid:
            raise MalformedJwtTokenError("Missing kid in header")

        jwks = self.get_jwks()

        # Check if the kid is in the jwks
        found = False
        for key in jwks:
            if key.get("kid") == kid:
                found = True
                break

        if not found:
            # Trigger refresh on all issuers
            for issuer in self.issuers:
                issuer.refresh()
            # Re-fetch
            jwks = self.get_jwks()

        key_set = KeySet.import_key_set({"keys": jwks})
        return key_set
