import json
import logging
import re
import time

import requests
from rich import traceback

from bots.data import config, db, marseygen_queue
from bots.clients.drama import DramaClient
from bots.clients.runpod import RunpodClient

traceback.install()

# Set up logging
logging.basicConfig(
    filename=f"{config['data_dir']}/logs/marseygen.log",
    filemode="a",
    format="%(name)s - %(levelname)s - %(message)s",
    level=logging.INFO,
)
logger = logging.getLogger("marseygen")


def main():
    if "marseygen_id" not in db:
        db["marseygen_id"] = 0
    if "marseygen_last_processed" not in db:
        db["marseygen_last_processed"] = 0

    runpod_client = RunpodClient(logger=logger)

    # If requests are still in the queue and
    # an instance is running, return and let it complete.
    if marseygen_queue and runpod_client.is_running("marseygen"):
        return

    drama_client = DramaClient(config["marseygen"]["token"], logger=logger)

    # Fetch new requests and add each to the queue.
    comments, newest_id = drama_client.fetch_new_comments(
        after=db["marseygen_last_processed"]
    )
    db["marseygen_last_processed"] = newest_id

    # Add new requests to queue.
    for comment in comments:
        prompts = re.findall(
            r"^!sd (.*)$", comment["body"], re.MULTILINE | re.IGNORECASE
        )
        prompts = prompts[:5]
        for prompt in prompts:
            prompt = prompt.replace("`", "")
            prompt = prompt.replace("marsey", "Marsey")
            db["marseygen_id"] += 1
            marseygen_queue[str(db["marseygen_id"])] = (comment, prompt)

    # Create instance if there are requests to be fulfilled.
    if marseygen_queue:
        runpod_client.create_instance(
            f"{config['data_dir']}/marseygen/create_runpod.json"
        )


if __name__ == "__main__":
    main()