xpost / atproto / models.py

import json
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any

from util.util import Result


URI = "at://"
URI_LEN = len(URI)


def cid_from_json(data: str | None) -> Result[str, str]:
    if data is None:
        return Result.err("Expected json, got None")

    try:
        return Result.ok(str(json.loads(data)["cid"]))
    except (json.JSONDecodeError, AttributeError, KeyError) as e:
        return Result.err(str(e))


class AtUri:
    @classmethod
    def record_uri(cls, uri: str) -> tuple[str, str, str]:
        if not uri.startswith(URI):
            raise ValueError(f"Ivalid record uri {uri}!")

        did, collection, rid = uri[URI_LEN:].split("/")
        if not (did and collection and rid):
            raise ValueError(f"Ivalid record uri {uri}!")

        return did, collection, rid


@dataclass(kw_only=True)
class StrongRef:
    uri: str
    cid: str

    def to_dict(self) -> dict[str, Any]:
        return {"uri": self.uri, "cid": self.cid}

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "StrongRef":
        return cls(uri=data["uri"], cid=data["cid"])


@dataclass(kw_only=True)
class ReplyRef:
    root: StrongRef
    parent: StrongRef

    def to_dict(self) -> dict[str, Any]:
        return {
            "root": self.root.to_dict(),
            "parent": self.parent.to_dict(),
        }


@dataclass(kw_only=True)
class FacetIndex:
    byte_start: int
    byte_end: int

    def to_dict(self) -> dict[str, int]:
        return {"byteStart": self.byte_start, "byteEnd": self.byte_end}


@dataclass(kw_only=True)
class FacetFeature(ABC):
    @abstractmethod
    def to_dict(self) -> dict[str, Any]:
        pass


@dataclass(kw_only=True)
class MentionFeature(FacetFeature):
    did: str

    def to_dict(self) -> dict[str, Any]:
        return {
            "$type": "app.bsky.richtext.facet#mention",
            "did": self.did,
        }


@dataclass(kw_only=True)
class LinkFeature(FacetFeature):
    uri: str

    def to_dict(self) -> dict[str, Any]:
        return {
            "$type": "app.bsky.richtext.facet#link",
            "uri": self.uri,
        }


@dataclass(kw_only=True)
class TagFeature(FacetFeature):
    tag: str

    def to_dict(self) -> dict[str, Any]:
        return {
            "$type": "app.bsky.richtext.facet#tag",
            "tag": self.tag,
        }


@dataclass(kw_only=True)
class Facet:
    index: FacetIndex
    features: list[FacetFeature] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return {
            "index": self.index.to_dict(),
            "features": [f.to_dict() for f in self.features],
        }


@dataclass(kw_only=True)
class ImageEmbed:
    image: bytes
    alt: str | None = None
    aspect_ratio: tuple[int, int] | None = None

    def to_dict(self, blob_ref: dict[str, Any]) -> dict[str, Any]:
        data: dict[str, Any] = {
            "image": blob_ref,
            "alt": self.alt or "",
        }
        if self.aspect_ratio:
            data["aspectRatio"] = {
                "width": self.aspect_ratio[0],
                "height": self.aspect_ratio[1],
            }
        return data


@dataclass(kw_only=True)
class VideoEmbed:
    video: bytes
    alt: str | None = None
    aspect_ratio: tuple[int, int] | None = None

    def to_dict(self, blob_ref: dict[str, Any]) -> dict[str, Any]:
        data: dict[str, Any] = {
            "$type": "app.bsky.embed.video",
            "video": blob_ref,
        }
        if self.alt:
            data["alt"] = self.alt
        if self.aspect_ratio:
            data["aspectRatio"] = {
                "width": self.aspect_ratio[0],
                "height": self.aspect_ratio[1],
            }
        return data


@dataclass(kw_only=True)
class RecordEmbed:
    record: StrongRef

    def to_dict(self) -> dict[str, Any]:
        return {
            "$type": "app.bsky.embed.record",
            "record": self.record.to_dict(),
        }


@dataclass(kw_only=True)
class RecordWithMediaEmbed:
    record: StrongRef
    media: ImageEmbed | VideoEmbed
    media_blob_ref: dict[str, Any]

    def to_dict(self) -> dict[str, Any]:
        media_data = self.media.to_dict(self.media_blob_ref)
        media_type = (
            "app.bsky.embed.images"
            if isinstance(self.media, ImageEmbed)
            else "app.bsky.embed.video"
        )
        return {
            "$type": "app.bsky.embed.recordWithMedia",
            "record": self.record.to_dict(),
            "media": {
                "$type": media_type,
                **media_data,
            },
        }


@dataclass(kw_only=True)
class SelfLabel:
    val: str

    def to_dict(self) -> dict[str, str]:
        return {"val": self.val}


@dataclass(kw_only=True)
class SelfLabels:
    values: list[SelfLabel] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return {
            "$type": "com.atproto.label.defs#selfLabels",
            "values": [v.to_dict() for v in self.values],
        }


@dataclass(kw_only=True)
class PostRecord:
    text: str
    created_at: str
    facets: list[Facet] | None = None
    embed: dict[str, Any] | None = None
    reply: ReplyRef | None = None
    langs: list[str] | None = None
    labels: SelfLabels | None = None

    def to_dict(self) -> dict[str, Any]:
        data: dict[str, Any] = {
            "$type": "app.bsky.feed.post",
            "text": self.text,
            "createdAt": self.created_at,
        }
        if self.facets:
            data["facets"] = [f.to_dict() for f in self.facets]
        if self.embed:
            data["embed"] = self.embed
        if self.reply:
            data["reply"] = self.reply.to_dict()
        if self.langs:
            data["langs"] = self.langs
        if self.labels:
            data["labels"] = self.labels.to_dict()
        return data


@dataclass(kw_only=True)
class CreateRecordResponse:
    uri: str
    cid: str
    commit: dict[str, Any] | None = None

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "CreateRecordResponse":
        return cls(
            uri=data.get("uri", ""),
            cid=data.get("cid", ""),
            commit=data.get("commit"),
        )


@dataclass(kw_only=True)
class RepostRecord:
    subject: StrongRef
    created_at: str

    def to_dict(self) -> dict[str, Any]:
        data: dict[str, Any] = {
            "$type": "app.bsky.feed.repost",
            "createdAt": self.created_at,
            "subject": self.subject.to_dict(),
        }
        return data


@dataclass(kw_only=True)
class ThreadGate:
    post: str = ""
    created_at: str
    allow: list[dict[str, Any]] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        data: dict[str, Any] = {
            "$type": "app.bsky.feed.threadgate",
            "post": self.post,
            "createdAt": self.created_at,
        }
        if self.allow is not None:
            data["allow"] = self.allow
        return data


@dataclass(kw_only=True)
class PostGate:
    post: str
    created_at: str
    detached_embedding_uris: list[str] | None = None
    embedding_rules: list[dict[str, Any]] | None = None

    def to_dict(self) -> dict[str, Any]:
        data: dict[str, Any] = {
            "$type": "app.bsky.feed.postgate",
            "post": self.post,
            "createdAt": self.created_at,
        }
        if self.detached_embedding_uris is not None:
            data["detachedEmbeddingUris"] = self.detached_embedding_uris
        if self.embedding_rules is not None:
            data["embeddingRules"] = self.embedding_rules
        return data