from __future__ import annotations
from typing import Any, Dict, List, Optional
from pydantic import BaseModel, Field

Json = Dict[str, Any]

class MediaBlob(BaseModel):
    mime_type: str
    data_b64: str
    meta: Json = Field(default_factory=dict)

class ContextObject(BaseModel):
    text: Optional[str] = None
    images: List[MediaBlob] = Field(default_factory=list)
    audio: List[MediaBlob] = Field(default_factory=list)
    data: Json = Field(default_factory=dict)
    citations: List[Json] = Field(default_factory=list)

    def merged(self, other: "ContextObject") -> "ContextObject":
        return ContextObject(
            text="\n".join([t for t in [self.text, other.text] if t]),
            images=[*self.images, *other.images],
            audio=[*self.audio, *other.audio],
            data={**self.data, **other.data},
            citations=[*self.citations, *other.citations],
        )
