# 工作流
import json
import os.path
import platform
import subprocess
import sys
import traceback
from hashlib import md5
from typing import Optional, IO, Union

from datetime import datetime, timedelta
from PyQt5 import QtGui
from PyQt5.QtCore import Qt, QThread, pyqtSignal
from PyQt5.QtWidgets import QWidget, QLabel, QApplication, QFrame, QVBoxLayout, QPushButton, \
    QSizePolicy, QMessageBox
from danmaku_xml_helper import get_file_start, diff_danmaku_files, NoDanmakuException
from config import load_config, \
    DANMAKU_EXEC, DANMAKU_SPEED, DANMAKU_FONT_NAME, DANMAKU_FONT_SIZE, \
    VIDEO_CLIP_EACH_SEC, VIDEO_CLIP_OVERFLOW_SEC, VIDEO_RESOLUTION, VIDEO_OUTPUT_DIR, \
    FFMPEG_EXEC, HANDBRAKE_EXEC, HANDBRAKE_PRESET_FILE, HANDBRAKE_PRESET, HANDBRAKE_ENCOPT


class Job:
    DANMAKU_ENCODE = 0
    PURE_SPLIT = 1

    def __init__(self):
        super(Job, self).__init__()
        self.video = None
        self.type = self.PURE_SPLIT
        self.danmaku: list[str] = []
        self.subtitles: list[str] = []

    def __repr__(self):
        return "Job对象：Video[{}];Type[{}]".format(self.video, self.type)


class WorkLabel(QLabel):
    def __init__(self, *args, **kwargs):
        super(WorkLabel, self).__init__(*args, **kwargs)
        self.workVideo = None
        self.workDanmaku: list[str] = []
        self.running: bool = False
        self.finished: bool = False
        self.init_ui()

    def _update_text(self):
        if self.workVideo is None:
            text = "请拖入视频"
        else:
            text = "工作视频：{}".format(self.workVideo)
        if len(self.workDanmaku) == 0:
            self.setText("{}\n分割视频模式".format(text))
        else:
            self.setText("{}\n压制弹幕模式\n基准弹幕：{}".format(text, "\n附加弹幕：".join(self.workDanmaku)))
        self.adjustSize()

    def mouseDoubleClickEvent(self, a0: QtGui.QMouseEvent) -> None:
        if self.finished:
            self.finished = False
        elif self.workVideo is None:
            self.parent().labelDestroy.emit(self)
            return self.deleteLater()
        else:
            self.workVideo = None
            self.workDanmaku = []
        self.init_ui()

    def init_ui(self):
        font = QtGui.QFont()
        font.setPixelSize(12)
        self.setFont(font)
        self.setFrameShape(QFrame.Box)
        self.setStyleSheet("""QLabel {
            background-color: white;
        }""")
        self.setAlignment(Qt.AlignTop)
        self.setWordWrap(True)
        self.setMinimumHeight(64)
        self.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        self._update_text()
        self.setAcceptDrops(True)

    def dragEnterEvent(self, a0: QtGui.QDragEnterEvent) -> None:
        if a0.mimeData().hasUrls():
            a0.accept()
        else:
            a0.ignore()
        print("Enter Label", a0.mimeData().urls())

    def dropEvent(self, a0: QtGui.QDropEvent) -> None:
        urls = a0.mimeData().urls()
        for url in urls:
            # 判断是否为视频
            path = url.path().strip("/")
            if path.endswith(".mp4") or path.endswith(".flv"):
                self.set_video(path)
            elif path.endswith(".xml"):
                self.add_danmaku(path)
            else:
                print("Unknown File", path)
        self._update_text()

    def set_video(self, file):
        if not os.path.isfile(file):
            raise FileNotFoundError(file)
        if file == self.workVideo:
            return print("Same Video", file)
        self.workVideo = file
        self._update_text()

    def add_danmaku(self, file: str) -> None:
        if not os.path.isfile(file):
            raise FileNotFoundError(file)
        if file in self.workDanmaku:
            return print("Already Added File", file)
        self.workDanmaku.append(file)
        self._update_text()

    def have_job(self) -> bool:
        return self.workVideo is not None and not self.running and not self.finished

    def get_job(self) -> Job:
        job = Job()
        job.video = self.workVideo
        if len(self.workDanmaku) > 0:
            job.type = Job.DANMAKU_ENCODE
            job.danmaku = self.workDanmaku
        else:
            job.type = Job.PURE_SPLIT
            job.danmaku = []
        return job

    def start_running(self) -> None:
        self.running = True
        self.finished = False
        self.setStyleSheet("""QLabel {
            color: white;
            background-color: red;
        }""")

    def stop_running(self) -> None:
        self.running = False
        self.finished = True
        self.setStyleSheet("""QLabel {
            color: white;
            background-color: green;
        }""")


