_service.py

"""Service of the platform module."""

import json
import time
from functools import cached_property
from http import HTTPStatus
from typing import Any

import httpx
from aignx.codegen.models import MeReadResponse as Me
from aignx.codegen.models import OrganizationReadResponse as Organization
from aignx.codegen.models import UserReadResponse as User
from loguru import logger
from pydantic import BaseModel, computed_field

from aignostics.constants import INTERNAL_ORGS
from aignostics.utils import BaseService, Health, user_agent

from ._authentication import get_token, remove_cached_token, verify_and_decode_token
from ._client import Client
from ._settings import Settings


class TokenInfo(BaseModel):
    """Class to store token information."""

    issuer: str  # iss
    issued_at: int  # iat
    expires_at: int  # exp
    scope: list[str]  # scope
    audience: list[str]  # aud
    authorized_party: str  # azp

    org_id: str  # org_id
    role: str  # https://audience/role

    @computed_field  # type: ignore[prop-decorator]
    @property
    def expires_in(self) -> int:
        """Calculate seconds until token expires.

        Returns:
            int: Number of seconds until the token expires. Negative if already expired.
        """
        return self.expires_at - int(time.time())

    @classmethod
    def from_claims(cls, claims: dict[str, Any]) -> "TokenInfo":
        """Create TokenInfo from JWT claims.

        Args:
            claims: JWT token claims dictionary.

        Returns:
            TokenInfo: Token information extracted from claims.
        """
        audience = claims["aud"] if isinstance(claims["aud"], list) else [claims["aud"]]

        return cls(
            issuer=claims["iss"],
            issued_at=claims["iat"],
            expires_at=claims["exp"],
            scope=claims["scope"].split(),
            audience=audience,
            authorized_party=claims["azp"],
            org_id=claims["org_id"],
            role=claims.get(audience[0] + "/role", "member"),
        )


class UserInfo(BaseModel):
    """Class to store info about the user."""

    role: str  # token.CLAIM_ROLE
    token: TokenInfo
    user: User
    organization: Organization

    @classmethod
    def from_claims_and_me(cls, claims: dict[str, Any], me: Me) -> "UserInfo":
        """Create UserInfo from JWT claims and optional auth0 userinfo.

        Args:
            claims (dict[str, Any]): JWT token claims dictionary.
            me (Me): Info about calling user and their oganisation.

        Returns:
            UserInfo: User information extracted from claims.
        """
        token = TokenInfo.from_claims(claims)
        return cls(
            role=token.role,
            token=token,
            user=me.user,
            organization=me.organization,
        )

    @cached_property
    def is_internal_user(self) -> bool:
        """Check if the user is an internal user.

        Returns:
            bool: True if it is an internal user, False otherwise.
        """
        return bool(self.organization and self.organization.name and self.organization.name.lower() in INTERNAL_ORGS)

    def model_dump_secrets_masked(self) -> dict[str, Any]:
        """Dump model to dict with sensitive organization and user secrets masked.

        Returns:
            dict[str, Any]: Dictionary representation with sensitive organization and user fields masked.
        """
        data = self.model_dump(mode="json")

        # Define mapping of data keys to their sensitive fields
        sensitive_fields_mapping = {
            "organization": [
                "aignostics_bucket_hmac_secret_access_key",
                "aignostics_logfire_token",
                "aignostics_sentry_dsn",
            ],
            "user": ["email"],
        }

        # Mask sensitive fields for each data section
        for data_key, secret_fields in sensitive_fields_mapping.items():
            if data.get(data_key):
                section_data = data[data_key]
                for field_name in secret_fields:
                    if section_data.get(field_name):
                        original_value = section_data[field_name]
                        section_data[field_name] = f"***MASKED({len(original_value)})***"

        return data

    def model_dump_json_secrets_masked(self) -> str:
        """Dump model to JSON with sensitive organization secrets masked.

        Returns:
            str: JSON representation with aignostics_bucket_hmac_access_key_id and
                 aignostics_bucket_hmac_secret_access_key masked.
        """
        return json.dumps(self.model_dump_secrets_masked())


