from __future__ import annotations

import re
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from enum import Enum, auto
from hashlib import pbkdf2_hmac
from typing import Any, Literal, Optional, Protocol

from pydantic import SecretStr
from starlette.responses import Response
from typing_extensions import TypeVar

from phoenix.config import (
    get_env_cookies_path,
    get_env_phoenix_use_secure_cookies,
)

ResponseType = TypeVar("ResponseType", bound=Response)


def compute_password_hash(*, password: SecretStr, salt: bytes) -> bytes:
    """
    Salts and hashes a password using PBKDF2, HMAC and SHA256.

    Args:
        password: the password to hash, wrapped in pydantic ``SecretStr``
        salt (bytes): the salt to use, must not be zero-length
    Returns:
        bytes: the hashed password
    """
    assert salt
    password_bytes = password.get_secret_value().encode("utf-8")
    return pbkdf2_hmac("sha256", password_bytes, salt, NUM_ITERATIONS)


def is_valid_password(*, password: SecretStr, salt: bytes, password_hash: bytes) -> bool:
    """
    Determines whether the password is valid by salting and hashing the password
    and comparing against the existing hash value.

    Args:
        password (str): the password to validate
        salt (bytes): the salt to use, must not be zero-length
        password_hash (bytes): the hash to compare against
    Returns:
        bool: True if the password is valid, False otherwise
    """
    assert salt
    return password_hash == compute_password_hash(password=password, salt=salt)


def sanitize_email(email: str) -> str:
    """
    Sanitizes an email address by trimming whitespace and converting to lowercase.

    Args:
        email (str): the email address to sanitize
    Returns:
        str: the sanitized email address
    """
    return email.strip().lower()


def validate_email_format(email: str) -> None:
    """
    Checks that the email has a valid format.

    Args:
        email (str): the email address to validate
    Returns:
        None
    Raises:
        ValueError: if the email address is invalid
    """
    if EMAIL_PATTERN.match(email) is None:
        raise ValueError("Invalid email address")


def validate_password_format(password: str) -> None:
    """
    Checks that the password has a valid format.
    """
    get_password_requirements().validate(password)


def get_password_requirements() -> _PasswordRequirements:
    """
    Returns the active password requirements based on configuration.

    When PHOENIX_ENABLE_STRONG_PASSWORD_POLICY is enabled, returns strong password
    requirements. Otherwise, returns the default minimal requirements.
    """
    from phoenix.config import get_env_enable_strong_password_policy

    if get_env_enable_strong_password_policy():
        return STRONG_PASSWORD_REQUIREMENTS
    return PASSWORD_REQUIREMENTS


def set_access_token_cookie(
    *, response: ResponseType, access_token: str, max_age: timedelta
) -> ResponseType:
    return _set_cookie(
        response=response,
        cookie_name=PHOENIX_ACCESS_TOKEN_COOKIE_NAME,
        cookie_max_age=max_age,
        samesite="strict",
        value=access_token,
    )


def set_refresh_token_cookie(
    *, response: ResponseType, refresh_token: str, max_age: timedelta
) -> ResponseType:
    return _set_cookie(
        response=response,
        cookie_name=PHOENIX_REFRESH_TOKEN_COOKIE_NAME,
        cookie_max_age=max_age,
        samesite="strict",
        value=refresh_token,
    )


def set_oauth2_state_cookie(
    *, response: ResponseType, state: str, max_age: timedelta
) -> ResponseType:
    return _set_cookie(
        response=response,
        cookie_name=PHOENIX_OAUTH2_STATE_COOKIE_NAME,
        cookie_max_age=max_age,
        samesite="lax",
        value=state,
    )


def set_oauth2_nonce_cookie(
    *, response: ResponseType, nonce: str, max_age: timedelta
) -> ResponseType:
    return _set_cookie(
        response=response,
        cookie_name=PHOENIX_OAUTH2_NONCE_COOKIE_NAME,
        cookie_max_age=max_age,
        samesite="lax",
        value=nonce,
    )


def set_oauth2_code_verifier_cookie(
    *, response: ResponseType, code_verifier: str, max_age: timedelta
) -> ResponseType:
    return _set_cookie(
        response=response,
        cookie_name=PHOENIX_OAUTH2_CODE_VERIFIER_COOKIE_NAME,
        cookie_max_age=max_age,
        samesite="lax",
        value=code_verifier,
    )


