import glob
import os
import sys
from hashlib import sha256
from pathlib import Path
import re
from mido import MidiFile, Message, tempo2bpm
from pprint import pprint
from dataclasses import dataclass, field
from dataclasses_json import dataclass_json, config as dj_config
from binascii import hexlify
from pebble import ProcessPool, ProcessExpired
import multiprocessing as mp
from concurrent.futures import TimeoutError
import json
import psycopg
from psycopg.errors import DeadlockDetected
from psycopg.types.json import Jsonb
from psycopg_pool import ConnectionPool
from typing import List
import traceback
import math
import signal
import random

dbpool = ConnectionPool('host=localhost dbname=composer_new_dataset_v3', min_size=2, max_size=2)

@dataclass_json
@dataclass
class Note:
    onBeat: float
    offBeat: float
    onVelocity: float
    offVelocity: float
    sourceTrack: int
    note: int

@dataclass
class AugmentEvent:
    evt: Message
    trackIdx: int
    absoluteTime: int
    beat: float

@dataclass_json
@dataclass
class Text:
    kind: str
    value: str

@dataclass_json
@dataclass
class Track:
    name: str
    text: str = field(default="", metadata=dj_config(exclude=lambda x: len(x) == 0))
    lyrics: str = field(default="", metadata=dj_config(exclude=lambda x: len(x) == 0))
    texts: List[Text] = field(default_factory=lambda: [], metadata=dj_config(exclude=lambda x: len(x) == 0))

OFF = 'off'

class Channel:
    def __init__(self, index):
        self.notes = [ OFF for i in range(128) ]
        self.note_tracks = [ 0 for i in range(128) ]
        self.vels = [ 0 for i in range(128) ]
        self.instruments = {}
        self.index = index
        self.program = None

    def push(self, ae):
        handler = getattr(self, ae.evt.type, None)
        if handler is None:
            return
        handler(ae)

    def note_on(self, ae):
        if ae.evt.velocity == 0:
            return self.note_off(ae)
        if self.notes[ae.evt.note] == OFF:
            self.notes[ae.evt.note] = ae.beat
            self.note_tracks[ae.evt.note] = ae.trackIdx
        self.vels[ae.evt.note] = ae.evt.velocity

    def note_off(self, ae):
        note = ae.evt.note
        if self.notes[note] == OFF:
            return
        if self.program not in self.instruments:
            self.instruments[self.program] = []
        self.instruments[self.program].append(
            Note(onBeat=self.notes[note],
                 offBeat=ae.beat,
                 onVelocity=self.vels[note],
                 offVelocity=ae.evt.velocity,
                 sourceTrack=self.note_tracks[note],
                 note=note)
        )
        self.notes[note] = OFF

    def program_change(self, ae):
        self.program = ae.evt.program

    def to_dict(self):
        return {
            'index': self.index,
            'isDrumIndex': self.index == 10,
            'instruments': [
                { 'program': k,
                  'notes': [ n.to_dict() for n in v ] }
                for k, v in sorted(self.instruments.items(), key=lambda x: -1 if x[0] is None else x[0])
            ]
        }

#MetaMessage('track_name', name='LuvTXG      ', time=0)
#MetaMessage('set_tempo', tempo=544464, time=0)
#MetaMessage('time_signature', numerator=4, denominator=4, clocks_per_click=24, notated_32nd_notes_per_beat=8, time=0)

class SkipException(Exception):
    pass

class ExistsException(Exception):
    pass

