data_structures.py

"""Mega API information
=====================

- This file contains definitions for some of the properties within the API.
- TypeDict objects are the raw data returned by the http requests to the API itself.
- The dataclasses are the internal representation of thoses objects

"""

from __future__ import annotations

import dataclasses
import time
from enum import IntEnum
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, ClassVar, Final, Self, TypeAlias, TypedDict

if TYPE_CHECKING:
    from collections.abc import Mapping
    from typing import NotRequired

    from typing_extensions import ReadOnly

NodeID: TypeAlias = str
UserID: TypeAlias = str
TimeStamp: TypeAlias = int

SharedKeys: TypeAlias = dict[NodeID, tuple[int, ...]]


class ByteSize(int):
    def human_readable(self) -> str:
        """(ex: '150.5MiB')"""
        scale = 1024
        me = float(self)
        for unit in ("B", "KiB", "MiB", "GiB", "TiB", "PiB"):
            if abs(me) < scale:
                if unit == "B":
                    return f"{me:0.0f}{unit}"
                return f"{me:0.1f}{unit}"
            me /= scale

        return f"{me:0.1f}EiB"

    def __repr__(self) -> str:
        return self.human_readable()


class NodeType(IntEnum):
    FILE = 0
    FOLDER = 1
    ROOT_FOLDER = 2
    INBOX = 3
    TRASH = 4


class StorageStatus(IntEnum):
    UNKNOWN = -9
    GREEN = 0  # there is storage available
    ORANGE = 1  # almost full
    RED = 2  #  full
    # CHANGE = 3,     # obsolete
    PAYWALL = 4  # storage is full and user didn't remedy despite of warnings


class NodeSerialized(TypedDict):
    h: NodeID  # ID
    p: NodeID  # Parent ID
    u: NotRequired[UserID]  # Owner (user ID), not present in transfer.it nodes
    t: ReadOnly[NodeType]
    a: str  # Serialized attributes
    ts: TimeStamp  # creation date

    k: NotRequired[str]  # Serialized node keys
    su: NotRequired[str]  # Share owner (user ID), only present present in shared (public) files / folder
    sk: NotRequired[str]  # Share key, only present present in shared (public) files / folder


class ShareKeySerialized(TypedDict):
    h: NodeID  # ID of node for this key
    k: str  # key
    ha: str  # ???


class ShareKeySerialized2(TypedDict):
    h: NodeID  # ID of node for this key
    u: str  # Owner (user ID)
    r: int
    ts: TimeStamp


class GetNodesResponse(TypedDict):
    f: list[NodeSerialized]
    ok: list[ShareKeySerialized]
    s: list[ShareKeySerialized2]


class AttributesSerialized(TypedDict):
    n: NotRequired[ReadOnly[str]]  # Name
    lbl: NotRequired[int]  # label
    fav: NotRequired[bool]  # favorited


class FileInfoSerialized(TypedDict):
    s: int  # size
    at: str  # Serialized attributes
    fa: str  # Media file attributes (thumb, audio or video)
    g: NotRequired[str]  # direct download URL


_FIELDS_CACHE: dict[type, tuple[str, ...]] = {}


def _fields(cls: type) -> tuple[str, ...]:
    if fields := _FIELDS_CACHE.get(cls):
        return fields
    fields = _FIELDS_CACHE[cls] = tuple(f.name for f in dataclasses.fields(cls))
    return fields


class _DictDumper:
    __dataclass_fields__: ClassVar[dict[str, dataclasses.Field[Any]]]

    def dump(self) -> dict[str, Any]:
        """Get a JSONable dict representation of this object"""
        return dataclasses.asdict(self)

    def _shallow_dump(self) -> dict[str, Any]:
        return {name: getattr(self, name) for name in _fields(type(self))}


class _DictParser:
    __dataclass_fields__: ClassVar[dict[str, dataclasses.Field[Any]]]

    @classmethod
    def _filter_dict(cls, data: Mapping[str, Any], /) -> dict[str, Any]:
        return {k: v for k, v in data.items() if k in _fields(cls)}

    @classmethod
    def parse(cls, data: Mapping[str, Any], /) -> Self:
        return cls(**cls._filter_dict(data))


@dataclasses.dataclass(slots=True, frozen=True)
class FileInfo(_DictDumper):
    name: str
    size: ByteSize
    url: str | None

    _at: str

    @classmethod
    def parse(cls, resp: FileInfoSerialized) -> FileInfo:
        return FileInfo(
            name="",
            size=ByteSize(resp["s"]),
            url=resp.get("g"),
            _at=resp["at"],
        )