def _set_cookie(
    *,
    response: ResponseType,
    cookie_name: str,
    cookie_max_age: timedelta,
    samesite: Literal["strict", "lax"],
    value: str,
) -> ResponseType:
    response.set_cookie(
        key=cookie_name,
        value=value,
        secure=get_env_phoenix_use_secure_cookies(),
        httponly=True,
        samesite=samesite,
        max_age=int(cookie_max_age.total_seconds()),
        path=get_env_cookies_path(),
    )
    return response


def delete_access_token_cookie(response: ResponseType) -> ResponseType:
    response.delete_cookie(key=PHOENIX_ACCESS_TOKEN_COOKIE_NAME)
    return response


def delete_refresh_token_cookie(response: ResponseType) -> ResponseType:
    response.delete_cookie(key=PHOENIX_REFRESH_TOKEN_COOKIE_NAME)
    return response


def delete_oauth2_state_cookie(response: ResponseType) -> ResponseType:
    response.delete_cookie(key=PHOENIX_OAUTH2_STATE_COOKIE_NAME)
    return response


def delete_oauth2_nonce_cookie(response: ResponseType) -> ResponseType:
    response.delete_cookie(key=PHOENIX_OAUTH2_NONCE_COOKIE_NAME)
    return response


def delete_oauth2_code_verifier_cookie(response: ResponseType) -> ResponseType:
    response.delete_cookie(key=PHOENIX_OAUTH2_CODE_VERIFIER_COOKIE_NAME)
    return response


@dataclass(frozen=True)
class _PasswordRequirements:
    """
    Password must be at least `length` characters long. Password must not contain whitespace
    characters. Password can contain only ASCII characters. The arguments `special_chars`,
    `digits`, `upper_case`, and `lower_case` control what category of characters will appear
    in the password. If set to True, at least one character from the corresponding category
    is guaranteed to appear. Special characters are characters from `!@#$%^&*()_+`, digits
    are characters from `0123456789`, and uppercase and lowercase characters are characters
    from the ASCII set of letters.

    Attributes:
        length (int): the minimum length of the password
        digits (bool): whether the password must contain at least one digit
        lower_case (bool): whether the password must contain at least one lowercase letter
        upper_case (bool): whether the password must contain at least one uppercase letter
        special_chars (bool): whether the password must contain at least one special character
    """

    length: int
    digits: bool = False
    lower_case: bool = False
    upper_case: bool = False
    special_chars: bool = False

    def validate(
        self,
        string: str,
        /,
        err_msg_subject: Literal["Password", "Phoenix secret"] = "Password",
    ) -> None:
        """
        Validates the password against the requirements.

        Args:
            string (str): the password to validate
            err_msg_subject (str, optional): the subject of the error message,
                defaults to "Password"
        Returns:
            None
        Raises:
            ValueError: if the password does not meet the requirements
        """
        if not string:
            raise ValueError(f"{err_msg_subject} must be non-empty")
        if any(char.isspace() for char in string):
            raise ValueError(f"{err_msg_subject} must not contain whitespace characters")
        if not string.isascii():
            raise ValueError(f"{err_msg_subject} must contain only ASCII characters")
        err_msg = []
        if len(string) < self.length:
            err_msg.append(f"must be at least {self.length} characters long")
        if self.digits and not any(char.isdigit() for char in string):
            err_msg.append("at least one digit")
        if self.lower_case and not any(char.islower() for char in string):
            err_msg.append("at least one lowercase letter")
        if self.upper_case and not any(char.isupper() for char in string):
            err_msg.append("at least one uppercase letter")
        if self.special_chars and not any(char in "!@#$%^&*()_+" for char in string):
            err_msg.append("at least one special character")
        if not err_msg:
            return
        if len(err_msg) > 1:
            err_text = f"{err_msg_subject} " + ", ".join(err_msg[:-1]) + ", and " + err_msg[-1]
        else:
            err_text = f"{err_msg_subject} {err_msg[0]}"
        raise ValueError(err_text)


