from collections.abc import Iterable, Sequence
from dataclasses import dataclass, field, fields
from datetime import datetime, timedelta, timezone
from typing import Any, Self

import jwt

from pyidentity.security.claims import Claim
from pyidentity.utils import normalize_datetime

STANDARD_CLAIMS: Sequence[str] = ("aud", "exp", "iat", "iss", "jti", "nbf", "sub")


def dump_claims_for_jwt(claims: Iterable[Claim]) -> dict[str, Any]:
    payload: dict[str, list | Any] = {}

    for claim in claims:
        value = payload.get(claim.type)

        if value is None:
            payload[claim.type] = claim.value
        else:
            if not isinstance(value, list):
                payload[claim.type] = [payload[claim.type]]

            payload[claim.type].append(claim.value)

    return payload


def load_claims_from_jwt(payload: dict[str, Any]) -> list[Claim]:
    claim_keys = payload.keys() - STANDARD_CLAIMS
    claims = []

    for key in claim_keys:
        value = payload[key]
        if isinstance(value, list):
            claims.extend(Claim(key, v) for v in value)
        else:
            claims.append(Claim(key, value))

        payload.pop(key)

    return claims


@dataclass
class JWTSecurityToken:
    signing_key: Any
    expires: datetime = field(metadata={"alias": "exp"})
    algorithm: str = field(default="HS256")
    audience: str | None = field(default=None, metadata={"alias": "aud"})
    claims: Iterable[Claim] = field(default_factory=list)
    issuer: str | None = field(default=None, metadata={"alias": "iss"})
    issuer_at: datetime = field(
        metadata={"alias": "iat"},
        default_factory=lambda: normalize_datetime(datetime.now(timezone.utc)),
    )
    not_before: datetime | None = field(default=None, metadata={"alias": "nbf"})
    subject: str | None = field(default=None, metadata={"alias": "sub"})
    jwt_id: str | None = field(default=None, metadata={"alias": "jti"})

    def __post_init__(self) -> None:
        if self.signing_key is None:
            raise jwt.InvalidKeyError()

        current_utc_timestamp = normalize_datetime(datetime.now(timezone.utc)).timestamp()

        if (
            isinstance(self.issuer_at, datetime)
            and (iat := normalize_datetime(self.issuer_at)).timestamp() <= current_utc_timestamp
        ):
            self.issuer_at = iat
        else:
            raise ValueError("iat must be a current or past time")

        if (
            isinstance(self.expires, datetime)
            and (exp := normalize_datetime(self.expires)).timestamp() >= current_utc_timestamp
        ):
            self.expires = exp
        else:
            raise ValueError("exp must be a datetime in the future")

        if self.not_before is not None:
            if (
                isinstance(self.not_before, datetime)
                and (nbf := normalize_datetime(self.not_before)).timestamp() <= current_utc_timestamp
            ):
                self.not_before = nbf
            else:
                raise ValueError("nbf must be a current or future time")

        if self.not_before and self.not_before >= self.expires:
            raise ValueError("exp must be after nbf")

    def encode(self) -> str:
        payload = dump_claims_for_jwt(self.claims)

        for _field in fields(self):
            if (value := getattr(self, _field.name)) and (alias := _field.metadata.get("alias")):
                payload[alias] = value

        return jwt.encode(payload=payload, key=self.signing_key, algorithm=self.algorithm)

    @classmethod
    def decode(
        cls,
        token: str | bytes,
        issuer_signing_key: Any,
        algorithms: Sequence[str] | None = None,
        options: dict[str, Any] | None = None,
        audience: str | Sequence[str] | None = None,
        issuer: str | Sequence[str] | None = None,
        leeway: float | timedelta = 0,
    ) -> Self:
        payload = jwt.decode(
            jwt=token,
            key=issuer_signing_key,
            algorithms=algorithms or ["HS256"],
            audience=audience,
            issuer=issuer,
            options=options,
            leeway=leeway,
        )
        audience = payload.pop("aud", None)
        expires = datetime.fromtimestamp(payload.pop("exp"), tz=timezone.utc)
        issuer_at = datetime.fromtimestamp(payload.pop("iat"), tz=timezone.utc)
        issuer = payload.pop("iss", None)
        subject = payload.pop("sub", None)
        not_before = payload.pop("nbf", None)
        jwt_id = payload.pop("jti", None)

        if not_before:
            not_before = datetime.fromtimestamp(payload.pop("nbf"), tz=timezone.utc)

        claims = load_claims_from_jwt(payload)
        return cls(
            signing_key="",
            audience=audience,
            expires=expires,
            not_before=not_before,
            issuer=issuer,
            issuer_at=issuer_at,
            subject=subject,
            jwt_id=jwt_id,
            claims=claims,
        )
