import io
import json
import re
import psycopg
from psycopg.rows import dict_row
from pprint import pprint
import openai
import queue
from contextlib import contextmanager
from concurrent.futures import ThreadPoolExecutor, as_completed
import sys
import threading

from tqdm import tqdm

slots = queue.Queue()
hosts = [
    'localhost:8080',
    'localhost:8080',
    'localhost:8081',
    'localhost:8081',
    'localhost:8082',
    '192.168.88.16:8082'
]
for h in hosts:
    slots.put(
        openai.OpenAI(
            base_url=f"http://{h}/v1", api_key="sk-XWptxvNyY6SLEnhD"
        )
    )
num_backends = len(hosts)


@contextmanager
def lease_slot():
    slot = slots.get()
    try:
        yield slot
    finally:
        slots.put(slot)


dbconnstr = "host=localhost dbname=composer_new_dataset_v3"

# with open("representative_file_descriptions.txt", "r") as f:
#     file_ids = [int(line.strip()) for line in f]


def control_to_space(text):
    return "".join(c if c.isprintable() else " " for c in text).strip()

def prepare_file(file_id):
    with psycopg.connect(dbconnstr, row_factory=dict_row) as conn, conn.cursor() as cur:
        cur.execute(
            "SELECT dataset_name, metadata FROM files WHERE id = %s", (file_id,)
        )
        file = cur.fetchone()

        cur.execute(
            "SELECT type, description FROM file_descriptions WHERE file_id = %s AND type IN ('file_name', 'imslp_text', 'musescore_text')",
            (file_id,),
        )
        file_descs = cur.fetchall()
        cur.execute("SELECT metadata FROM instruments WHERE file_id = %s", (file_id,))
        instrument_metas = [m["metadata"] for m in cur.fetchall()]
        cur.execute(
            "SELECT name, lyrics, texts FROM tracks WHERE file_id = %s", (file_id,)
        )
        tracks = cur.fetchall()

        prompt_suf = f"Dataset name: {file['dataset_name']}\n"
        file_names = [
            control_to_space(f["description"])
            for f in file_descs
            if f["type"] == "file_name" and control_to_space(f["description"]) != ""
        ]
        if file_names:
            prompt_suf += "File names (aliases):\n " + "\n ".join(file_names) + "\n"
        file_metadata = file["metadata"]
        if file_metadata is None:
            file_metadata = {}
        mu_meta = file_metadata.get("meta", {})
        if "composer" in mu_meta:
            prompt_suf += f"Composer: {control_to_space(mu_meta['composer'])}\n"
        if "workTitle" in mu_meta:
            prompt_suf += f"Work title: {control_to_space(mu_meta['workTitle'])}\n"
        mu_descs = [
            control_to_space(f["description"])
            for f in file_descs
            if f["type"] == "musescore_text"
            and control_to_space(f["description"]) != ""
        ]
        if mu_descs:
            prompt_suf += (
                "Musescore extended metadata:\n " + "\n ".join(mu_descs) + "\n"
            )

        track_names = [
            control_to_space(t["name"])
            for t in tracks
            if t.get("name", None) is not None and control_to_space(t["name"]) != ""
        ]
        instrument_meta = [
            (
                control_to_space(
                    m.get("instrument")
                    + (
                        (" (" + m.get("instrumentId") + ")")
                        if "instrumentId" in m
                        else ""
                    )
                )
            )
            for m in instrument_metas
            if m is not None and "instrument" in m
        ]
        if instrument_meta:
            prompt_suf += (
                "Instrument metadata:\n "
                + "\n ".join(f"{i+1}. {m}" for i, m in enumerate(instrument_meta))
                + "\n"
            )
        elif track_names:
            prompt_suf += (
                "Track names:\n "
                + "\n ".join(f"{i+1}. {t}" for i, t in enumerate(track_names))
                + "\n"
            )
        track_lyrics = "\n".join(
            [
                control_to_space(line)
                for t in tracks
                if t["lyrics"] is not None
                for line in t["lyrics"].splitlines()
                if control_to_space(line) != ""
            ]
        )
        if track_lyrics:
            prompt_suf += (
                "Track lyrics (first 500 bytes):\n "
                + "\n ".join(track_lyrics[:500].split("\n"))
                + "\n"
            )

        # print("File:")
        # pprint(file)
        # print("File descriptions:")
        # pprint(file_descs)
        # print("Instrument metadata:")
        # pprint(instrument_metas)
        # print("Track metadata:")
        # pprint(tracks)
        # print("---\n\n")

        # prompt = "Write 5 short English prompts that will be used to train a generative AI model to produce short segments of symbolic music. "
        # prompt += "The prompts should be specific to a training example (a MIDI file or musical score), metadata for which is provided below. "
        # prompt += "They should be phrased as a user would enter them into a single-line text input: terse, not as full sentences. "
        # prompt += "Based on the metadata, infer the type of music, the composer, and the title of the work. The prompts should use any relevant "
        # prompt += "information from the metadata to guide the AI model in generating music segments, including artist name, work title, genre, and "
        # prompt += "stylistic / emotional descriptors. Not all metadata is accurate, but the majority of it is correct for a given file. "
        # prompt += "Be careful drawing inference from file names, as they may not be accurate. Use a majority vote approach when inferring the "
        # prompt += "identity of the work from file names, and always prefer metadata from within the file over file names when available. "
        # prompt += "Briefly describe your reasoning for each prompt. "
        # prompt += "\n\n"

        prompt = "In several steps, we will use metadata extracted from a symbolic music file to create a set of English text prompts "
        prompt += "that will be used to train a music-generating AI model. The metadata may include information about the composer, "
        prompt += "the title of the work, the genre, and stylistic or emotional descriptors. Not all metadata will be accurate. In particular, "
        prompt += 'a minority of file names under the optional "File names (aliases)" section may contain errors. Prefer to draw inferences '
        prompt += 'from metadata in other sections, or use a majority vote approach when inferring the identity of the work from "File names (aliases)". '
        prompt += "First, try to infer the type of music, the composer/artist/performer, the title of the work, and any relevant emotional or stylistic attributes from the metadata provided. "
        prompt += "You need not generate any prompts or ask follow-up questions at this stage. "
        prompt += "Here is the metadata extracted from the file:\n\n"

        return file_id, prompt + prompt_suf

