"""
Functions for managing FastAPI request context, translating into a suitable form for OwsClient.
This module is shamelessly lifted from fansifter-common.
"""

import base64
import json
from collections.abc import Iterable, Mapping
from contextvars import ContextVar, Token
from dataclasses import asdict, dataclass
from typing import Any, ClassVar, Dict
from uuid import UUID


def exclude(d: dict[str, Any], *, keys: Iterable[str]) -> dict[str, Any]:
    return {i: d[i] for i in d if i not in keys}


_correlation_id: ContextVar[str | None] = ContextVar("correlation_id", default=None)


def get_correlation_id() -> str | None:
    """Get the current correlation id."""
    return _correlation_id.get()


def set_correlation_id(correlation_id: str | None) -> Token[Any]:
    """Set the current correlation id."""
    return _correlation_id.set(correlation_id)


def reset_correlation_id(token: Token[Any]) -> None:
    """Reset the current correlation id."""
    _correlation_id.reset(token)


@dataclass
class RequestContext:
    """Request context."""

    PROFILES_CLAIM: ClassVar[str] = "https://grass.theorchard.com/profiles"
    USER_METADATA_CLAIM: ClassVar[str] = "https://grass.theorchard.com/user_metadata"
    REQUESTOR_SERVICE_NAME: ClassVar[str] = "ows_contributor"

    requestor_service_name: str | None = None
    identity_id: str | None = None
    identity_uuid: str | None = None
    profile_type: str | None = None
    profile_id: int | None = None
    profile_uuid: UUID | None = None
    authorization: str | None = None

    def dict(
        self, exclude_empty: bool = False, exclude_keys: Iterable[str] | None = None
    ) -> Dict[str, Any]:
        d = asdict(self)
        if exclude_keys is not None:
            return exclude(d, keys=exclude_keys)
        if exclude_empty:
            return {k: v for k, v in d.items() if v}
        return d

    @classmethod
    def from_headers(cls, headers: Mapping[str, str]) -> "RequestContext":
        authorization: str | None = headers.get("Authorization")
        claims = cls._decode_jwt_payload(cls._extract_bearer_token(authorization))

        # profile context headers
        profile_type = headers.get("Orchard-Profile-Type")
        profile_id = headers.get("Orchard-Profile-Id")
        profile_uuid = headers.get("Orchard-Profile-Uuid")

        claim_profiles = claims.get(cls.PROFILES_CLAIM, [])
        first_profile = claim_profiles[0] if claim_profiles else {}
        if not profile_type:
            profile_type = first_profile.get("profile_type")
        if not profile_id:
            profile_id = first_profile.get("profile_id")

        # identity
        identity_id = headers.get("Orchard-Identity-Id")
        identity_uuid = headers.get("Orchard-Identity-Uuid")

        user_metadata = claims.get(cls.USER_METADATA_CLAIM, {})
        token_identity_uuid = user_metadata.get("orchardIdentityId")
        if token_identity_uuid:
            if not identity_id:
                identity_id = token_identity_uuid
            if not identity_uuid:
                identity_uuid = token_identity_uuid

        return cls(
            requestor_service_name=cls.REQUESTOR_SERVICE_NAME,
            identity_id=identity_id,
            identity_uuid=identity_uuid,
            profile_type=profile_type,
            profile_id=cls._parse_profile_id(profile_id),
            profile_uuid=UUID(profile_uuid) if profile_uuid else None,
            authorization=authorization,
        )

    @staticmethod
    def _parse_profile_id(profile_id: Any) -> int | None:
        profile_id_str = str(profile_id) if profile_id is not None else ""
        return int(profile_id_str) if profile_id_str.isnumeric() else None

    @staticmethod
    def _extract_bearer_token(authorization: str | None) -> str | None:
        if not authorization:
            return None

        prefix = "Bearer "
        if authorization.startswith(prefix):
            return authorization[len(prefix) :].strip()
        return None

    @staticmethod
    def _decode_jwt_payload(token: str | None) -> Dict[str, Any]:
        if not token:
            return {}

        token_parts = token.split(".")
        if len(token_parts) < 2:
            return {}

        payload = token_parts[1]
        payload += "=" * (-len(payload) % 4)

        try:
            decoded_payload = base64.urlsafe_b64decode(payload.encode("ascii"))
            return json.loads(decoded_payload)
        except ValueError, TypeError, json.JSONDecodeError:
            return {}


_request_context: ContextVar[RequestContext | None] = ContextVar(
    "request_context", default=None
)


def get_request_context() -> RequestContext | None:
    """Get the current request context."""
    return _request_context.get()


def set_request_context(
    request_context: RequestContext | None,
) -> Token[Any]:
    """Set the current request context."""
    return _request_context.set(request_context)


def reset_request_context(token: Token[Any]) -> None:
    """Reset the current request context."""
    _request_context.reset(token)


def request_context_from_headers(headers: Mapping[str, str]) -> RequestContext:
    """Create a request context from headers."""
    return RequestContext.from_headers(headers)