import glob
import json
from hashlib import sha256
from pebble import ProcessPool, ProcessExpired
import multiprocessing as mp
import time

from psycopg.types.json import Jsonb
from psycopg.errors import DeadlockDetected
from collect import dbpool, ExistsException, main, SkipException

file_metadata_keys = ['meta', 'timeSignatureChanges', 'sectionStartEvents']
instrument_metadata_keys = ['instrument', 'instrumentId', 'drumPitchMapping']

start_timestamp = int(time.time())

def deep_equal(a, b):
    if isinstance(a, (str, int, float, bool)):
        return a == b
    if isinstance(a, list):
        return isinstance(b, list) and len(a) == len(b) and all(deep_equal(u, v) for u, v in zip(a, b))
    if isinstance(a, dict):
        return isinstance(b, dict) and len(a) == len(b) and all(k in b and deep_equal(v, b[k]) for k, v in a.items())
    return False

def process(cur, path, path_hash, data):
    obj = json.loads(data.decode('utf-8'))

    insert_notes = 0
    for track in obj["tracks"]:
        insert_notes += len(track["notes"])
    if insert_notes < 100:
        return
   
    time_sigs = obj['timeSignatureChanges']
    time_sig = time_sigs[0]['signature'] if len(time_sigs) > 0 else '4/4'

    file_meta = { k: obj[k] for k in file_metadata_keys if k in obj }
    cur.execute("insert into files (tempo, time_signature, path, hash, metadata, dataset_name) values (120, %s, %s, %s, %s, 'musescore') on conflict (hash) do nothing returning id", (time_sig, str(path), path_hash, Jsonb(file_meta)))

    rows = cur.fetchall()
    if len(rows) == 0:
        raise ExistsException('File already in')

    file_id = rows[0][0]

    cur.execute("insert into file_names (file_id, name, source) values (%s, %s, 'collect_musescore')", (file_id, str(path)))

    for text in obj["text"]:
        cur.execute("insert into file_descriptions (file_id, description, script_start_timestamp, type) values (%s, %s, %s, 'musescore_text') on conflict do nothing", (file_id, text, start_timestamp))

    for i, track in enumerate(obj["tracks"]):
        if len(track["notes"]) == 0:
            continue
        inst_meta = { k: track[k] for k in instrument_metadata_keys if k in track }
        cur.execute("insert into instruments (file_id, metadata) values (%s, %s) returning id", (file_id, Jsonb(inst_meta)))
        track["inst_id"] = cur.fetchone()[0]
        cur.execute("insert into tracks (file_id, track_index, name) values (%s, %s, %s) returning id", (file_id, i, track["name"]))
        track["track_id"] = cur.fetchone()[0]

        cur.execute("insert into track_descriptions (track_id, description, script_start_timestamp) values (%s, %s, %s)", (track["track_id"], track["name"], start_timestamp))

    for track in obj["tracks"]:
        cur.executemany("insert into notes (instrument_id, track_id, pitch, \"start\", \"end\", velocity, off_velocity) values (%s, %s, %s, %s, %s, %s, 127)",
                        ((track["inst_id"],
                          track["track_id"],
                          note["pitch"],
                          note["start"],
                          note["end"],
                          max(0, min(127, int(127 * note["velocity"]))))
                         for note in track["notes"]))

def process_insert_musescore(path_hash_data):
    for i in range(100):
        try:
            with dbpool.connection() as conn:
                conn.autocommit = True
                with conn.transaction(), conn.cursor() as cur:
                    path, path_hash, data = path_hash_data
                    process(cur, path, path_hash, data)
        except DeadlockDetected as e:
            if i < 99:
                print('retry on deadlock')
                continue
        except ExistsException:
            pass
        break
    else:
        raise RuntimeError('insert retries exceeded')

if __name__ == '__main__':
    mp.set_start_method('spawn')
    main('.*\.json$', process_insert_musescore)