memestream/main.py

import os
import asyncio
import json
from urllib3.util import parse_url
from pathlib import Path
from io import BytesIO
from dataclasses import dataclass
from uuid import uuid4, UUID
from typing import Dict, Any, List

import aiohttp

from aiobotocore.session import get_session as aiobotocore_get_session
from types_aiobotocore_s3.client import S3Client

from PIL import Image
from dotenv import dotenv_values, load_dotenv

from google.cloud import vision as google_vision
from google.protobuf.json_format import MessageToDict, MessageToJson
from google.cloud.vision_v1 import types
from minio import Minio
import meilisearch
from urllib3.util import Url, url

import uploader

# need to change the env for the google api calls
load_dotenv()
config = dotenv_values(".env")

DATA_DIR = './data/memestream'
MEILI_INDEX = 'memes'
MEILI_URL = 'http://localhost:7700'

CHANNEL_IDS = [
    367936426962190336, # TPM shitposts
    1095060717117374687, # aella memes
]

headers = {
    'Authorization': config['DISCORD_AUTH'],
}

# TODO
# [ ] Backfill last N images from a channel
# [ ] Set oldest_img_id for channel
# [ ] save img as json + imagefile
# [ ] check extensions (ignore mp4 maybe)

@dataclass
class MemeDocument:
    id: UUID
    url: Url
    text: str
    tags: List[str]

@dataclass
class DiscordAttachmentMessage:
    id: UUID
    msg_id: int
    url: str
    upload_url: str
    guild_name: str
    channel_name: str
    vision_json: str | None
    message_json: Dict[str, Any]
    guild_json: Dict[str, Any]
    channel_json: Dict[str, Any]

    def document_dir(self):
        return f'{DATA_DIR}/guilds/{self.guild_name}/{self.channel_name}'

    def file_dir(self):
        return f"{DATA_DIR}/imgs/"

    def detect_text(self, vision):
        request = {
            'image': { 'source': { 'image_uri': self.url } },
            'features': [types.Feature(type=types.Feature.Type.TEXT_DETECTION)],
        }

        response: types.AnnotateImageResponse = vision.annotate_image(request=request)
        self.vision_json = json.loads(types.AnnotateImageResponse.to_json(response))

        if len(response.text_annotations) > 0:
            self.text = response.text_annotations[0].description

        return request

    def persist(self):
        with open(f'{self.document_dir()}/{self.id}.json', 'w') as f:
            json.dump(vars(self), f, indent=4)

def filter_message_img_urls(messages):
    url_lists = map(
        lambda msg: list(map(lambda attachment: (parse_url(attachment['url']), msg), msg['attachments']))
    , messages)
    return [item for group in url_lists for item in group]

async def download_channel_attachments(session, storage, vision, channel_id):
    async def fetch_channel_messages(channel_id):
        url = f"https://discord.com/api/v9/channels/{channel_id}/messages"
        async with session.get(url, headers=headers) as res:
            return await res.json()

    async def fetch_guild_details(guild_id):
        url = f"https://discord.com/api/v9/guilds/{guild_id}"
        async with session.get(url, headers=headers) as res:
            data = await res.json()

            keys_to_remove = ['emojis', 'stickers', 'roles']
            [data.pop(key, None) for key in keys_to_remove]

            return data

    async def fetch_channel(channel_id):
        url = f"https://discord.com/api/v9/channels/{channel_id}"
        async with session.get(url, headers=headers) as res:
            return await res.json()

    async def download_attachment(
        session: aiohttp.ClientSession,
        storage: S3Client,
        url: Url,
        message,
    ):
        # there's a big param of some kind after the image url that we remove
        if not url.path:
            raise RuntimeError(f"url missing path url={url}")

        storage_bucket = config['UPLOAD_BUCKET']
        if not storage_bucket:
            raise RuntimeError(f"config missing UPLOAD_BUCKET")

        (filename, ext) = os.path.splitext(os.path.basename(url.path))
        # TODO
        # FILTER WANTED TYPES

        id = uuid4()

        async with session.get(url.url) as response:
            if response.status == 200:

                # make a webp conversion for obvious image types
                convert_extensions = ['jpg', 'png', 'jpeg', 'webp']
                res_bytes = await response.read()
                object_name: str

                # trim the .
                if ext.lower()[1:] in convert_extensions:
                    image = Image.open(BytesIO(res_bytes))

                    webp_output = BytesIO()
                    image.save(webp_output, format='WebP')
                    webp_output.seek(0)

                    object_name = f"{id}.webp"
                    print(f"uploading object_name={object_name} filename={filename} ext={ext}")
                    return await storage.put_object(Bucket=storage_bucket, Key=object_name, Body=webp_output)

                object_name = f"{id}{ext}"
                print(f"uploading object_name={object_name} filename={filename} ext={ext}")
                return await storage.put_object(Bucket=storage_bucket, Key=object_name, Body=res_bytes)

                # add discord id to server processed msgs list
                # if there then nothing to do

                # return MemeDocument(
                #     id=id,
                #     url=self.upload_url,
                #     text=self.text or '',
                #     tags=[self.guild_name, self.channel_name],
                # )

    channel = await fetch_channel(channel_id)
    guild = await fetch_guild_details(channel['guild_id'])
    print(f"Fetching messages for channel guild={guild['name']} name={channel['name']} id={channel['id']}")
    messages = await fetch_channel_messages(channel_id)
    urls_and_messages = filter_message_img_urls(messages)

    tasks = [download_attachment(session, storage, url, message) for (url, message) in urls_and_messages]
    await asyncio.gather(*tasks)

    # [attachment.detect_text(vision) for attachment in attachments]
    # [attachment.persist() for attachment in attachments]

    # client = meilisearch.Client(MEILI_URL, config['MEILI_MASTER_KEY'])

    # # have to vars the dataclass out for meili
    # documents = [vars(attachment.to_meme()) for attachment in attachments]
    # client.index(MEILI_INDEX).add_documents(documents)

async def check_bucket_exists(storage, bucket_name):
    try:
        await storage.head_bucket(Bucket=bucket_name)
        return True
    except storage.exceptions.ClientError as e:
        error_code = e.response.get('Error', {}).get('Code')
        if error_code == '404':
            return False  # Bucket does not exist
        else:
            raise  # Reraise other exceptions

async def main():
    assert config['MINIO_URL']

    vision = google_vision.ImageAnnotatorClient()

    # storage = Minio(config['MINIO_URL'],
    #     access_key=config['MINIO_ACCESS_KEY'],
    #     secret_key=config['MINIO_SECRET_KEY'],
    # )

    aioboto = aiobotocore_get_session()
    async with aioboto.create_client(
        's3',
        endpoint_url=config['STORAGE_URL'],
        aws_access_key_id=config['STORAGE_ACCESS_KEY'],
        aws_secret_access_key=config['STORAGE_SECRET_KEY'],
    ) as storage:
        storage: S3Client

        storage_bucket = config['UPLOAD_BUCKET']
        assert storage_bucket

        # Make the bucket if it doesn't exist.
        exists = await check_bucket_exists(storage, storage_bucket)
        if not exists:
            await storage.create_bucket(Bucket=storage_bucket)
            print("Created bucket", storage_bucket)
        else:
            print("Bucket", storage_bucket, "already exists")

        async with aiohttp.ClientSession() as session:
            channels = [download_channel_attachments(session, storage, vision, channel_id) for channel_id in CHANNEL_IDS]
            memes = await asyncio.gather(*channels)


asyncio.run(main())