class HomePage(QWidget):
    showMessageBox = pyqtSignal(str, str)
    labelDestroy = pyqtSignal(WorkLabel)
    processCurTime = pyqtSignal(str)
    processSpeed = pyqtSignal(str)
    processCurTime2 = pyqtSignal(str)
    processSpeed2 = pyqtSignal(str)

    def __init__(self):
        super(HomePage, self).__init__()
        self.layout = None
        self.labels: list[WorkLabel] = []
        self.worker: WorkerThread
        self.btn_start: QPushButton
        self.showMessageBox.connect(self.on_show_message_box_info)
        self.labelDestroy.connect(self.on_label_destroy)
        self.processCurTime.connect(self.on_process_cur_time_change)
        self.processSpeed.connect(self.on_process_speed_change)
        self.process_cur_time = "-"
        self.process_speed = "-"
        self.cur_clip_duration = 0
        self.init_ui()

    def init_ui(self):
        layout = QVBoxLayout()
        self.layout = QVBoxLayout()
        layout.addLayout(self.layout)

        btn_start = QPushButton(self)
        btn_start.setText("开始")
        btn_start.clicked.connect(self.handle_do)
        btn_start.setGeometry(0, 560, 400, 40)
        layout.addWidget(btn_start, 1, Qt.AlignBottom)
        self.setLayout(layout)

        self.btn_start = btn_start
        self.resize(400, 600)
        self.setWindowTitle("录播工作流")
        self.setAcceptDrops(True)
        self.show()

    def add_label(self):
        label1 = WorkLabel(self)
        self.layout.addWidget(label1)
        self.labels.append(label1)
        return label1

    def handle_do(self):
        if len(self.labels) == 0:
            return self.on_show_message_box_warn("提示", "请添加任务")
        for label in self.labels:
            if not label.have_job():
                continue
            _thread = WorkerThread(self, label)
            _thread.start()
            self.worker = _thread
            _thread.started.connect(self.on_worker_start)
            _thread.finished.connect(self.on_worker_stop)
            break

    def dragEnterEvent(self, a0: QtGui.QDragEnterEvent) -> None:
        if a0.mimeData().hasUrls():
            a0.accept()
        else:
            a0.ignore()

    def dropEvent(self, a0: QtGui.QDropEvent) -> None:
        label1 = self.add_label()
        label1.dropEvent(a0)

    def on_show_message_box_info(self, title: str, content: str):
        QMessageBox.information(self, title, content, QMessageBox.Yes)

    def on_show_message_box_warn(self, title: str, content: str):
        QMessageBox.warning(self, title, content, QMessageBox.Yes)

    def on_label_destroy(self, label: WorkLabel):
        if label in self.labels:
            self.labels.remove(label)

    def on_worker_start(self):
        self.btn_start.setDisabled(True)
        self.btn_start.setText("正在处理")

    def on_process_cur_time_change(self, s: str) -> None:
        if self.process_cur_time == s:
            return
        self.process_cur_time = s
        self.update_btn_process_text()

    def on_process_speed_change(self, s: str) -> None:
        if self.process_speed == s:
            return
        self.process_speed = s
        self.update_btn_process_text()

    def update_btn_process_text(self):
        if self.process_cur_time != "" and self.process_speed != "":
            self.btn_start.setText("{}@{}".format(self.process_cur_time, self.process_speed))
        else:
            self.btn_start.setText("Working")

    def on_worker_stop(self):
        self.btn_start.setDisabled(False)
        self.btn_start.setText("开始")
        self.worker = None
        self.handle_do()


