import os
import asyncio
import aiohttp
import json
import meilisearch

from dataclasses import dataclass
from typing import Dict, Any

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

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

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 MemeMessage:
    id: int
    url: str
    guild_name: str
    channel_name: str

    text: str | None
    vision_json: str | None

    # maybe don't need these
    # kind of extra
    message_json: Dict[str, Any]
    guild_json: Dict[str, Any]
    channel_json: Dict[str, Any]

    async def save_img(self, session: aiohttp.ClientSession):
        output_dir = self.output_dir()

        # there's a big param of some kind after the image url that we remove
        url_filename = os.path.basename(self.url).split('?')[0]
        (filename, ext) = os.path.splitext(url_filename)
        print(f"filename={filename} ext={ext}")
        os.makedirs(output_dir, exist_ok=True);

        async with session.get(self.url) as response:
            if response.status == 200:
                with open(f'{output_dir}/{self.id}{ext}', 'wb') as f:
                    f.write(await response.read())

    def output_dir(self):
        return f'./data/{self.guild_name}/{self.channel_name}'

    def annotate_image_request(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.output_dir()}/{self.id}.json', 'w') as f:
            json.dump(vars(self), f, indent=4)

    def to_meili_document(self):
        return {
            'id': self.id,
            'text': self.text,
            'url': self.url,
        }

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

async def download_channel_imgs(session, 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()

    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 = filter_message_img_urls(messages)

    memes = list(map(
        lambda msg: MemeMessage(
            id=msg[1]['id'],
            url=msg[0],
            guild_name=guild['name'],
            channel_name=channel['name'],
            guild_json=guild,
            channel_json=channel,
            message_json=msg[1],
            vision_json=None,
            text=None,
        )
    , urls))

    tasks = [meme.save_img(session) for meme in memes]
    await asyncio.gather(*tasks)

    [meme.annotate_image_request(vision) for meme in memes]
    [meme.persist() for meme in memes]

    client = meilisearch.Client('http://localhost:7700', 'aSampleMasterKey')
    documents = [meme.to_meili_document() for meme in memes]
    client.index('memes').add_documents(documents)

async def main():
    vision = google_vision.ImageAnnotatorClient()
    async with aiohttp.ClientSession() as session:
        channels = [download_channel_imgs(session, vision, channel_id) for channel_id in CHANNEL_IDS]
        memes = await asyncio.gather(*channels)

asyncio.run(main())