import discord
import os
import json
import time
import io
import random
import logging

import lib.log as log
import lib.db as db

from lib.helpers import *
from lib.settings import *
from lib.parser import *
from lib.comfyui import *

async def on_message_or_reaction(client, obj):
    global RUN_TIME

    msg     = None
    chl     = None
    user    = None
    author  = None
    roles   = None
    rxn     = None

    msg_types = [discord.MessageType.default]

    # Handle reactions:
    if isinstance(obj, discord.RawReactionActionEvent):
        chl     = await client.fetch_channel(obj.channel_id)
        msg     = await chl.fetch_message(obj.message_id)
        user    = await client.fetch_user(obj.user_id)
        author  = await client.fetch_user(obj.message_author_id)
        roles   = obj.member.roles
        rxn     = obj

        msg_types.append(discord.MessageType.reply)

    # Handle direct messages from users:
    if isinstance(obj, discord.Message):
        msg     = obj
        chl     = obj.channel
        user    = msg.author
        author  = msg.author
        roles   = msg.author.roles
        rxn     = None

    if user == client.user:         return
    if msg.type not in msg_types:   return
    if user.bot:                    return

    #
    #
    #

    settings = get_settings()

    CONFIG_DIR = settings["config_dir"]

    chl_topic_parts     = []
    chl_topic_part_1    = ""
    if chl.topic is not None:
        chl_topic_parts     = chl.topic.split(",")
        chl_topic_part_1    = chl_topic_parts[0]

    using_workflow_path = None

    for path in [f"{CONFIG_DIR}workflows/{chl_topic_part_1}"]:
        try_path = f"{path}.json"
        if os.path.isfile(try_path):
            using_workflow_path = try_path
            log.write(f"Workflow {log.colors.fg.lightcyan}{try_path}")

    if using_workflow_path is None:
        return

    #
    #
    #

    setting_paths = [
        f"{CONFIG_DIR}guild_defaults/{chl.guild.id}",
        f"{CONFIG_DIR}guild_defaults/{chl.guild.name}",
        f"{CONFIG_DIR}category_defaults/{chl.category.id}",
        f"{CONFIG_DIR}category_defaults/{chl.category.name}",
        f"{CONFIG_DIR}channel_defaults/{chl.id}",
        f"{CONFIG_DIR}channel_defaults/{chl.name}",
    ]

    for i in chl_topic_parts:
        chl_topic_part  = i.strip()
        setting_paths   = setting_paths + [f"{CONFIG_DIR}defaults/{chl_topic_part}"]

    for r in roles:
        setting_paths = setting_paths + [f"{CONFIG_DIR}role_defaults/{r.name}"]

    setting_paths = setting_paths + [
        f"{CONFIG_DIR}user_defaults/{user.id}",
        f"{CONFIG_DIR}user_defaults/{user.name}",
    ]

    for path in setting_paths:
        try_path = f"{path}.json"
        log.write(f"Seeking {try_path}")
        if os.path.isfile(try_path):
            settings = merge_dicts(settings, read_json(try_path, {}))
            log.write(f"Merging {log.colors.fg.lightcyan}{try_path}")

    #
    #
    #

    ALLOW_REPEAT    = settings["allow_repeat"]
    SHOW_REPEAT     = settings["show_repeat"]
    REPEAT_EMOJI    = settings["repeat_emoji"]
    WAITING_EMOJI   = settings["waiting_emoji"]

    if rxn is not None:
        if not ALLOW_REPEAT:
            return

    # If allowed_users is specified, only listen to listed users:
    if "allowed_users" in settings.keys():
        allowed_users = settings["allowed_users"]
        if isinstance(allowed_users, list):
            if str(user.id) not in allowed_users and str(user.name) not in allowed_users:
                log.write(f"User {user.name} or their ID not in list of allowed users", log.colors.fg.lightred)
                return

    # If ignored_users is specified, ignore listed users:
    if "ignored_users" in settings.keys():
        ignored_users = settings["ignored_users"]
        if isinstance(ignored_users, list):
            if str(user.id) in ignored_users or str(user.name) in ignored_users:
                log.write(f"User {user.name} or their ID is explicitly being ignored", log.colors.fg.lightred)
                return

    # Read the found .json file:
    workflow_json = ""
    with open(using_workflow_path, "r") as file:
        workflow_json = file.read()

    # Break the user message into prompt parameters:
    params = get_prompt_parameters(msg.content, settings)

    for k in params.keys():
        v       = params[k]
        k       = k.upper()
        label   = f"__{k}__"
        log.write(f"Replacing {log.colors.fg.yellow}{label}{log.colors.reset} with {log.colors.fg.white}{v}")

    log.write_prompt(client, msg, chl, user, author, roles, rxn)

    try:
        time_start = time.perf_counter()

        # Indicate to the user something is happening:
        await msg.add_reaction(WAITING_EMOJI)
        await chl.typing()

        repeat_n_times  = 1 if "repeat_n_times" not in params.keys() else params["repeat_n_times"]
        all_attachments = []
        execution_time  = 0

        for i in range(0, repeat_n_times):
            workflow_json_clone = workflow_json
            params_clone        = params.copy()

            # If the seed is not specified, generate one:
            if "seed" not in params.keys() or params["seed"] is None:
                new_seed                = random.randint(1, 999_999_999_999_999)
                params_clone["seed"]    = new_seed
                log.write(f"Replacing {log.colors.fg.yellow}__SEED__{log.colors.reset} with {log.colors.fg.white}{new_seed}")

            # Make replacements, e.g.
            # __WIDTH__ to value of params["width"]
            # __POSITIVE__ to value of params["positive"]
            for k in params_clone.keys():
                v = params_clone[k]
                k = k.upper()

                if v is None:
                    continue

                workflow_json_clone = re.sub(rf"__{k}__", str(v), workflow_json_clone)

            # Must be valid JSON:
            workflow = json.loads(workflow_json_clone)
            response = await get_comfyui_generations(settings["api_url"], workflow)

            all_attachments = all_attachments + response["images"]
            execution_time  = execution_time + response["execution_time"]

        bytes_generated     = 0
        images_generated    = 0

        # Process 8 attachments at a time per one Discord message:
        while True:
            attachments_buffer  = all_attachments[:8]
            discord_files       = []
            
            for a in attachments_buffer:
                bytes_generated     = bytes_generated + len(a)
                images_generated    = images_generated + 1
                discord_file        = discord.File(io.BytesIO(a), filename="file.png")
                discord_files.append(discord_file)

            if len(discord_files) < 1: 
                discord_files = None

            post = await chl.send(files=discord_files, content=msg.content, reference=msg)
            del all_attachments[:8]

            if len(all_attachments) < 1:
                break

        if ALLOW_REPEAT:
            if SHOW_REPEAT:
                await msg.add_reaction(REPEAT_EMOJI)
                await post.add_reaction(REPEAT_EMOJI)
        
        await msg.remove_reaction(WAITING_EMOJI, client.user)

        time_end    = time.perf_counter()
        time_taken  = time_end - time_start

        log.write(f"{time_taken:0.2f}s taken.", log.colors.fg.lightgreen)

        #
        #
        #

        db.insert("prompts", {
            "handler_name":     settings["handler_name"],
            "api_url":          settings["api_url"],
            "bot_id":           client.user.id,
            "bot_name":         get_legacy_username(client.user),
            "guild_id":         chl.guild.id,
            "guild_name":       chl.guild.name,
            "channel_id":       chl.id,
            "channel_name":     chl.name,
            "author_id":        author.id,
            "author_name":      get_legacy_username(author),
            "user_id":          user.id,
            "user_name":        get_legacy_username(user),
            "prompt":           msg.content,
            "workflow_path":    using_workflow_path,
            "wait_time":        time_taken,
            "execution_time":   execution_time,
            "images_generated": images_generated,
            "bytes_generated":  bytes_generated
        })

    finally:
        lock = False