@dataclasses.dataclass(slots=True, frozen=True)
class Crypto(_DictDumper):
    key: tuple[int, int, int, int]
    iv: tuple[int, int]
    meta_mac: tuple[int, int]

    full_key: tuple[int, int, int, int, int, int, int, int]
    share_key: tuple[int, ...] | None

    @classmethod
    def compose(
        cls,
        key: tuple[int, ...],
        iv: tuple[int, ...],
        meta_mac: tuple[int, ...],
        node_type: NodeType = NodeType.FILE,
    ) -> Crypto:
        if node_type is NodeType.FILE:
            full_key = (
                key[0] ^ iv[0],
                key[1] ^ iv[1],
                key[2] ^ meta_mac[0],
                key[3] ^ meta_mac[1],
                *iv,
                *meta_mac,
            )
        else:
            full_key = *key, *iv, *meta_mac

        return Crypto(key, iv, meta_mac, full_key, None)  # pyright: ignore[reportArgumentType]

    @classmethod
    def decompose(
        cls,
        full_key: tuple[int, ...],
        node_type: NodeType = NodeType.FILE,
        share_key: tuple[int, ...] | None = None,
    ) -> Crypto:
        iv = full_key[4:6]
        meta_mac = full_key[6:8]
        key = full_key
        if node_type is NodeType.FILE:
            key = (
                key[0] ^ iv[0],
                key[1] ^ iv[1],
                key[2] ^ meta_mac[0],
                key[3] ^ meta_mac[1],
            )

        return Crypto(key, iv, meta_mac, full_key, share_key)  # pyright: ignore[reportArgumentType]

    @classmethod
    def from_dump(cls, dump: dict[str, Any]) -> Self:
        share_key = dump.pop("share_key")
        crypto = {k: tuple(v) for k, v in dump.items()}
        return cls(**crypto, share_key=tuple(share_key) if share_key else None)


# We populate attrs and crypto after instance creation
@dataclasses.dataclass(slots=True, frozen=True, order=True)
class Node(_DictDumper):
    id: NodeID
    parent_id: NodeID
    owner: UserID
    type: NodeType
    attributes: Attributes
    created_at: TimeStamp
    keys: MappingProxyType[UserID, str] = dataclasses.field(compare=False)
    share_owner: UserID | None
    share_key: str | None

    _a: str
    _crypto: Crypto

    @property
    def is_file(self) -> bool:
        return self.type is NodeType.FILE

    @property
    def is_folder(self) -> bool:
        return self.type is NodeType.FOLDER

    @classmethod
    def parse(cls, node: NodeSerialized) -> Node:
        owner = node.get("u", "")
        if k := node.get("k"):
            if owner:
                keys = dict(key_pair.split(":", 1) for key_pair in k.split("/") if ":" in key_pair)

            else:
                keys = {owner: k}
        else:
            keys: dict[str, str] = {}

        return Node(
            id=node["h"],
            parent_id=node["p"],
            owner=owner,
            created_at=node["ts"],
            type=NodeType(node["t"]),
            keys=MappingProxyType(keys),
            share_owner=node.get("su"),
            share_key=node.get("sk"),
            _a=node["a"],
            attributes=None,  # pyright: ignore[reportArgumentType]
            _crypto=None,  # pyright: ignore[reportArgumentType]
        )

    @classmethod
    def from_dump(cls, dump: dict[str, Any], /) -> Self:
        dump = dump | dict(  # noqa: C408
            type=NodeType[str(dump["type"]).upper()],
            attributes=Attributes(**dump["attributes"]) if dump["attributes"] else None,
            keys=MappingProxyType(dump["keys"]),
            _crypto=Crypto.from_dump(dump["_crypto"]) if dump["_crypto"] else None,
        )

        return cls(**dump)

    def dump(self) -> dict[str, Any]:
        """Get a JSONable dict representation of this object"""
        me = self._shallow_dump()
        me["type"] = self.type.name.lower()
        me["attributes"] = self.attributes.dump() if self.attributes else {}
        me["keys"] = dict(self.keys)
        me["_crypto"] = self._crypto.dump() if self._crypto else None
        return me


_LABELS: Final = "", "red", "orange", "yellow", "green", "blue", "purple", "grey"


@dataclasses.dataclass(slots=True, frozen=True)
class Attributes(_DictDumper):
    name: str
    label: str = ""
    favorited: bool = False

    @classmethod
    def parse(cls, attrs: AttributesSerialized) -> Self:
        return cls(
            name=attrs.get("n", ""),
            label=_LABELS[attrs.get("lbl", 0)],
            favorited=bool(attrs.get("fav")),
        )

    def serialize(self) -> AttributesSerialized:
        return {  # pyright: ignore[reportReturnType]
            key: value
            for key, value in [
                ("n", self.name),
                ("lbl", _LABELS.index(self.label)),
                ("fav", self.favorited),
            ]
            if value
        }


@dataclasses.dataclass(slots=True, frozen=True)
class AccountBalance(_DictDumper):
    amount: float
    currency: str

    @classmethod
    def parse(cls, balance: list[tuple[float, str]] | None) -> Self:
        amount, currency = balance[0] if balance else (0.0, "EUR")
        return cls(float(amount), str(currency))


