"""JWT utilities using Ed25519 asymmetric encryption.

This module provides functions for encoding and decoding JSON Web Tokens (JWT)
using Ed25519 asymmetric encryption algorithm.
"""

from datetime import datetime, timedelta
from pathlib import Path
from typing import Any

import jwt as pyjwt
from cryptography.hazmat.primitives.asymmetric.ed25519 import (
    Ed25519PrivateKey,
    Ed25519PublicKey,
)

from hezor_common.security.signature import (
    deserialize_private_key,
    deserialize_public_key,
)


def encode_jwt(
    private_key: Ed25519PrivateKey,
    payload: dict[str, Any],
    algorithm: str = "EdDSA",
    expires_in: int | None = None,
) -> str:
    """Encode a JWT token using Ed25519 private key.

    Parameters
    ----------
    private_key : Ed25519PrivateKey
        The private key to use for signing
    payload : dict[str, Any]
        The payload data to encode (must be JSON-serializable)
    algorithm : str, optional
        The algorithm to use (default: "EdDSA")
    expires_in : int | None, optional
        Token expiration time in seconds (default: None, no expiration)

    Returns
    -------
    str
        The encoded JWT token

    Examples
    --------
    >>> from hezor_common.security import generate_key_pair
    >>> private_key, public_key = generate_key_pair()
    >>> payload = {"user_id": "123", "role": "admin"}
    >>> token = encode_jwt(private_key, payload)
    >>> isinstance(token, str)
    True
    >>> len(token.split("."))
    3
    """
    # Add expiration if specified
    if expires_in is not None:
        payload = payload.copy()
        payload["exp"] = datetime.utcnow() + timedelta(seconds=expires_in)

    return pyjwt.encode(payload, private_key, algorithm=algorithm)


def decode_jwt(
    public_key: Ed25519PublicKey,
    token: str,
    algorithms: list[str] | None = None,
    verify: bool = True,
) -> dict[str, Any]:
    """Decode and verify a JWT token using Ed25519 public key.

    Parameters
    ----------
    public_key : Ed25519PublicKey
        The public key to use for verification
    token : str
        The JWT token to decode
    algorithms : list[str] | None, optional
        List of allowed algorithms (default: ["EdDSA"])
    verify : bool, optional
        Whether to verify the signature (default: True)

    Returns
    -------
    dict[str, Any]
        The decoded payload

    Raises
    ------
    jwt.InvalidTokenError
        If the token is invalid or verification fails

    Examples
    --------
    >>> from hezor_common.security import generate_key_pair
    >>> private_key, public_key = generate_key_pair()
    >>> payload = {"user_id": "123", "role": "admin"}
    >>> token = encode_jwt(private_key, payload)
    >>> decoded = decode_jwt(public_key, token)
    >>> decoded["user_id"]
    '123'
    """
    if algorithms is None:
        algorithms = ["EdDSA"]

    options = {"verify_signature": verify}
    return pyjwt.decode(token, public_key, algorithms=algorithms, options=options)


def encode_jwt_with_pem(
    private_pem: bytes,
    payload: dict[str, Any],
    password: bytes | None = None,
    algorithm: str = "EdDSA",
    expires_in: int | None = None,
) -> str:
    """Encode a JWT token using PEM-encoded private key.

    This is a convenience function that deserializes the PEM key and encodes
    the JWT in one step.

    Parameters
    ----------
    private_pem : bytes
        PEM-encoded private key bytes
    payload : dict[str, Any]
        The payload data to encode
    password : bytes | None, optional
        Password for decrypting the private key (default: None)
    algorithm : str, optional
        The algorithm to use (default: "EdDSA")
    expires_in : int | None, optional
        Token expiration time in seconds (default: None)

    Returns
    -------
    str
        The encoded JWT token

    Examples
    --------
    >>> from hezor_common.security import generate_key_pair, serialize_private_key
    >>> private_key, _ = generate_key_pair()
    >>> private_pem = serialize_private_key(private_key)
    >>> payload = {"user_id": "123"}
    >>> token = encode_jwt_with_pem(private_pem, payload)
    >>> isinstance(token, str)
    True
    """
    private_key = deserialize_private_key(private_pem, password=password)
    return encode_jwt(private_key, payload, algorithm=algorithm, expires_in=expires_in)


