from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, Literal

import jwt

from pyidentity.security.authentication.abc import IAuthenticationHandler
from pyidentity.security.authentication.jwt.token import JWTSecurityToken
from pyidentity.security.authentication.result import AuthenticationResult
from pyidentity.security.authentication.utils import (
    get_authentication_cookie_name,
    get_authorization_scheme_param,
    join_value_from_chunks,
    split_value_to_chunks,
)
from pyidentity.security.claims import ClaimsIdentity, ClaimsPrincipal
from pyidentity.utils import normalize_datetime

if TYPE_CHECKING:
    from collections.abc import Iterable, Sequence

    from pyidentity.http.context import HttpContext


@dataclass
class TokenValidationParameters:
    """Contains a set of parameters that are used by a JWTAuthenticationHandler when validating a security token."""

    issuer_signing_key: Any
    valid_algorithms: Sequence[str] = field(default_factory=lambda: ["HS256"])
    valid_audiences: str | Iterable[str] | None = field(default=None)
    valid_issuers: str | Sequence[str] | None = field(default=None)
    verify_exp: bool = field(default=True)
    verify_nbf: bool = field(default=True)
    verify_iat: bool = field(default=True)
    verify_aud: bool = field(default=True)
    verify_iss: bool = field(default=True)
    verify_sub: bool = field(default=True)
    verify_jti: bool = field(default=True)
    require_claims: list[str] = field(default_factory=list)
    leeway: float | timedelta = field(default=0)

    @property
    def options(self) -> dict[str, Any]:
        return {
            "verify_exp": self.verify_exp,
            "verify_nbf": self.verify_nbf,
            "verify_iat": self.verify_iat,
            "verify_aud": self.verify_aud,
            "verify_iss": self.verify_iss,
            "verify_sub": self.verify_sub,
            "verify_jti": self.verify_jti,
            "require": self.require_claims,
        }


@dataclass
class CookieAuthenticationOptions:
    name: str | None = field(default=None)
    """A string that will be the cookies key."""
    max_age: timedelta = field(default=lambda: timedelta(days=7))
    """An integer that defines the lifetime of the cookies in seconds."""
    expires: timedelta | None = field(default=None)
    """Timedelta, which defines the interval until the cookies expires."""
    path: str = field(default="/")
    """A string that specifies the subset of routes to which the cookies will apply."""
    domain: str | None = field(default=None)
    """A string that specifies the domain for which the cookies is valid."""
    httponly: bool = field(default=True)
    """A bool indicating that the cookies cannot be accessed via JavaScript through *Document.cookies* property,
    the *XMLHttpRequest* or Request APIs."""
    secure: bool = field(default=True)
    """A bool indicating that the cookies will only be sent to the server
    if request is made using SSL and the HTTPS protocol."""
    samesite: Literal["lax", "strict", "none"] = field(default="lax")
    """A string that specifies the samesite strategy for the cookies.
    Valid values are 'lax', 'strict' and 'none'. Defaults to 'lax'."""


class JWTBaseAuthenticationHandler(IAuthenticationHandler):
    def __init__(self, validation_parameters: TokenValidationParameters):
        self._validation_parameters = validation_parameters

    def get_token(self, context: HttpContext, scheme: str) -> str | None:
        pass

    async def authenticate(self, context: HttpContext, scheme: str) -> AuthenticationResult:
        token = self.get_token(context, scheme)
        if token is None:
            return AuthenticationResult(ClaimsPrincipal(), {})

        try:
            jwt_token = JWTSecurityToken.decode(
                token,
                issuer_signing_key=self._validation_parameters.issuer_signing_key,
                algorithms=self._validation_parameters.valid_algorithms,
                options=self._validation_parameters.options,
                audience=self._validation_parameters.valid_audiences,
                issuer=self._validation_parameters.valid_issuers,
                leeway=self._validation_parameters.leeway,
            )
        except jwt.PyJWTError:
            return AuthenticationResult(ClaimsPrincipal(), {})
        else:
            identity = ClaimsIdentity("AuthenticationTypes.Federation", jwt_token.claims)
            return AuthenticationResult(ClaimsPrincipal([identity]), {})

    async def sign_in(
        self,
        context: HttpContext,
        scheme: str,
        principal: ClaimsPrincipal,
        **properties: Any,
    ) -> None:
        pass

    async def sign_out(self, context: HttpContext, scheme: str) -> None:
        pass


class JWTBearerAuthenticationHandler(JWTBaseAuthenticationHandler):
    def __init__(self, validation_parameters: TokenValidationParameters) -> None:
        super().__init__(validation_parameters)

    def get_token(self, context: HttpContext, scheme: str) -> str | None:
        authorization = context.request.headers.get("Authorization")
        scheme, token = get_authorization_scheme_param(authorization)
        if not authorization or scheme.lower() != "bearer":
            return None
        return token


class JWTCookieAuthenticationHandler(JWTBaseAuthenticationHandler):
    def __init__(
        self,
        secret_key: Any,
        validation_parameters: TokenValidationParameters,
        algorithm: str = "HS256",
        audience: str | None = None,
        issuer: str | None = None,
        cookie_options: CookieAuthenticationOptions | None = None,
    ) -> None:
        self._secret_key = secret_key
        self._algorithm = algorithm
        self._audience = audience
        self._issuer = issuer
        self._cookie_options = cookie_options or CookieAuthenticationOptions()
        super().__init__(validation_parameters)

    def get_token(self, context: HttpContext, scheme: str) -> str | None:
        token_key = get_authentication_cookie_name(self._cookie_options.name, scheme)
        return join_value_from_chunks(token_key, context.request.cookies)

    async def sign_in(
        self,
        context: HttpContext,
        scheme: str,
        principal: ClaimsPrincipal,
        **properties: Any,
    ) -> None:
        context.response.headers["Cache-Control"] = "no-cache,no-store"
        context.response.headers["Pragma"] = "no-cache"
        is_persistent = properties.get("is_persistent")
        token_key = get_authentication_cookie_name(self._cookie_options.name, scheme)

        max_age = self._cookie_options.expires  # timedelta or None
        if is_persistent:
            max_age = normalize_datetime(datetime.now(timezone.utc) + self._cookie_options.expires_default)
        elif max_age is not None:
            max_age = normalize_datetime(datetime.now(timezone.utc) + self._cookie_options.expires)

        token = JWTSecurityToken(
            signing_key=self._secret_key,
            expires=max_age,
            algorithm=self._algorithm,
            audience=self._audience,
            issuer=self._issuer,
            claims=principal.claims,
        ).encode()

        for key, value in split_value_to_chunks(token_key, token).items():
            context.response.set_cookie(
                key=key,
                value=value,
                max_age=max_age,
                path=self._cookie_options.path,
                domain=self._cookie_options.domain,
                secure=self._cookie_options.secure,
                httponly=self._cookie_options.httponly,
                samesite=self._cookie_options.samesite,
            )

    async def sign_out(self, context: HttpContext, scheme: str) -> None:
        token_key = get_authentication_cookie_name(self._cookie_options.name, scheme)
        for key in context.request.cookies:
            if key.startswith(token_key):
                context.response.delete_cookie(key)
