import sys
import re
from PyQt5 import QtCore, QtWidgets
import tkinter as tk
from tkinter import filedialog
import os
from collections import OrderedDict as odict
from io import SEEK_CUR
from struct import unpack

class MainWindow(QtWidgets.QMainWindow):
    def __init__(self):
        QtWidgets.QWidget.__init__(self)
        self.resize(360, 100)
        self.centralwidget = QtWidgets.QWidget(self)
        self.centralwidget.setObjectName("centralwidget")
        sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Fixed, QtWidgets.QSizePolicy.Fixed)
        sizePolicy.setHorizontalStretch(0)
        sizePolicy.setVerticalStretch(0)
        self.PrefixLabel = QtWidgets.QLabel(self.centralwidget)
        self.PrefixLabel.setGeometry(QtCore.QRect(10, 10, 41, 20))
        self.PrefixLabel.setObjectName("PrefixLabel")
        self.PrefixLineEdit = QtWidgets.QLineEdit(self.centralwidget)
        self.PrefixLineEdit.setGeometry(QtCore.QRect(50, 10, 181, 20))
        self.PrefixLineEdit.setObjectName("PrefixLineEdit")
        self.PrefixLineEdit.setText("EP")
        self.StartFromLabel = QtWidgets.QLabel(self.centralwidget)
        self.StartFromLabel.setGeometry(QtCore.QRect(250, 10, 51, 20))
        self.StartFromLabel.setObjectName("StartFromLabel")
        self.StartFromSpinBox = QtWidgets.QSpinBox(self.centralwidget)
        self.StartFromSpinBox.setGeometry(QtCore.QRect(310, 10, 41, 20))
        self.StartFromSpinBox.setObjectName("StartFromSpinBox")
        self.StartFromSpinBox.setMinimum(1)
        self.StartFromSpinBox.setMaximum(99)
        self.OpenChaptersButton = QtWidgets.QPushButton(self.centralwidget)
        self.OpenChaptersButton.setGeometry(QtCore.QRect(40, 50, 200, 23))
        self.OpenChaptersButton.setObjectName("OpenChaptersButton")
        self.CloseButton = QtWidgets.QPushButton(self.centralwidget)
        self.CloseButton.setGeometry(QtCore.QRect(245, 50, 75, 23))
        self.CloseButton.setObjectName("CloseButton")
        self.setCentralWidget(self.centralwidget)
        self.menubar = QtWidgets.QMenuBar(self)
        self.menubar.setGeometry(QtCore.QRect(0, 0, 361, 20))
        self.menubar.setObjectName("menubar")
        self.setMenuBar(self.menubar)
        self.statusbar = QtWidgets.QStatusBar(self)
        self.statusbar.setObjectName("statusbar")
        self.setStatusBar(self.statusbar)
        self.CloseButton.clicked.connect(self.CloseApp)
        self.OpenChaptersButton.clicked.connect(self.OpenChapters)
        self.retranslateUi(self)
        self.setFixedSize(self.size())

    def retranslateUi(self, MainWindow):
        self.setWindowTitle("Chapter Trimmer")
        self.CloseButton.setText("Close")
        self.PrefixLabel.setText("Prefix:")
        self.StartFromLabel.setText("Start from:")
        self.OpenChaptersButton.setText("Open and Write chapters")

    def OpenChapters(self):
        root = tk.Tk() # workaround
        root.withdraw() # workaround
        file_path = filedialog.askopenfilename(filetypes = (("Blu-ray playlists", "*.mpls"), ("All files", "*.*")))
        directory = filedialog.askdirectory()
        if len(file_path) > 0 and len(directory) > 0:
            mplsfile = open(file_path, 'rb')
            chapterdata = load(mplsfile)
            mplsfile.close()
            filesnum = len(chapterdata)
            filePrefix = self.PrefixLineEdit.text()
            fileNumber = int(self.StartFromSpinBox.text())
            for file in chapterdata:
                times = file['times']
                timeDict = {}
                titleDict = {}
                namecount = 1
                chaptersList = []
                index = 0
                if len(times) > 0:
                    startpos = 0
                    for time in times:
                        if startpos is 0:
                            startpos = time
                        nameFin = ''
                        timeFin = self.MsecToTimeString2(float((time - startpos) / 45000.))
                        namecount += 1
                        last = index is len(times) - 1
                        chaptersList.append('CHAPTER{index:02d}={time}\nCHAPTER{index:02d}NAME={name}'.format(index=index, time=timeFin, name=nameFin))
                        if last:
                            f = open('{dir}/{prefix}{number:02d}.txt'.format(dir=directory, prefix=filePrefix, number=fileNumber), 'w')
                            chaptersList.pop()
                            f.write('\n'.join(chaptersList))
                            f.close()
                            fileNumber += 1
                            index = 0
                            chaptersList.clear()
                        index += 1

    def TimeStringToMsec(self, string):
        matches = re.findall(r'(\d+):(\d+):(\d+)\.(\d+)', string, re.I)
        hours, mins, secs, msecs = matches[0]
        time = (int(msecs) + (int(secs) * 1000) + (int(mins) * 60 * 1000) + (int(hours) * 60 * 60 * 1000))
        return time

    def MsecToTimeString2(self, sec):
        secs = sec % 60
        mins = sec // 60 % 60
        hours = sec // 60 // 60 % 24
        time = r'{:02d}:{:02d}:{:06.03f}'.format(int(hours), int(mins), secs)
        return time

    def MsecToTimeString(self, msec):
        msecs = msec % 1000
        secs = msec // 1000 % 60
        mins = msec // 1000 // 60 % 60
        hours = msec // 1000 // 60 // 60 % 24
        time = r'{:02d}:{:02d}:{:02d}.{:03d}'.format(hours, mins, secs, msecs)
        return time

    def CloseApp(self):
        sys.exit(self)

def load(f, fix_overlap=True):
    def int_be(data):
        return {
            1: ord,
            2: lambda b: unpack('>H', b)[0],
            4: lambda b: unpack('>I', b)[0]
        }[len(data)](data)
    f.seek(8)
    addr_items, addr_marks = int_be(f.read(4)), int_be(f.read(4))
    f.seek(addr_items + 6)
    item_count = int_be(f.read(2))
    f.seek(2, SEEK_CUR)
    def read_item():
        block_size = int_be(f.read(2))
        item = odict()
        item['name'] = f.read(5).decode()
        f.seek(7, SEEK_CUR)
        item['times'] = [int_be(f.read(4)), int_be(f.read(4))]
        f.seek(block_size-20, SEEK_CUR)
        return item
    items = [read_item() for _ in range(item_count)]
    f.seek(addr_marks + 4)
    mark_count = int_be(f.read(2))
    def read_mark():
        f.seek(2, SEEK_CUR)
        index = int_be(f.read(2))
        time = int_be(f.read(4))
        f.seek(6, SEEK_CUR)
        return index, time
    for _ in range(mark_count):
        index, time = read_mark()
        if time > items[index]['times'][-2]:
            items[index]['times'].insert(-1, time)
    if fix_overlap:
        b = None
        for item in items:
            a, b = b, item['times']
        if a and b[0] < a[-1] < b[-1]:
            a[-1] = b[0]
        if len(b) > 1 and b[-1] - b[-2] < 90090:
            b.pop()
    return items

if __name__ == "__main__":
    app = QtWidgets.QApplication(sys.argv)
    w = MainWindow()
    w.show()
    sys.exit(app.exec_())