# Services derived from BaseService and exported by modules via their __init__.py are automatically registered
# with the system module, enabling for dynamic discovery of health, info and further functionality.
class Service(BaseService):
    """Service of the application module."""

    _settings: Settings

    def __init__(self) -> None:
        """Initialize service."""
        super().__init__(Settings)  # automatically loads and validates the settings

    async def info(self, mask_secrets: bool = True) -> dict[str, Any]:
        """Determine info of this service.

        Args:
            mask_secrets (bool): Whether to mask sensitive information in the output.

        Returns:
            dict[str,Any]: The info of this service.
        """
        user_info = None
        try:
            user_info = self.get_user_info()
        except RuntimeError:
            message = "Failed to retrieve user info for system info."
            logger.warning(message)
        return {
            "userinfo": (user_info.model_dump_secrets_masked() if mask_secrets else user_info.model_dump(mode="json"))
            if user_info
            else None,
        }

    @staticmethod
    def _health_from_response(response: httpx.Response) -> Health:
        """Map a PAPI health response to a Health status.

        Handles non-200 status codes, unparseable bodies, and the three recognised
        ``status`` values (``"UP"``, ``"DEGRADED"``, ``"DOWN"``).

        Args:
            response: httpx response from the ``/health`` endpoint.

        Returns:
            Health: ``UP``, ``DEGRADED``, or ``DOWN`` derived from the response.
        """
        if response.status_code != HTTPStatus.OK:
            logger.error("Aignostics Platform API returned '{}'", response.status_code)
            return Health(
                status=Health.Code.DOWN, reason=f"Aignostics Platform API returned status '{response.status_code}'"
            )

        try:
            body = response.json()
        except Exception:
            return Health(status=Health.Code.DOWN, reason="Aignostics Platform API returned unparseable response")

        api_status = body.get("status")
        if api_status == "UP":
            return Health(status=Health.Code.UP)
        if api_status == "DEGRADED":
            reason = body.get("reason") or "Aignostics Platform API is DEGRADED"
            logger.warning("Aignostics Platform API is DEGRADED: {}", reason)
            return Health(status=Health.Code.DEGRADED, reason=reason)
        return Health(
            status=Health.Code.DOWN,
            reason=f"Aignostics Platform API returned unknown status '{api_status}'",
        )

    async def _determine_api_public_health(self) -> Health:
        """Determine healthiness and reachability of Aignostics Platform API.

        - Checks if health endpoint is reachable and returns 200 OK
        - Parses the response body to detect DEGRADED status
        - Uses httpx for a direct connection check without authentication

        Returns:
            Health: The healthiness of the Aignostics Platform API via basic unauthenticated request.
        """
        try:
            async with httpx.AsyncClient() as client:
                response = await client.get(
                    f"{self._settings.api_root}/health",
                    headers={"User-Agent": user_agent()},
                    timeout=self._settings.health_timeout,
                )
            return self._health_from_response(response)
        except Exception as e:
            logger.exception("Issue with Aignostics Platform API")
            return Health(status=Health.Code.DOWN, reason=f"Issue with Aignostics Platform API: '{e}'")

    async def _determine_api_authenticated_health(self) -> Health:
        """Determine healthiness and reachability of Aignostics Platform API via authenticated request.

        Parses the response body to detect DEGRADED status.

        Returns:
            Health: The healthiness of the Aignostics Platform API when trying to reach via authenticated request.
        """
        try:
            token = get_token(use_cache=True)
            async with httpx.AsyncClient() as client:
                response = await client.get(
                    f"{self._settings.api_root}/health",
                    headers={
                        "User-Agent": user_agent(),
                        "Authorization": f"Bearer {token}",
                    },
                    timeout=self._settings.health_timeout,
                )
            return self._health_from_response(response)
        except Exception as e:
            logger.exception("Issue with Aignostics Platform API")
            return Health(status=Health.Code.DOWN, reason=f"Issue with Aignostics Platform API: '{e}'")

    async def health(self) -> Health:
        """Determine health of this service.

        Returns:
            Health: The health of the service.
        """
        return Health(
            status=Health.Code.UP,
            components={
                "api_public": await self._determine_api_public_health(),
                "api_authenticated": await self._determine_api_authenticated_health(),
            },
        )

    @staticmethod
    def login(relogin: bool = False) -> bool:
        """Login.

        Args:
            relogin (bool): If True, forces a re-login even if a token is cached.

        Returns:
            bool: True if successfully logged in, False if login failed
        """
        if relogin:
            Service.logout()
        try:
            _ = get_token()
            return True
        except RuntimeError as e:
            message = f"Error during login: {e!s}"
            logger.exception(message)
            return False

    @staticmethod
    def logout() -> bool:
        """Logout if authenticated.

        Deletes the cached authentication token if existing.

        Returns:
            bool: True if successfully logged out, False if not logged in.
        """
        return remove_cached_token()

    @staticmethod
    def get_user_info(relogin: bool = False) -> UserInfo:
        """Get user information from authentication token.

        Args:
            relogin (bool): If True, forces a re-login even if a token is cached.

        Returns:
            UserInfo | None: User information if successfully authenticated, None if login failed.

        Raises:
            RuntimeError: If the token cannot be verified or decoded.
        """
        if relogin:
            Service.logout()
        try:
            return UserInfo.from_claims_and_me(verify_and_decode_token(get_token()), Client().me())  # pyright: ignore[reportArgumentType]
        except RuntimeError as e:
            message = f"Error during login: {e!s}"
            logger.exception(message)
            raise