# -*- coding: utf-8 -*-
import os
import random
import subprocess
from pathlib import Path
from typing import List, Tuple
from concurrent.futures import ThreadPoolExecutor
from pydub import AudioSegment
import tempfile
import time
from datetime import datetime
import re
import json

from auto_media_publisher.utils.logger import get_logger
from auto_media_publisher.config.conf_video_taijiniao import *
from auto_media_publisher.generators.media_generator_base import MediaGenerator


class VideoGeneratorTaijiniao(MediaGenerator):
    """
    参照 VideoGeneratorMashup 为 太极鸟青少年体能培训 机构定制的类。
    
    """
    required_keys = ['user_id', 'scheduled_time','proxy_full']
    def __init__(self, *args, **kwargs):
        if args:
            raise TypeError("不支持位置参数，请全部使用关键字参数")

        for key in self.required_keys:
            if key not in kwargs:
                raise ValueError(f"缺少必要参数：{key}")
            setattr(self, key, kwargs[key])

        # 允许设置额外字段
        for key, value in kwargs.items():
            if key not in self.required_keys:
                setattr(self, key, value)
                
        self.logger = get_logger(self.__class__.__name__,"admin")
    

    def get_random_mp4_files(self,video_dir: str, num_files: int) -> List[str]:
        all_mp4s = []
        for root, dirs, files in os.walk(video_dir):
            for file in files:
                if file.lower().endswith((".mp4", ".mov")):
                    all_mp4s.append(os.path.join(root, file))
        return random.sample(all_mp4s, min(num_files, len(all_mp4s)))


    def get_video_files_and_segments(self,video_dir: str, num_segs=10) -> List[str]:
        """
        返回一个列表，列表长度 == num_segs
        当视频文件数量 > num_segs，取前 num_segs 个文件各一个 segment，
        当文件数量 <= num_segs，随机选文件，文件可以重复，凑够 num_segs。
        """
        all_mp4s = self.get_random_mp4_files(video_dir,num_segs)

        if len(all_mp4s) >= num_segs:
            # 直接取前num_segs个文件
            return all_mp4s[:num_segs]
        else:
            # 不够10个，随机挑选文件（允许重复），直到凑够num_segs
            chosen = []
            for _ in range(num_segs):
                chosen.append(random.choice(all_mp4s))
            return chosen

    def get_random_text_line(self,txt_dir: str) -> str:
        lines = []
        for root, dirs, files in os.walk(txt_dir):
            for file in files:
                if file.lower().endswith(".txt"):
                    with open(os.path.join(root, file), 'r', encoding='utf-8') as f:
                        lines.extend([line.strip() for line in f if line.strip()])
        return random.choice(lines) if lines else ""

    def get_all_audio_files(self,audio_dir: str) -> List[str]:
        audio_files = []
        for root, dirs, files in os.walk(audio_dir):
            for file in files:
                if file.lower().endswith((".mp3", ".wav")):
                    audio_files.append(os.path.join(root, file))
        return audio_files

    def get_random_speed(self) -> float:
        roll = random.random()
        if roll < 0.15:
            return 0.5
        elif roll < 0.30:
            return 1.2
        else:
            return 1.0

    def extract_segment_cmd(self,input_path: str, start: float, duration: float, speed: float, resolution: Tuple[int, int], flip: bool = False, zoom: bool = False) -> List[str]:
        filters = []
        # 缩放
        filters.append(f"scale={resolution[0]}:{resolution[1]}")
        # 镜头推进：zoompan 滤镜（保持帧率和缩放同步）
        if zoom:
            # 使用 zoompan 实现从中心推进的效果，zoom=1.0→1.2 线性过渡
            # 注意：需设置帧率和持续帧数（duration × fps），否则默认1fps
            zoom_filter = (
                f"zoompan="
                f"z='zoom+0.002':"
                f"x='iw/2-(iw/zoom/2)':"
                f"y='ih/2-(ih/zoom/2)':"
                f"d=1:fps=25:"
                f"s={resolution[0]}x{resolution[1]}"
            )

            filters.append(zoom_filter)

        # 变速
        if speed != 1.0:
            filters.append(f"setpts={1/speed}*PTS")
        # 水平翻转
        if flip:
            filters.append("hflip")

        vf = ",".join(filters)

        # "-q:v", "1" 质量最高 ，31 最差
        return [
            "ffmpeg", "-loglevel", "error", "-ss", str(start), "-t", str(duration), "-i", input_path,
            "-an", "-vf", vf,
            #"-f", "matroska", "-codec:v", "mpeg1video", "-q:v", "1", "pipe:1"
            "-f", "mpegts", "-c:v", "libx264", "-preset", "veryfast", "-crf", "18", "pipe:1"
        ]


        # 以下是原代码，使用ts格式，质量稍差
        '''
        return [
            "ffmpeg", "-loglevel", "error", "-ss", str(start), "-t", str(duration), "-i", input_path,
            "-an", "-vf", vf,
            "-f", "mpegts", "-codec:v", "mpeg1video", "-q:v", "3", "pipe:1"
        ]
        '''


    def process_video_segment(self,input_path: str, duration_range: Tuple[float, float], resolution: Tuple[int, int]) -> bytes:
        try:
            duration = random.uniform(*duration_range)
            speed = self.get_random_speed()
            start = random.uniform(0, max(1.0, self.get_video_duration(input_path) - duration))

            flip = random.random() > 0.5
            zoom = random.random() > 0.5

            cmd = self.extract_segment_cmd(input_path, start, duration, speed, resolution, flip=flip, zoom=zoom)
            return subprocess.check_output(cmd)
        except subprocess.CalledProcessError as e:
            self.logger.error(f"提取视频片段失败：{input_path}，错误：{e}")
            return b""


    def get_video_duration(self,path: str) -> float:
        cmd = [
            "ffprobe", "-v", "error", "-show_entries",
            "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", path
        ]
        output = subprocess.check_output(cmd).decode().strip()
        return float(output)

    def merge_segments_to_ts(self,segment_datas: List[bytes]) -> bytes:
        return b"".join(segment_datas)


    import re

    def wrap_text_cn_advanced(self,text, max_chars=20):
        # 常用中文标点符号，含中英文括号引号等
        punctuation = r'[，。！？；：“”‘’（）()、,.!?;:"\'\[\]]'
        
        # 先判断空格
        if ' ' in text:
            words = text.split()
            lines = []
            current_line = ""
            for w in words:
                if len(current_line) + len(w) + 1 <= max_chars:
                    current_line += (w + " ")
                else:
                    lines.append(current_line.strip())
                    current_line = w + " "
            lines.append(current_line.strip())
            return "\n".join(lines)
        else:
            # 没空格，先用正则把文本按标点符号切成小段，标点本身保留
            parts = re.split(f'({punctuation})', text)
            
            # parts 里奇数位是标点，偶数位是文本
            segments = []
            buffer = ""
            for i, part in enumerate(parts):
                if part == "":
                    continue
                if re.match(punctuation, part):
                    buffer += part  # 标点加到buffer末尾
                    segments.append(buffer)
                    buffer = ""
                else:
                    buffer += part
            if buffer:
                segments.append(buffer)
            
            # 现在segments里每个元素是一个小段（带标点）
            lines = []
            current_line = ""
            for seg in segments:
                if len(current_line) + len(seg) <= max_chars:
                    current_line += seg
                else:
                    if current_line:
                        lines.append(current_line)
                    # 如果单个segment比max_chars还长，则强制切断
                    while len(seg) > max_chars:
                        lines.append(seg[:max_chars])
                        seg = seg[max_chars:]
                    current_line = seg
            if current_line:
                lines.append(current_line)
            
            return "\n".join(lines)


    def compose_final_video(
        self,
        video_dir: str,
        text_line: str,
        audio_dir: str,
        logo_path: str,
        output_path: str,
        num_segments: int,
        segment_duration: Tuple[float, float],
        max_workers: int,
        resolution: Tuple[int, int],
        font_path: str,
        font_size: int,
        text_color: str
    ):
        #selected_medias = get_random_mp4_files(video_dir, num_segments)
        selected_medias = self.get_video_files_and_segments(video_dir,num_segments)

        audio_files = self.get_all_audio_files(audio_dir)
        audio_file = random.choice(audio_files) if audio_files else None
      

        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = [executor.submit(self.process_video_segment, str(v), segment_duration, resolution) for v in selected_medias]
            segments = [f.result() for f in futures if f.result()]

        ts_data = self.merge_segments_to_ts(segments)
        
        self.cleanup_old_temp_files()


        with tempfile.NamedTemporaryFile(delete=False, suffix=".ts") as temp_ts:
            temp_ts.write(ts_data)
            temp_ts.flush()

            #filters = [f"movie={logo_path} [logo]; [0:v][logo] overlay=30:30:format=auto [v1]; [v1] drawtext=fontfile={font_path}:text='{text_line}':fontsize={font_size}:fontcolor={text_color}:x=40:y=300 [vout]"]
            
            filters = [
                f"movie={logo_path}, scale=iw*0.3:-1, format=rgba, colorchannelmixer=aa=0.8 [logo];"
                f"[0:v][logo] overlay=100:100:format=auto [v1];"
                f"[v1] drawtext=fontfile={font_path}:text='{text_line}':fontsize={font_size}:fontcolor={text_color}:"
                f"box=1:boxcolor=black@0.5:boxborderw=5:x=(w-text_w)/2:y=300 [vout]"
            ]

            cmd = ["ffmpeg", "-loglevel", "error","-y", "-i", temp_ts.name]
            if audio_file:
                required_duration = num_segments * segment_duration[1]
                def get_valid_audio_offset(audio_path, max_attempts=5):
                    sound = AudioSegment.from_file(audio_path)
                    total_duration = len(sound) / 1000  # in seconds

                    for _ in range(max_attempts):
                        offset = random.uniform(0, total_duration - required_duration)
                        if (total_duration - offset) >= required_duration:
                            return round(offset, 2)
                    return 0.0  # fallback to beginning
                cmd += ["-ss", str(get_valid_audio_offset(audio_file)), "-i", audio_file, "-filter_complex", "; ".join(filters), "-map", "[vout]", "-map", "1:a"]
            else:
                cmd += ["-filter_complex", "; ".join(filters), "-map", "[vout]"]

            #cmd += ["-shortest", "-c:v", "libx264", "-preset", "slow", "-c:a", "aac", output_path]
            cmd += ["-shortest", "-c:v", "libx264", "-crf", "18", "-preset", "veryfast", "-c:a", "aac", output_path]
            # CPU: libx264
            # Nivida GPU : h264_nvenc 或 hevc_nvenc
            # AMD GPU: h264_amf 或 hevc_amf


            subprocess.run(cmd, check=True)
            time.sleep(2)  # 等待操作系统释放文件句柄
            try:
                os.unlink(temp_ts.name)
            except PermissionError:
                self.logger.warn(f"Warning: Could not delete temp file {temp_ts.name} because it is in use.")


    def extract_cover_frame(self,input_path, output_path, frame_number=10):
        cmd = [
            "ffmpeg",
            "-y",  # 自动覆盖输出文件
            "-loglevel", "error",
            "-i", input_path,
            "-vf", f"select=eq(n\\,{frame_number})",
            "-vframes", "1",
            "-q:v", "1",
            output_path
        ]
        subprocess.run(cmd, check=True)
        

    # 示例调用（你可以改为 CLI 参数解析或 GUI）
    def generate(self):
        try:
            tmp_datetime = datetime.strptime(self.scheduled_time, "%Y-%m-%d %H:%M:%S")
            scheduled_time_str  = tmp_datetime.strftime("%Y%m%d%H%M%S")
            scheduled_time_pub = tmp_datetime.strftime("%Y-%m-%d %H:%M")
            self.logger.info(f"正在生成 media：{self.user_id},{self.scheduled_time} .mp4")

            output_mp4_path =f"{TAIJINIAO_DIR_VIDEO_DST}/{self.user_id}/{scheduled_time_str}.mp4"
            output_png_path =f"{TAIJINIAO_DIR_VIDEO_DST}/{self.user_id}/{scheduled_time_str}.png"
            output_json_path = f"{TAIJINIAO_DIR_VIDEO_DST}/{self.user_id}/{scheduled_time_str}.json"

            text_line = self.wrap_text_cn_advanced(self.get_random_text_line(TAIJINIAO_DIR_TEXT), max_chars=450)  # 换行符用 \n 传给ffmpeg需要写成 \\n

            self.compose_final_video(
                video_dir=TAIJINIAO_DIR_VIDEO_SRC,
                text_line=text_line,
                audio_dir=TAIJINIAO_DIR_AUDIO,
                logo_path=TAIJINIAO_VIDEO_LOGO_PATH,
                output_path=output_mp4_path,
                num_segments=TAIJINIAO_VIDEO_SEGMENT_COUNT,
                segment_duration=TAIJINIAO_VIDEO_SEGMENT_DURATION,
                max_workers=4,
                resolution=TAIJINIAO_VIDEO_RESOLUTION,
                font_path=TAIJINIAO_VIDEO_FONT_PATH,
                font_size=TAIJINIAO_VIDEO_FONT_SIZE,
                text_color="yellow@0.8"
            )
            self.extract_cover_frame(output_mp4_path,output_png_path)
            video_json ={
                "user": self.user_id,
                "title": "太极鸟体适能",
                "description": text_line,
                "path_media": output_mp4_path,
                "tags": ["太极鸟","太极鸟体适能","基础运动技能"],
                "scheduled_time": scheduled_time_pub,
                "location": "上海城创金融科技国际产业园",
                "category": "健身",
                "path_cover": output_png_path,
                "proxy": self.proxy_full
            }
            # 保存
            with open(output_json_path, "w", encoding="utf-8") as f:
                json.dump(video_json, f, ensure_ascii=False, indent=2)
            
            return True
        except Exception as e:
            self.logger.error(f"taijiniao_media_generator error: {e}")
            return False