class WorkerThread(QThread):
    def __init__(self, app: HomePage, label: WorkLabel):
        super(WorkerThread, self).__init__()
        self.app = app
        self.label = label

    def get_video_real_duration(self, filename) -> float:
        ffmpeg_process = subprocess.Popen([
            FFMPEG_EXEC, "-hide_banner", "-progress", "-", "-loglevel", "error", "-i", filename, "-c", "copy", "-f",
            "null", "-"
        ], **subprocess_args(True)) 
        _duration_str = self.handle_ffmpeg_output(ffmpeg_process.stdout)
        ffmpeg_process.wait()
        return duration_str_to_float(_duration_str)

    def run(self) -> None:
        try:
            self.label.start_running()
            job = self.label.get_job()
            if job.type == Job.DANMAKU_ENCODE:
                self.run_danmaku_encode(job)
            elif job.type == Job.PURE_SPLIT:
                self.quick_split_video(job.video)
        except Exception as e:
            print(e)
            print(traceback.format_exc())
        finally:
            self.label.stop_running()

    def run_danmaku_encode(self, job: Job):
        base_danmaku = job.danmaku.pop(0)
        time_shift = 0
        try:
            base_start_ts = get_file_start(base_danmaku)
        except NoDanmakuException:
            return
        new_subtitle_name = danmaku_to_subtitle(base_danmaku, time_shift)
        job.subtitles.append(new_subtitle_name)
        for danmaku in job.danmaku:
            try:
                time_shift = diff_danmaku_files(base_danmaku, danmaku)
            except NoDanmakuException:
                continue
            new_subtitle_name = danmaku_to_subtitle(danmaku, time_shift)
            job.subtitles.append(new_subtitle_name)
        # 压制
        self.encode_video_with_subtitles(job.video, job.subtitles, base_start_ts)
        for _f in job.subtitles:
            if os.path.isfile(_f):
                os.remove(_f)

    def encode_video_with_subtitles(self, orig_filename: str, subtitles: list[str], base_ts: float):
        current_dt = create_dt = datetime.fromtimestamp(base_ts)
        process = get_encode_process_use_handbrake(orig_filename,
                                                    subtitles,
                                                    os.path.join(VIDEO_OUTPUT_DIR, "{}.mp4".format(current_dt)))
        self.handle_handbrake_output(process.stdout)
        process.wait()

    def quick_split_video(self, file):
        if not os.path.isfile(file):
            raise FileNotFoundError(file)
        file_name = os.path.split(file)[-1]
        _create_dt = os.path.splitext(file_name)[0]
        create_dt = datetime.strptime(_create_dt, "%Y%m%d_%H%M")
        duration = self.get_video_real_duration(file)
        current_sec = 0
        if duration < VIDEO_CLIP_EACH_SEC:
            print("[-]Less than each sec, skip")
            return False
        while current_sec < duration:
            if (current_sec + VIDEO_CLIP_OVERFLOW_SEC * 2) > duration:
                print("[-]Less than 2 overflow sec, skip")
                break
            current_dt = (create_dt + timedelta(seconds=current_sec)).strftime("%Y%m%d_%H%M_")
            print("CUR_DT", current_dt)
            print("BIAS_T", current_sec)
            split_process = subprocess.Popen([
                FFMPEG_EXEC, "-y", "-hide_banner", "-progress", "-", "-loglevel", "error",
                "-ss", str(current_sec),
                "-i", file, "-c:v", "copy", "-f", "mp4", "-c:a", "aac",
                "-t", str(VIDEO_CLIP_EACH_SEC + VIDEO_CLIP_OVERFLOW_SEC),
                "-fflags", "+genpts", "-shortest", "-movflags", "faststart",
                os.path.join(VIDEO_OUTPUT_DIR, "{}.mp4".format(current_dt))
            ], **subprocess_args(True))
            self.handle_ffmpeg_output(split_process.stdout)
            split_process.wait()
            current_sec += VIDEO_CLIP_EACH_SEC
        return True

    def handle_handbrake_output(self, stdout: Optional[IO[bytes]]):
        if stdout is None:
            print("[!]STDOUT is null")
            return
        json_body = ""
        json_start = False
        while True:
            line = stdout.readline()
            if line == b"":
                break
            if json_start:
                json_body += line.strip().decode("UTF-8")
                if line.startswith(b"}"):
                    json_start = False
                    status_payload = json.loads(json_body)
                    if status_payload["State"] == "WORKING":
                        self.app.processCurTime.emit("ETA: {Hours:02d}:{Minutes:02d}:{Seconds:02d}".format_map(status_payload["Working"]))
                        self.app.processSpeed.emit("{Rate:.2f}FPS".format_map(status_payload["Working"]))
                    elif status_payload["State"] == "WORKDONE":
                        break
                continue
            if line.startswith(b"Progress:"):
                json_start = True
                json_body = "{"
        self.app.processSpeed.emit("")
        self.app.processCurTime.emit("")

    def handle_ffmpeg_output(self, stdout: Optional[IO[bytes]]) -> str:
        out_time = "0:0:0.0"
        if stdout is None:
            print("[!]STDOUT is null")
            return out_time
        speed = "0"
        while True:
            line = stdout.readline()
            if line == b"":
                break
            if line.strip() == b"progress=end":
                # 处理完毕
                break
            if line.startswith(b"out_time="):
                out_time = line.replace(b"out_time=", b"").decode().strip()
                self.app.processCurTime.emit(out_time)
            if line.startswith(b"speed="):
                speed = line.replace(b"speed=", b"").decode().strip()
                self.app.processSpeed.emit(speed)
        self.app.processSpeed.emit("")
        self.app.processCurTime.emit("")
        print("[ ]Speed:", out_time, "@", speed)
        return out_time


