memestream/main.py

import os
import asyncio
import json
import itertools

from urllib3.util import parse_url
from pathlib import Path
from io import BytesIO
from dataclasses import dataclass
from uuid import uuid4, UUID

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
import meilisearch
from urllib3.util import Url, url

from sqlalchemy import create_engine
from db.migrate import migrate_db
from db.attachment import Attachment
from db.meme import Meme

# 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]

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

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

    return request

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 upload_object(object_name, body):
        storage_bucket = config['STORAGE_BUCKET']
        if not storage_bucket:
            raise RuntimeError(f"config missing STORAGE_BUCKET")
        return await storage.put_object(Bucket=storage_bucket, Key=object_name, Body=body)

    async def download_attachment(
        session: aiohttp.ClientSession,
        storage: S3Client,
        attachment: Attachment
    ):
        async with session.get(attachment.url) as response:
            if response.status == 200:
                # make a webp conversion for obvious image types
                res_bytes = await response.read()

                (filename, ext) = attachment.file_details()
                object_name = attachment.object_name()
                print(f"uploading object_name={object_name} filename={filename} ext={ext}")

                if attachment.convert_to_webp():
                    image = Image.open(BytesIO(res_bytes))
                    webp_output = BytesIO()
                    image.save(webp_output, format='WebP')
                    webp_output.seek(0)
                    return await upload_object(object_name=object_name, body=webp_output)

                return await upload_object(object_name=object_name, body=res_bytes)

    def create_attachments(message, channel, guild):
        urls = [parse_url(attachment['url']) for attachment in message['attachments']]

        def generate_attachment(url: Url):
            if not url.path:
                return None

            (filename, ext) = os.path.splitext(os.path.basename(url.path))
            extensions = ['jpg', 'png', 'jpeg', 'webp', 'mov', 'gif', 'mp4']
            if ext.lower()[1:] not in extensions:
                return None

            return Attachment(
                id=uuid4(),
                msg_id=int(message['id']),
                message=message,
                guild=guild,
                url=url.url,
                channel=channel
            )

        return [generate_attachment(url) for url in urls]

    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)

    # at this point we have the channel messages, the guild (server) and channel info

    # first thing we have to do is parse out the urls in the msg attachments
    # once we have those we can generate where we should put them in our storage
    # should only persist them later once the processing is done
    attachment_groups = [create_attachments(message, channel, guild) for message in messages]

    # super wack flatmap and filter
    attachments = list(filter(None, itertools.chain.from_iterable(attachment_groups)))

    # NOW
    # -------------------
    # insert attachment
    # upload to s3
    # fetch vision for relevant attachments
    # create meme from attachment and vision
    # insert meme
    # add documents to meili

    # 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['STORAGE_URL']
    assert config['POSTGRES_URL']

    vision = google_vision.ImageAnnotatorClient()

    # db setup
    db_url = config['POSTGRES_URL']
    db_engine = create_engine(db_url, echo=True)
    migrate_db(db_engine, db_url)

    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['STORAGE_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())