def process(path):
    try:
        mid = MidiFile(path)
    except EOFError:
        raise SkipException(f'File ended early: {path}')
    except Exception as e:
        if re.match('.*(list index|MThd|must be in range|decode key|EOFError|running status|MTrk)', str(e)):
            raise SkipException(f'Failed to parse (acceptable error): {path}')
        raise e

    if mid.type == 2:
        raise SkipException(f'Skipping type 2 file: {path}')

    tempo_evt = None
    time_signature_evt = None

    aug_evts = []

    tracks = []

    for i, track in enumerate(mid.tracks):
        time = 0
        track_out = Track(track.name.replace('\x00', ''))
        tracks.append(track_out)
        for evt in track:
            time += evt.time
            aug_evts.append(AugmentEvent(evt, i, time, time / mid.ticks_per_beat))
            if evt.type == 'text':
                if track_out.text != '':
                    track_out.text += evt.text.replace('\x00', '')
                else:
                    alltexts = [ t.text for t in track_out.texts if t.kind == 'text' and len(t.text) > 0 ]
                    alltexts.append(evt.text.replace('\x00', ''))
                    if len(alltexts) > 20: # probably lyrics
                        track_out.text = ''.join(alltexts)
                        track_out.texts = [ t for t in track_out.texts if t.kind != 'text' ]
            elif evt.type == 'lyrics':
                track_out.lyrics += evt.text.replace('\x00', '')
            elif str(evt.type) in ['copyright', 'marker', 'cue_marker']:
                track_out.texts.append(Text(evt.type, evt.text.replace('\x00', '')))
            elif str(evt.type) in ['instrument_name', 'device_name']:
                track_out.texts.append(Text(evt.type, evt.name.replace('\x00', '')))

    aug_evts = list(sorted(aug_evts, key=lambda e: e.absoluteTime))

    channels = {}

    for ae in aug_evts:
        if ae.evt.type == 'set_tempo':
            #if tempo_evt is not None and tempo_evt.tempo != ae.evt.tempo:
            #    raise Exception(f'Tempo change - skipping')
            if tempo_evt is None:
                tempo_evt = ae.evt
        elif ae.evt.type == 'time_signature':
            if time_signature_evt is not None and not (time_signature_evt.numerator == ae.evt.numerator and time_signature_evt.denominator == ae.evt.denominator):
                raise SkipException('Time signature change - skipping')
            time_signature_evt = ae.evt
        if ae.evt.is_meta or ae.evt.type == 'sysex':
            continue
        chan = ae.evt.channel
        if chan not in channels:
            channels[chan] = Channel(chan)
        channels[chan].push(ae)

    live_channels = [ c for i, c in sorted(channels.items(), key=lambda x: x[0]) if len(c.instruments) > 0 ]
    return {
        'tempo': tempo2bpm(tempo_evt.tempo) if tempo_evt is not None else 120,
        'time_signature': f'{time_signature_evt.numerator}/{time_signature_evt.denominator}' if time_signature_evt is not None else '4/4',
        'tracks': [ t.to_dict() for t in tracks ],
        'channels': [ c.to_dict() for c in live_channels ]
    }

def hash_notes(first_beat, ns):
    def quantize(v):
        return int(v * 192) // 192
    ns = sorted(ns, key=lambda n: (n["onBeat"], n["offBeat"], n["note"], n["onVelocity"]))
    preimage = []
    for n in ns:
        preimage.extend([
            str(quantize(n["onBeat"] - first_beat)),
            str(quantize(n["offBeat"] - first_beat)),
            str(int(n["note"])),
            str(int(n["onVelocity"]))
        ])
    preimage = "\n".join(preimage).encode('utf-8')
    h = sha256(preimage).digest()
    return int.from_bytes(h[:8], "big", signed=True)