def duration_str_to_float(duration_str) -> float:
    _duration = datetime.strptime(duration_str, "%H:%M:%S.%f") - datetime(1900, 1, 1)
    return _duration.total_seconds()


def base_ts_to_filename(start_ts: float, is_mp4=False) -> str:
    base_start = datetime.fromtimestamp(start_ts)
    if is_mp4:
        return base_start.strftime("%Y%m%d_%H%M.mp4")
    else:
        return base_start.strftime("%Y%m%d_%H%M.flv")


def danmaku_to_subtitle(file: Union[os.PathLike[str], str], time_shift: float):
    new_subtitle_name = md5(file.encode("utf-8")).hexdigest() + ".ass"
    process = subprocess.Popen((
        DANMAKU_EXEC, "--ignore-warnings",
        "-r", str(VIDEO_RESOLUTION), "-s", str(DANMAKU_SPEED), "-f", "5",
        "-S", str(DANMAKU_FONT_SIZE), "-N", str(DANMAKU_FONT_NAME), "--showmsgbox", "FALSE",
        "-O", "255", "-L", "1", "-D", "0",
        "-o", "ass", new_subtitle_name, "-i", file, "-t", str(time_shift)
    ), **subprocess_args(True))
    process.wait()
    return new_subtitle_name


def get_encode_process_use_handbrake(orig_filename: str, subtitles: list[str], new_filename: str):
    print("[+]Use HandBrakeCli")
    encode_process = subprocess.Popen([
        HANDBRAKE_EXEC, *_common_handbrake_setting(),
        "--preset-import-file", HANDBRAKE_PRESET_FILE, "--preset", HANDBRAKE_PRESET,
        "-i", orig_filename, "-x", HANDBRAKE_ENCOPT,
        "--ssa-file", ",".join(i for i in subtitles),
        "--ssa-burn", ",".join("%d" % (i+1) for i in range(len(subtitles))),
        "-o",
        new_filename
    ], **subprocess_args(True))
    return encode_process