class AccountStatsSerialized(TypedDict):
    # The NotRequired attributes are only available on PRO accounts

    mstrg: int  # maximum storage allowance
    bt: int  # "Base time age",  number of seconds since the start of the current quota buckets
    tah: list[int]  # The free IP-based quota buckets, 6 entries for 6 hours
    tar: int  # IP transfer reserved
    rua: int  # Actor reserved quota
    ruo: int  # Owner reserved quota
    cstrg: int  # total account storage usage
    cstrgn: dict[
        NodeID, list[int]
    ]  # NodeId -> [bytes, num_of_files, num_of_folders, versioned_bytes, num_versioned_files]

    balance: list[tuple[float, str]]
    uslw: (
        int  # The percentage (x100) indicating the limit at which the user is 'nearly' over. 98% for PRO, 90% for free.
    )
    usl: int  # User storage status
    subs: list[str]
    plans: list[str]
    features: list[str]

    caxfer: NotRequired[int]  # PRO transfer quota consumed by the user
    tuo: int  # Transfer usage by the owner on quota which hasn't yet been committed back to the API DB. Supplements caxfer
    csxfer: NotRequired[int]  # PRO transfer quota served to others
    tua: int  # Transfer usage served to other users which hasn't yet been committed back to the API DB. Supplements csxfer
    mxfer: NotRequired[int]  # maximum transfer allowance
    srvratio: float  # Ratio of PRO transfer quota that is able to be served to others
    suntil: NotRequired[TimeStamp]  # Expiration time of the currently active plan


@dataclasses.dataclass(slots=True, frozen=True)
class StorageMetrics(_DictDumper):
    bytes_used: ByteSize
    files: int
    folders: int

    @classmethod
    def parse(cls, metrics: list[int]) -> Self:
        bytes_used, files, folders = metrics[0:3]
        return cls(ByteSize(bytes_used), files, folders)


@dataclasses.dataclass(slots=True, frozen=True)
class StorageQuota(_DictDumper):
    used: ByteSize
    max: ByteSize

    percent: int
    threshold: int

    @property
    def ratio(self) -> float:
        return self.used / self.max

    @property
    def is_full(self) -> bool:
        return self.ratio >= 1

    @property
    def is_almost_full(self) -> bool:
        return self.ratio >= self.threshold

    @classmethod
    def parse(cls, data: Mapping[str, Any]) -> Self:
        max, used, threshold = map(ByteSize, (data["mstrg"], data["cstrg"], data["uslw"]))

        return cls(
            used=used,
            max=max,
            percent=int(used / max * 100),
            threshold=threshold // 100,
        )


@dataclasses.dataclass(slots=True, frozen=True)
class ProTransferQuota:
    used: ByteSize
    used_by_others: ByteSize
    max: ByteSize
    srv_ratio: float


@dataclasses.dataclass(slots=True, frozen=True)
class AccountStats(_DictParser, _DictDumper):
    storage: StorageQuota
    transfer_quota: ProTransferQuota | None
    balance: AccountBalance
    metrics: Mapping[NodeID, StorageMetrics]
    subs: tuple[str, ...]
    plans: tuple[str, ...]
    features: tuple[str, ...]
    storage_status: StorageStatus
    plan_expires: int | None
    current_quota_start_time: int
    _raw: AccountStatsSerialized = dataclasses.field(compare=False)

    def serialize(self) -> AccountStatsSerialized:
        return self._raw.copy()

    @classmethod
    def parse(cls, data: AccountStatsSerialized) -> Self:  # pyright: ignore[reportIncompatibleMethodOverride]
        transfer = None
        if max_transfer := data.get("mxfer"):
            assert "caxfer" in data
            assert "csxfer" in data

            transfer = ProTransferQuota(
                used=ByteSize(data["caxfer"]),
                used_by_others=ByteSize(data["caxfer"]),
                max=ByteSize(max_transfer),
                srv_ratio=data["srvratio"],
            )

        return cls(
            storage=StorageQuota.parse(data),
            balance=AccountBalance.parse(data.get("balance")),
            metrics={node_id: StorageMetrics.parse(stats) for node_id, stats in data["cstrgn"].items()},
            storage_status=StorageStatus(data["usl"]),
            plan_expires=data.get("suntil"),
            subs=tuple(data["subs"]),
            plans=tuple(data["plans"]),
            features=tuple(data["features"]),
            transfer_quota=transfer,
            current_quota_start_time=int(time.time() - data["bt"]),
            _raw=data,
        )


class UserResponse(TypedDict):
    u: UserID
    since: TimeStamp  # timestamp of account creation
    ipcc: str  # IP country code (ex: US)

    # These fields are not available when using a temp account (anonymous login)
    email: NotRequired[str]
    emails: NotRequired[list[str]]
    pemails: NotRequired[list[str]]
    name: NotRequired[str]