# -*- coding: utf-8 -*-
from concurrent.futures import ThreadPoolExecutor, FIRST_COMPLETED, wait
import subprocess
import ffmpeg
import re
from tqdm import tqdm

from ddans.native.hook import NHook
from ddans.native.log import NLog
from ddans.native.system import NSystem
from ddans.pbar import TQDM_FORMAT

ValidFormats = ["mp4", "mkv", 'avi', "wmv", "rm", "ts", "mov", "mpeg", "rmvb"]

DefaultVcodec = "libx264"

Quality720 = {
    0: 2500,
    1: 2000,
    2: 1500,
}

Quality1080 = {
    0: 5000,
    1: 4000,
    2: 3000,
}

QualityAudio = {
    0: 320,
    1: 192,
    2: 128,
}


class NFfmpeg:

    @staticmethod
    def get_video_duration(input_file):
        if not NSystem.exists(input_file):
            NLog.error(f"{input_file} 不存在")
            return 0
        try:
            probe = ffmpeg.probe(input_file)
            return float(probe['format']['duration'])
        except Exception as e:
            NLog.error(f"get video duration {e}")
            return 0

    @staticmethod
    def get_video_resolution(url):
        try:
            # 调用 ffprobe 获取视频信息
            probe = NFfmpeg.get_probe(url)
            if not probe:
                return None

            # 查找视频流
            video_streams = [
                stream for stream in probe['streams']
                if stream['codec_type'] == 'video'
            ]
            if not video_streams:
                return None  # 没有视频流

            # 取第一个视频流的分辨率（通常是最主要的）
            stream = video_streams[0]
            width = int(stream['width'])
            height = int(stream['height'])
            return width, height

        # except ffmpeg.Error as e:
        #     NLog.error(f"get_video_resolution  {e.stderr.decode('utf-8')}")
        #     return None
        except Exception as e:
            NLog.error(f"get_video_resolution {e}")
            return None

    @staticmethod
    def get_probe(input_file):
        try:
            return ffmpeg.probe(input_file)
        except Exception as e:
            NLog.error(f"get_probe {e}")
            return None

    @staticmethod
    def compress_video(input_file,
                       output_file,
                       vcodec=DefaultVcodec,
                       acodec='aac',
                       video_bitrate='2500k',
                       resolution='1280x720'):
        try:
            if not NSystem.exists(input_file):
                NLog.error(f"{input_file} 不存在")
                return

            _, ext = NSystem.split_name_ext(input_file)
            if ext.lower() not in ValidFormats:
                NLog.error(f"视频压缩不支持该格式: {ext}")
                return

            NSystem.ensure_dir(output_file)

            # 设置输出参数
            (ffmpeg.input(input_file).output(
                output_file,
                vcodec=vcodec,
                acodec=acodec,
                video_bitrate=video_bitrate,
                vf=f'scale={resolution}').overwrite_output().run(quiet=True))
            NLog.success(f'compress {input_file} to {output_file}')
        except ffmpeg.Error as e:
            NLog.error(f'compress: {e.stderr.decode()}')
        except Exception as e:
            NLog.error(f'compress: {e}')
            pass

    @staticmethod
    def compress_video_ex(input_file,
                          output_file,
                          vcodec=DefaultVcodec,
                          acodec='aac',
                          resolution=720,
                          quality=0,
                          title="",
                          delete_file=False):

        try:
            if not NSystem.exists(input_file):
                NLog.error(f"{input_file} 不存在")
                return False

            name, ext = NSystem.split_name_ext(input_file)
            if not NSystem.isvalid_ext(input_file, ValidFormats):
                NLog.error(f"视频压缩不支持${ext}格式")
                return False

            probe = NFfmpeg.get_probe(input_file)
            if probe is None or 'format' not in probe:
                return False

            total_size = float(probe['format'].get('size', 0))
            total_duration = float(probe['format'].get('duration', 0))
            # total_size = NSystem.getsize(input_file)
            if total_size is None or total_size <= 0:
                return False

            # total_duration = NFfmpeg.get_video_duration(input_file)
            if total_duration is None or total_duration <= 0:
                return False

            # 查找视频流
            video_streams = [
                stream for stream in probe['streams']
                if stream['codec_type'] == 'video'
            ]
            if not video_streams:
                return None  # 没有视频流

            # 取第一个视频流的分辨率（通常是最主要的）
            stream = video_streams[0]
            width = int(stream.get('width', 0))
            height = int(stream.get('height', 0))

            min_size = min(width, height)
            if not min_size:
                return None

            NSystem.ensure_dir(output_file)
            is_gpu = vcodec != DefaultVcodec

            max_resolution = resolution
            # 分辨率小于等720，但设置1080时，调整成720
            if min_size <= 720 and resolution > 720:
                max_resolution = 720

            if max_resolution <= 720:
                bitrate = Quality720.get(quality, Quality720.get(0))
            else:
                bitrate = Quality1080.get(quality, Quality1080.get(0))

            video_bitrate = f'{bitrate}k'
            minrate = video_bitrate
            maxrate_num = int(
                min(int(re.match(r"\d+", minrate).group()) * 2, 10000))
            bufsize_num = int(maxrate_num * 2)
            cmd = [
                'ffmpeg', '-y', '-i', input_file, '-map', '0:v', '-map', '0:a',
                '-vcodec', vcodec, '-acodec', acodec
            ]
            if is_gpu:
                cmd.extend(['-b:v', video_bitrate, '-preset', 'veryslow'])
            else:
                # '-b:v', video_bitrate,
                # '-minrate': f'{minrate}'
                # '-crf', '23'
                # '-preset', 'slow',
                cmd.extend(['-b:v', video_bitrate, '-preset', 'slow'])
            cmd.extend([
                '-maxrate', f'{maxrate_num}k', '-bufsize', f'{bufsize_num}k',
                '-movflags', 'faststart', '-vf', f'scale=-1:{max_resolution}',
                output_file
            ])
            process = subprocess.Popen(cmd,
                                       bufsize=0,
                                       universal_newlines=True,
                                       shell=False,
                                       stderr=subprocess.PIPE,
                                       encoding="utf-8",
                                       text=True)

            total_size_mb = round(total_size / (1024 * 1024), 2)  # 转换为MB
            desc = title if NHook.isvalid_str(title) else name
            with tqdm(total=total_size_mb,
                      unit='MB',
                      unit_scale=False,
                      ncols=80,
                      bar_format=TQDM_FORMAT,
                      leave=False,
                      colour="cyan",
                      desc=f'{desc}') as pbar:

                while process.poll() is None:
                    try:
                        line = process.stderr.readline().strip()
                        if not NHook.isvalid_str(line):
                            continue
                        # duration_res = re.search(
                        #     r'Duration: (?P<duration>\S+)', line)
                        # if duration_res is not None:
                        #     duration = duration_res.groupdict()\
                        #                ['duration']
                        #     duration = re.sub(r',', '', duration)
                        result = re.search(r'time=(?P<time>\S+)', line)
                        if result is not None:
                            elapsed_time = result.groupdict()['time']

                            currentTime = NFfmpeg.get_seconds(elapsed_time)
                            estimated_size = (currentTime /
                                              total_duration) * total_size
                            estimated_mb = round(
                                estimated_size / (1024 * 1024), 2)
                            pbar.update(estimated_mb - pbar.n)
                        # result = re.search(r'size=\s*(?P<size>\d+)ki?b',
                        #                    line.lower())
                    # if result is not None:
                    #     estimated_size = float(result.groupdict()['size'])
                    #     estimated_mb = round(estimated_size / 1024, 2)
                    #     pbar.update(estimated_mb - pbar.n)
                    except Exception:
                        pass
            process.wait()
            if process.returncode != 0:
                # NLog.error('compress_video: failed')
                return False
            else:
                if delete_file:
                    NSystem.remove_file(input_file)
                # NLog.success(f'compress_video {input_file} to {output_file}')
                return True
        except ffmpeg.Error as e:
            NLog.error(f'compress: {e.stderr.decode()}')
            return False
        except Exception as e:
            NLog.error(f'compress: {e}')
            return False

    # @staticmethod
    # def get_seconds(time):
    #     h = int(time[0:2])
    #     # print("时：" + str(h))
    #     m = int(time[3:5])
    #     # print("分：" + str(m))
    #     s = int(time[6:8])
    #     # print("秒：" + str(s))
    #     ms = int(time[9:12])
    #     # print("毫秒：" + str(ms))
    #     ts = (h * 60 * 60) + (m * 60) + s + (ms / 1000)
    #     return ts
    @staticmethod
    def get_seconds(time_str):
        if ':' not in time_str:
            return 0
        """将时间字符串转换为秒数"""
        h, m, s = map(float, time_str.split(':'))
        return h * 3600 + m * 60 + s

    @staticmethod
    def ffmpeg_command(*args):
        """ 运行 FFmpeg 并返回输出 """
        try:
            cmd = ("ffmpeg", ) + args
            result = NSystem.run(*cmd)
            return result.stdout + result.stderr  # 可能信息在 stderr
        except FileNotFoundError:
            return ""  # FFmpeg 未安装

    @staticmethod
    def has_nvidia_gpu():
        """ 判断是否有 NVIDIA 显卡 """
        try:
            result = NSystem.run("nvidia-smi")
            # 如果返回结果包含 "NVIDIA-SMI" 或有显卡信息，说明有 NVIDIA 显卡
            return 'NVIDIA-SMI' in result.stdout
        except FileNotFoundError:
            # 如果没有安装 nvidia-smi 或没有 NVIDIA 显卡，会抛出 FileNotFoundError
            return False

    @staticmethod
    def get_available_encoders():
        """ 运行 FFmpeg 检测可用的 GPU 编码器 """
        output = NFfmpeg.ffmpeg_command("-encoders")
        encoders = {
            "nvenc":
            NSystem.is_windows and NFfmpeg.has_nvidia_gpu()
            and "hevc_nvenc" in output,
            "qsv":
            NSystem.is_windows and "hevc_qsv" in output,
            "vaapi":
            False,  # "h264_vaapi" in output,
            "videotoolbox":
            NSystem.is_darwin and "videotoolbox" in output
        }
        return encoders

    @staticmethod
    def select_hevc_encoder():
        available_encoders = NFfmpeg.get_available_encoders()
        """ 根据可用性选择最合适的 GPU 编码器 """
        if available_encoders["nvenc"]:
            return "hevc_nvenc"  # NVIDIA
        elif available_encoders["qsv"]:
            return "hevc_qsv"  # Intel QSV
        elif available_encoders["vaapi"]:
            return "hevc_amf"  # AMD VAAPI
        elif available_encoders["videotoolbox"]:
            return "hevc_videotoolbox"  # macOS
        else:
            return DefaultVcodec  # 纯软件编码（兼容模式）

    @staticmethod
    def select_encoder():
        available_encoders = NFfmpeg.get_available_encoders()
        """ 根据可用性选择最合适的 GPU 编码器 """
        if available_encoders["nvenc"]:
            return "h264_nvenc"  # NVIDIA
        elif available_encoders["qsv"]:
            return "h264_qsv"  # Intel QSV
        elif available_encoders["vaapi"]:
            return "h264_amf"  # AMD VAAPI
        elif available_encoders["videotoolbox"]:
            return "h264_videotoolbox"  # macOS
        else:
            return DefaultVcodec  # 纯软件编码（兼容模式）

    @staticmethod
    def compress_videos(dir="",
                        max=2,
                        resolution=720,
                        quality=0,
                        delete_file=False,
                        gpu=False):
        root = NSystem.get_full_path(dir)
        # print("compress_videos: %s" % root)
        if not NSystem.exists(root):
            NLog.error("compress directory not found!")
            return

        files = NSystem.listdir(root)
        arrs = list(
            filter(
                lambda i: NSystem.isfile(i) and NSystem.isvalid_ext(
                    i, ValidFormats), files))

        if len(arrs) <= 0:
            NLog.error("compress none")
            return

        futures = []
        vcodec = NFfmpeg.select_encoder() if gpu else DefaultVcodec
        with ThreadPoolExecutor(max_workers=max) as executor:
            for i in range(len(arrs)):
                file = arrs[i]
                old_name, ext = NSystem.split_name_ext(file)
                output_file = NSystem.join(root, 'out', f'F-{old_name}.mp4')
                futures.append(
                    executor.submit(NFfmpeg.compress_video_ex,
                                    file,
                                    output_file,
                                    delete_file=delete_file,
                                    quality=quality,
                                    vcodec=vcodec,
                                    resolution=resolution,
                                    title=str(i + 1)))

            pbar = tqdm(total=len(futures),
                        desc="compress",
                        unit='',
                        ncols=80,
                        bar_format=TQDM_FORMAT,
                        colour="red")
            while futures:
                done, futures = wait(futures, return_when=FIRST_COMPLETED)
                for future in done:
                    try:
                        future.result()  # 检查异常，但不阻塞
                    except Exception:
                        # NLog.error(f'视频压缩异常: {e}')
                        pass
                    pbar.update(1)
            pbar.close()

        NLog.success("compress: end")

    @staticmethod
    def export_audio(input_file,
                     output_file,
                     acodec='libmp3lame',
                     quality=1,
                     title="",
                     delete_file=False):

        try:
            if not NSystem.exists(input_file):
                NLog.error(f"{input_file} 不存在")
                return False

            name, ext = NSystem.split_name_ext(input_file)
            if not NSystem.isvalid_ext(input_file, ValidFormats):
                NLog.error(f"不支持${ext}格式")
                return False

            probe = NFfmpeg.get_probe(input_file)
            if probe is None or 'format' not in probe:
                return False

            total_size = float(probe['format'].get('size', 0))
            total_duration = float(probe['format'].get('duration', 0))
            # total_size = NSystem.getsize(input_file)
            if total_size is None or total_size <= 0:
                return False

            # total_duration = NFfmpeg.get_video_duration(input_file)
            if total_duration is None or total_duration <= 0:
                return False

            # 查找音频流
            # audio_streams = [
            #     stream for stream in probe['streams']
            #     if stream['codec_type'] == 'audio'
            # ]
            # if not audio_streams:
            #     return None  # 没有视频流

            # # 取第一个视频流的分辨率（通常是最主要的）
            # stream = audio_streams[0]
            # audio_bitrate = int(stream.get('bit_rate', 0))
            bitrate = QualityAudio.get(quality, QualityAudio.get(1, 0))

            min_bitrate = min(QualityAudio.get(0), bitrate)
            if not min_bitrate:
                return None

            NSystem.ensure_dir(output_file)

            # 设置音频比特率等参数
            cmd = [
                'ffmpeg', '-y', '-i', input_file, '-vn', '-acodec', acodec,
                '-b:a', f'{min_bitrate}k', '-vn', output_file
            ]
            process = subprocess.Popen(cmd,
                                       bufsize=0,
                                       universal_newlines=True,
                                       shell=False,
                                       stderr=subprocess.PIPE,
                                       encoding="utf-8",
                                       text=True)

            total_size_mb = round(total_size / (1024 * 1024), 2)  # 转换为MB
            desc = title if NHook.isvalid_str(title) else name
            with tqdm(total=total_size_mb,
                      unit='MB',
                      unit_scale=False,
                      ncols=80,
                      bar_format=TQDM_FORMAT,
                      leave=False,
                      colour="cyan",
                      desc=f'{desc}') as pbar:

                while process.poll() is None:
                    try:
                        line = process.stderr.readline().strip()
                        if not NHook.isvalid_str(line):
                            continue
                        # duration_res = re.search(
                        #     r'Duration: (?P<duration>\S+)', line)
                        # if duration_res is not None:
                        #     duration = duration_res.groupdict()\
                        #                ['duration']
                        #     duration = re.sub(r',', '', duration)
                        result = re.search(r'time=(?P<time>\S+)', line)
                        if result is not None:
                            elapsed_time = result.groupdict()['time']

                            currentTime = NFfmpeg.get_seconds(elapsed_time)
                            estimated_size = (currentTime /
                                              total_duration) * total_size
                            estimated_mb = round(
                                estimated_size / (1024 * 1024), 2)
                            pbar.update(estimated_mb - pbar.n)
                    except Exception:
                        pass
            process.wait()
            if process.returncode != 0:
                # NLog.error('compress_video: failed')
                return False
            else:
                if delete_file:
                    NSystem.remove_file(input_file)
                # NLog.success(f'compress_video {input_file} to {output_file}')
                return True
        except ffmpeg.Error as e:
            NLog.error(f'export: {e.stderr.decode()}')
            return False
        except Exception as e:
            NLog.error(f'export: {e}')
            return False

    @staticmethod
    def export_audios(dir="", max=2, quality=1):
        root = NSystem.get_full_path(dir)
        if not NSystem.exists(root):
            NLog.error("export directory not found!")
            return

        files = NSystem.listdir(root)
        arrs = list(
            filter(
                lambda i: NSystem.isfile(i) and NSystem.isvalid_ext(
                    i, ValidFormats), files))

        if len(arrs) <= 0:
            NLog.error("export none")
            return

        futures = []
        with ThreadPoolExecutor(max_workers=max) as executor:
            for i in range(len(arrs)):
                file = arrs[i]
                old_name, ext = NSystem.split_name_ext(file)
                output_file = NSystem.join(root, 'out', f'F-{old_name}.mp3')
                futures.append(
                    executor.submit(NFfmpeg.export_audio,
                                    file,
                                    output_file,
                                    quality=quality,
                                    title=str(i + 1)))

            pbar = tqdm(total=len(futures),
                        desc="export audios",
                        unit='',
                        ncols=80,
                        bar_format=TQDM_FORMAT,
                        colour="red")
            while futures:
                done, futures = wait(futures, return_when=FIRST_COMPLETED)
                for future in done:
                    try:
                        future.result()  # 检查异常，但不阻塞
                    except Exception:
                        # NLog.error(f'视频压缩异常: {e}')
                        pass
                    pbar.update(1)
            pbar.close()

        NLog.success("export_audios: end")