def process_insert_txn(res, cur, path, path_hash):
    cur.execute("insert into files (tempo, time_signature, path, hash, dataset_name) values (%s, %s, %s, %s, %s) on conflict (hash) do nothing returning id",
                (res['tempo'], res['time_signature'], str(path), path_hash, str(path).split('/')[1]))

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

    file_id = rows[0][0]

    track_ids = []
    for i, track in enumerate(res['tracks']):
        cur.execute("insert into tracks (file_id, track_index, name, text, lyrics, texts) values (%s, %s, %s, %s, %s, %s)",
                    (file_id,
                     i,
                     track['name'],
                     track.get('text', None),
                     track.get('lyrics', None),
                     Jsonb(track.get('texts', []))))
        track_ids.append(cur.fetchone()[0])
    for channel in res['channels']:
        for instrument in channel['instruments']:
            notes = instrument["notes"]
            if len(notes) == 0:
                continue

            cur.execute("insert into instruments (file_id, channel_index, program) values (%s, %s, %s) returning id",
                        (file_id, channel["index"], instrument["program"]))
            instrument_id = cur.fetchone()[0]

            track_idxs = set(note["sourceTrack"] for note in notes)

            for track_idx in track_idxs:
                track_id = track_ids[track_idx]
                track_notes = [n for n in notes if n["sourceTrack"] == track_idx]
                first_beat = int(4 * (track_notes[0]["onBeat"] // 4))
                last_beat = int(math.ceil(max(note["offBeat"] for note in track_notes)))
                clip_length = 20
                clip_stride = 16

                def add_clip(start, end):
                    clip_notes = [
                        n for n in track_notes if (
                            n["onBeat"] >= start and
                            n["offBeat"] <= end
                        )
                    ]
                    if len(clip_notes) == 0:
                        return
                    h = hash_notes(first_beat, clip_notes)
                    cur.execute("insert into clips (hash, file_id, instrument_id, track_id, \"start\", \"length\") values (%s, %s, %s, %s, %s, %s) on conflict (instrument_id, track_id, start, length) do nothing",
                                (h, file_id, instrument_id, track_id, start, end - start))

                #for pos in range(first_beat, last_beat + clip_stride, clip_stride):
                #    add_clip(pos, pos + clip_length)

                #add_clip(first_beat, int(4 * math.ceil(max(n["offBeat"] for n in track_notes) / 4)))

                cur.executemany("insert into notes (instrument_id, track_id, pitch, \"start\", \"end\", velocity, off_velocity) values (%s, %s, %s, %s, %s, %s, %s)",
                                ((instrument_id, track_id,
                                  note["note"],
                                  note["onBeat"],
                                  note["offBeat"],
                                  note["onVelocity"],
                                  note["offVelocity"])
                                 for note in track_notes))

def process_insert(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, _ = path_hash_data

                    cur.execute('select id from files where hash = %s', (path_hash,))
                    rows = cur.fetchall()
                    if len(rows) == 0:
                        res = process(path)
                        if all(len(inst["notes"]) == 0 for channel in res["channels"] for inst in channel["instruments"]):
                            return
                        process_insert_txn(res, cur, path, path_hash)
                    cur.execute("insert into file_names (file_id, name, source) values (%s, %s, 'collect') on conflict (file_id, name) do nothing",
                                (rows[0][0], str(path)))
        except DeadlockDetected as e:
            if i < 99:
                print('retry on deadlock')
                continue
        except ExistsException:
            pass
        break
    else:
        raise RuntimeError('insert retries exceeded')

def main(glob, process_insert):
    def generate_paths():
        for p in sys.argv[1:]:
            for path in Path(p).rglob('*'):
                if not re.match(glob, str(path), flags=re.IGNORECASE):
                    continue
                if not os.path.isfile(path):
                    continue
                with open(path, 'rb') as f:
                    data = f.read()
                path_hash_hex = hexlify(sha256(data).digest()).decode('utf-8')
                yield path, path_hash_hex, data

    with ProcessPool(max_workers=12) as pool:
        future = pool.map(process_insert, generate_paths(), timeout=30)
        iterator = future.result()
        while True:
            try:
                result = next(iterator)
            except SkipException:
                pass
            except StopIteration:
                break
            except TimeoutError as error:
                print("function took longer than %d seconds" % error.args[1])
            except ProcessExpired as error:
                print("%s. Exit code: %d" % (error, error.exitcode))
            except Exception as error:
                print(error)
                print(error.traceback)  # Python's traceback of remote process

if __name__ == '__main__':
    mp.set_start_method('spawn')
    main('.*\.midi?$', process_insert)