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

from sqlalchemy import create_engine
from db.migrate import migrate_db

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

    vision = google_vision.ImageAnnotatorClient()

    # db setup
    db_url = config['POSTGRES_URL']
    db_engine = create_engine(db_url)
    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['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())