# Create a set of arguments which make a ``subprocess.Popen`` (and
# variants) call work with or without Pyinstaller, ``--noconsole`` or
# not, on Windows and Linux. Typical use::
#
#   subprocess.call(['program_to_run', 'arg_1'], **subprocess_args())
#
# When calling ``check_output``::
#
#   subprocess.check_output(['program_to_run', 'arg_1'],
#                           **subprocess_args(False))
def subprocess_args(include_stdout=True):
    # The following is true only on Windows.
    if hasattr(subprocess, 'STARTUPINFO'):
        # On Windows, subprocess calls will pop up a command window by default
        # when run from Pyinstaller with the ``--noconsole`` option. Avoid this
        # distraction.
        si = subprocess.STARTUPINFO()
        si.dwFlags |= subprocess.STARTF_USESHOWWINDOW
        # Windows doesn't search the path by default. Pass it an environment so
        # it will.
        env = os.environ
    else:
        si = None
        env = None

    # ``subprocess.check_output`` doesn't allow specifying ``stdout``::
    #
    #   Traceback (most recent call last):
    #     File "test_subprocess.py", line 58, in <module>
    #       **subprocess_args(stdout=None))
    #     File "C:\Python27\lib\subprocess.py", line 567, in check_output
    #       raise ValueError('stdout argument not allowed, it will be overridden.')
    #   ValueError: stdout argument not allowed, it will be overridden.
    #
    # So, add it only if it's needed.
    if include_stdout:
        ret = {'stdout': subprocess.PIPE}
    else:
        ret = {}

    # On Windows, running this from the binary produced by Pyinstaller
    # with the ``--noconsole`` option requires redirecting everything
    # (stdin, stdout, stderr) to avoid an OSError exception
    # "[Error 6] the handle is invalid."
    ret.update({'stdin': subprocess.PIPE,
                'startupinfo': si,
                'env': env})
    return ret


def check_exec(name: Union[os.PathLike[str], str]) -> bool:
    if is_windows():
        check_process = subprocess.Popen([
            "where.exe", name
        ], **subprocess_args(True))
        check_process.wait()
        return len(check_process.stdout.readlines()) > 0
    elif is_linux():
        check_process = subprocess.Popen([
            "which", name
        ])
        check_process.wait()
        return check_process.returncode == 0
    else:
        return False


def is_windows() -> bool:
    return platform.system().lower() == "windows"


def is_linux() -> bool:
    return platform.system().lower() == "linux"


def check_all_prerequisite():
    if not check_exec(DANMAKU_EXEC):
        input("弹幕处理工具不存在")
        exit(1)
    if not check_exec(FFMPEG_EXEC):
        input("FFMPEG工具不存在")
        exit(1)
    if not check_exec(HANDBRAKE_EXEC):
        input("HANDBRAKE工具不存在")
        exit(1)


def _common_handbrake_setting():
    return (
        "--json",
        "--crop-mode", "none", "--no-comb-detect", "--no-bwdif", "--no-decomb", "--no-detelecine", "--no-hqdn3d",
        "--no-nlmeans", "--no-chroma-smooth", "--no-unsharp", "--no-lapsharp", "--no-deblock", "--align-av"
    )


def _common_ffmpeg_setting():
    return (
        "-y", "-hide_banner", "-progress", "-", "-loglevel", "error",
    )


def main():
    check_all_prerequisite()
    app = QApplication(sys.argv)
    page = HomePage()
    sys.exit(app.exec_())


if __name__ == '__main__':
    load_config()
    main()