from pydantic import BaseModel, validator
from typing import Union, List, Optional
from datetime import datetime
from PIL import Image
from io import BytesIO

class Tweet(BaseModel):
    account: str
    url: str
    time: datetime
    text: str
    images: Optional[Union[List[Image.Image], List[BytesIO], List[bytes]]]
    video: Optional[str]
    
    @validator('images')
    def parse_images(cls, v):
        parsed_elem_list = list()
        for elem in v:
            if isinstance(elem, bytes):
                elem = BytesIO(elem)
            if isinstance(elem, BytesIO):
                elem = Image.open(elem)
            parsed_elem_list.append(elem)
        return parsed_elem_list
    
    class Config: # allows subtypes to be e.g. PIL.Image.Image without error
        arbitrary_types_allowed = True
    
    def as_db_dict(self):
        """Convert to database dictionary"""
        d = self.dict()
        d['time'] = d['time'].strftime('%Y-%m-%d %H:%M:%S')
        img_list = list()
        for img in d['images']:
            if isinstance(img, BytesIO):
                img = Image.open(img)
            buf = BytesIO()
            img.save(buf, format='PNG')
            img_list.append(buf.getvalue())
        d['images'] = img_list
        return d