def ai_file(file_id, prompt0):
    m0 = [
        {
            "role": "system",
            "content": "You are an AI assistant for machine learning research. Your top priority is to produce accurate, well-reasoned responses that satisfy user requests.",
        },
        {"role": "user", "content": prompt0},
    ]

    with lease_slot() as client:
        completion = client.chat.completions.create(
            model="gpt-3.5-turbo",
            extra_body={"mirostat": 2},
            messages=m0,
        )

        m1 = m0 + [completion.choices[0].message]
        m1 += [
            {
                "role": "user",
                "content": "".join(
                    [
                        "Now, generate 5 distinct prompts for the file. ",
                        'Each prompt should be terse, as a user would enter in a single-line prompt input text box, for example: "groovy upbeat bassline michael jackson". ',
                        "Write the prompts separated by newlines and without any preamble, prefix, or formatting. ",
                    ]
                ),
            }
        ]
        completion = client.chat.completions.create(
            model="gpt-3.5-turbo",
            extra_body={"mirostat": 2},
            messages=m1,
        )

    prompts = [
        re.sub(r"^\s*[0-9]+[.):]", "", line).replace("<|eot_id|>", "").strip()
        for line in completion.choices[0].message.content.splitlines()
    ]
    prompts = [p for p in prompts if p != ""]

    return file_id, prompts

script_start_timestamp = 1716949249

ai_q = queue.Queue(maxsize=100)
result_q = queue.Queue(maxsize=100)

def prep_thd_entry():
    try:
        with psycopg.connect(dbconnstr, row_factory=dict_row) as conn, conn.cursor() as cur:
            cur.execute(
                """
                SELECT id FROM files
                WHERE EXISTS (SELECT FROM extracted_clips WHERE file_id = files.id AND symbolic_length IS NOT NULL)
                AND NOT EXISTS (SELECT FROM file_descriptions WHERE file_id = files.id AND type = 'llama_prompt')
                AND dataset_name <> 'imslp'
                ORDER BY rand_order DESC
                """
            )
            for row in cur:
                try:
                    ai_q.put(prepare_file(row["id"]))
                except Exception as e:
                    print(f'prep_thd: {e}')
    finally:
        for _ in range(num_backends):
            ai_q.put(None)

def write_thd_entry():
    dead_thds = 0
    while True:
        try:
            x = result_q.get()
            if x is None:
                dead_thds += 1
                if dead_thds == num_backends:
                    break
            file_id, prompts = x
            print(f'{file_id}: {prompts}')
            with psycopg.connect(dbconnstr, row_factory=dict_row) as conn:
                conn.autocommit = True
                with conn.transaction(), conn.cursor() as cur:
                    cur.executemany(
                        "INSERT INTO file_descriptions ( file_id, script_start_timestamp, type, description ) VALUES (%s, %s, 'llama_prompt', %s)",
                        ((file_id, script_start_timestamp, p) for p in prompts),
                    )
        except Exception as e:
            print(f'write_thd: {e}')

def ai_thd_entry():
    while True:
        x = ai_q.get()
        if x is None:
            result_q.put(None)
            break
        try:
            result_q.put(ai_file(*x))
        except Exception as e:
            print(f'ai_thd: {e}')

prep_thd = threading.Thread(target=prep_thd_entry, daemon=True)
prep_thd.start()
write_thd = threading.Thread(target=write_thd_entry, daemon=True)
write_thd.start()
ai_thds = []
for _ in range(num_backends):
    ai_thds.append(threading.Thread(target=ai_thd_entry, daemon=True))
    ai_thds[-1].start()

prep_thd.join()
for t in ai_thds:
    t.join()
write_thd.join()

# Remove all control characters from the text

# bitmidi, geocities, lmd_full:
# file_descriptions -> file_name
# a file may have multiple file_name descriptions, some are wrong, but usually the majority are correct for a given file
# track_metadata -> name, lyrics, texts
#   texts -> kind in (marker, ...)
# some files have name and author info encoded in track names

# imslp:
# just file_descriptions -> imslp_text

# musescore:
# file -> metadata -> meta -> composer, workTitle
# file_descriptions -> musescore_text
# instrument -> metadata -> instrument, instrumentId