def decode_jwt_with_pem(
    public_pem: bytes,
    token: str,
    algorithms: list[str] | None = None,
    verify: bool = True,
) -> dict[str, Any]:
    """Decode and verify a JWT token using PEM-encoded public key.

    This is a convenience function that deserializes the PEM key and decodes
    the JWT in one step.

    Parameters
    ----------
    public_pem : bytes
        PEM-encoded public key bytes
    token : str
        The JWT token to decode
    algorithms : list[str] | None, optional
        List of allowed algorithms (default: ["EdDSA"])
    verify : bool, optional
        Whether to verify the signature (default: True)

    Returns
    -------
    dict[str, Any]
        The decoded payload

    Examples
    --------
    >>> from hezor_common.security import generate_key_pair, serialize_public_key
    >>> private_key, public_key = generate_key_pair()
    >>> public_pem = serialize_public_key(public_key)
    >>> payload = {"user_id": "123"}
    >>> token = encode_jwt(private_key, payload)
    >>> decoded = decode_jwt_with_pem(public_pem, token)
    >>> decoded["user_id"]
    '123'
    """
    public_key = deserialize_public_key(public_pem)
    return decode_jwt(public_key, token, algorithms=algorithms, verify=verify)


def encode_jwt_with_file(
    private_key_path: str | Path,
    payload: dict[str, Any],
    password: bytes | None = None,
    algorithm: str = "EdDSA",
    expires_in: int | None = None,
) -> str:
    """Encode a JWT token using private key from file.

    This is a convenience function that reads the PEM key from file and encodes
    the JWT in one step.

    Parameters
    ----------
    private_key_path : str | Path
        Path to PEM-encoded private key file
    payload : dict[str, Any]
        The payload data to encode
    password : bytes | None, optional
        Password for decrypting the private key (default: None)
    algorithm : str, optional
        The algorithm to use (default: "EdDSA")
    expires_in : int | None, optional
        Token expiration time in seconds (default: None)

    Returns
    -------
    str
        The encoded JWT token

    Examples
    --------
    >>> from pathlib import Path
    >>> from hezor_common.security import generate_key_pair, serialize_private_key
    >>> private_key, _ = generate_key_pair()
    >>> # Save to temporary file
    >>> path = Path("/tmp/jwt_key.pem")
    >>> path.write_bytes(serialize_private_key(private_key))
    >>> payload = {"user_id": "123", "role": "admin"}
    >>> token = encode_jwt_with_file(path, payload)
    >>> isinstance(token, str)
    True
    """
    private_pem = Path(private_key_path).read_bytes()
    return encode_jwt_with_pem(
        private_pem,
        payload,
        password=password,
        algorithm=algorithm,
        expires_in=expires_in,
    )


def decode_jwt_with_file(
    public_key_path: str | Path,
    token: str,
    algorithms: list[str] | None = None,
    verify: bool = True,
) -> dict[str, Any]:
    """Decode and verify a JWT token using public key from file.

    This is a convenience function that reads the PEM key from file and decodes
    the JWT in one step.

    Parameters
    ----------
    public_key_path : str | Path
        Path to PEM-encoded public key file
    token : str
        The JWT token to decode
    algorithms : list[str] | None, optional
        List of allowed algorithms (default: ["EdDSA"])
    verify : bool, optional
        Whether to verify the signature (default: True)

    Returns
    -------
    dict[str, Any]
        The decoded payload

    Examples
    --------
    >>> from pathlib import Path
    >>> from hezor_common.security import generate_key_pair, serialize_public_key
    >>> private_key, public_key = generate_key_pair()
    >>> # Save to temporary file
    >>> path = Path("/tmp/jwt_pub.pem")
    >>> path.write_bytes(serialize_public_key(public_key))
    >>> payload = {"user_id": "123"}
    >>> token = encode_jwt(private_key, payload)
    >>> decoded = decode_jwt_with_file(path, token)
    >>> decoded["user_id"]
    '123'
    """
    public_pem = Path(public_key_path).read_bytes()
    return decode_jwt_with_pem(public_pem, token, algorithms=algorithms, verify=verify)


# Convenient aliases for common operations
encode = encode_jwt_with_file
decode = decode_jwt_with_file