DEFAULT_ADMIN_USERNAME = "admin"
DEFAULT_ADMIN_EMAIL = "admin@localhost"
DEFAULT_ADMIN_PASSWORD = "admin"
DEFAULT_SYSTEM_USERNAME = "system"
DEFAULT_SYSTEM_EMAIL = "system@localhost"
DEFAULT_SECRET_LENGTH = 32
DEFAULT_PASSWORD_RESET_TOKEN_EXPIRY_MINUTES = 15
DEFAULT_ACCESS_TOKEN_EXPIRY_MINUTES = 10
DEFAULT_REFRESH_TOKEN_EXPIRY_MINUTES = 60 * 24 * 7
"""The default length of a secret key in bytes."""
EMAIL_PATTERN = re.compile(r"^[^@\s]+@[^@\s]+[.][^@\s]+\Z")
"""The regular expression pattern for a valid email address."""
NUM_ITERATIONS = 10_000
"""The number of iterations to use for the PBKDF2 key derivation function."""
MIN_PASSWORD_LENGTH = 4
"""The minimum length of a password."""
STRONG_PASSWORD_LENGTH = 12
"""The minimum length of a password when the strong password policy is enabled."""
PASSWORD_REQUIREMENTS = _PasswordRequirements(length=MIN_PASSWORD_LENGTH)
"""The requirements for a valid password."""
STRONG_PASSWORD_REQUIREMENTS = _PasswordRequirements(
    length=STRONG_PASSWORD_LENGTH,
    digits=True,
    lower_case=True,
    upper_case=True,
    special_chars=True,
)
"""The requirements for a valid password when the strong password policy is enabled."""
REQUIREMENTS_FOR_PHOENIX_SECRET = _PasswordRequirements(
    length=DEFAULT_SECRET_LENGTH, digits=True, lower_case=True
)
"""The requirements for the Phoenix secret key."""
JWT_ALGORITHM = "HS256"
"""The algorithm to use for the JSON Web Token."""
PHOENIX_ACCESS_TOKEN_COOKIE_NAME = "phoenix-access-token"
"""The name of the cookie that stores the Phoenix access token."""
PHOENIX_REFRESH_TOKEN_COOKIE_NAME = "phoenix-refresh-token"
"""The name of the cookie that stores the Phoenix refresh token."""
PHOENIX_OAUTH2_STATE_COOKIE_NAME = "phoenix-oauth2-state"
"""The name of the cookie that stores the state used for the OAuth2 authorization code flow."""
PHOENIX_OAUTH2_NONCE_COOKIE_NAME = "phoenix-oauth2-nonce"
"""The name of the cookie that stores the nonce used for the OAuth2 authorization code flow."""
PHOENIX_OAUTH2_CODE_VERIFIER_COOKIE_NAME = "phoenix-oauth2-code-verifier"
"""The name of the cookie that stores the PKCE code verifier for OAuth2."""
DEFAULT_OAUTH2_LOGIN_EXPIRY_MINUTES = 15
"""
The default amount of time in minutes that can elapse between the initial
redirect to the IDP and the invocation of the callback URL during the OAuth2
authorization code flow.
"""


class Token(str): ...


class ClaimSetStatus(Enum):
    VALID = auto()
    INVALID = auto()
    EXPIRED = auto()


@dataclass(frozen=True)
class TokenAttributes: ...


@dataclass(frozen=True)
class ClaimSet:
    issuer: Optional[Any] = None
    "Analog of `iss` claim in JWT RFC7519: https://datatracker.ietf.org/doc/html/rfc7519#section-4.1.1"
    subject: Optional[Any] = None
    "Analog of `sub` claim in JWT RFC7519: https://datatracker.ietf.org/doc/html/rfc7519#section-4.1.2"
    audience: Optional[Any] = None
    "Analog of `aud` claim in JWT RFC7519: https://datatracker.ietf.org/doc/html/rfc7519#section-4.1.3"
    not_before: Optional[datetime] = None
    "Analog of `nbf` claim in JWT RFC7519: https://datatracker.ietf.org/doc/html/rfc7519#section-4.1.5"
    issued_at: Optional[datetime] = None
    "Analog of `iat` claim in JWT RFC7519: https://datatracker.ietf.org/doc/html/rfc7519#section-4.1.6"
    expiration_time: Optional[datetime] = None
    "Analog of `exp` claim in JWT RFC7519: https://datatracker.ietf.org/doc/html/rfc7519#section-4.1.4"
    token_id: Optional[Any] = None
    "Analog of `jti` claim in JWT RFC7519: https://datatracker.ietf.org/doc/html/rfc7519#section-4.1.7"
    attributes: Optional[TokenAttributes] = None
    "Application/domain-specific claims"

    @property
    def status(self) -> ClaimSetStatus:
        # Per JWT RFC 7519 Section 4.1.4, the expiration time identifies the time
        # "on or after which" the JWT must not be accepted. Use <= for inclusive check.
        # https://datatracker.ietf.org/doc/html/rfc7519#section-4.1.4
        if (
            self.expiration_time
            and self.expiration_time.timestamp() <= datetime.now(timezone.utc).timestamp()
        ):
            return ClaimSetStatus.EXPIRED
        if self.token_id is not None and self.subject is not None:
            return ClaimSetStatus.VALID
        return ClaimSetStatus.INVALID


class CanReadToken(Protocol):
    async def read(self, token: Token) -> Optional[ClaimSet]: ...
