import av
import torch
import threading, queue
from threading import current_thread
import numpy as np
import insightface
import cv2
import time
from insightface.app.common import Face
from gfpgan.utils import GFPGANer
from .typing import Frame # 你已有定义
from auto_media_publisher.config.conf_base import FACE_ENHANCER_MODEL_PATH,FACE_SWAPPER_MODEL_PATH
from insightface.app import FaceAnalysis

SIMILAR_FACE_DISTANCE: float = 0.45     # < 0.4 同一人 0.4 -0.6 大概是同一人 >0.6 不是同一个

def _select_device():
    if torch.cuda.is_available():
        return 'cuda:0'
    elif torch.backends.mps.is_available():
        return 'mps'
    else:
        return 'cpu'
    
NUM_WORKERS = 1

#face_swapper = insightface.model_zoo.get_model(FACE_SWAPPER_MODEL_PATH, providers=['CUDAExecutionProvider'])
#face_enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=_select_device())

frame_q = queue.Queue(maxsize=10)
processed_q = queue.Queue(maxsize=10)

THREAD_SEMAPHORE =threading.Semaphore(4)  # 不限制了。

def enhance_face(thread_face_enhancer,target_face, temp_frame:np.ndarray) -> np.ndarray:
    start_x, start_y, end_x, end_y = map(int, target_face['bbox'])
    padding_x = int((end_x - start_x) * 0.5)
    padding_y = int((end_y - start_y) * 0.5)
    start_x = max(0, start_x - padding_x)
    start_y = max(0, start_y - padding_y)
    end_x = max(0, end_x + padding_x)
    end_y = max(0, end_y + padding_y)
    temp_face = temp_frame[start_y:end_y, start_x:end_x]
    #print(temp_face.size)
    if temp_face.size:
        _, _, temp_face = thread_face_enhancer.enhance(
                temp_face,
                paste_back=True
        )
        temp_frame[start_y:end_y, start_x:end_x] = temp_face
    return temp_frame
    

def decode_video(video_path):
    container = av.open(video_path)
    for idx, frame in enumerate(container.decode(video=0)):
        img = frame.to_ndarray(format='rgb24')
        frame_q.put((idx, img))
    for _ in range(NUM_WORKERS):
        frame_q.put((None, None))

def process_frame_loop(new_face):
    start_time = time.time()
    thread_face_analyser = FaceAnalysis(name='buffalo_l', root='./', providers=['CUDAExecutionProvider'])
    thread_face_analyser.prepare(ctx_id=3, det_size=(640, 640))
    #thread_face_enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=_select_device())
    thread_face_enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=_select_device())
    thread_face_swapper = insightface.model_zoo.get_model(FACE_SWAPPER_MODEL_PATH, providers=['CUDAExecutionProvider'])

    while True:
        idx, frame = frame_q.get()
        if idx is None:
            processed_q.put((None, None))
            break
        try:
            many_faces = thread_face_analyser.get(frame)
            if many_faces:
                old_face = many_faces[0]
                frame = thread_face_swapper.get(frame, old_face, new_face, paste_back=True)

                for sinlge_face in many_faces:
                    frame_ = enhance_face(thread_face_enhancer,sinlge_face, frame)                

            processed_q.put((idx, frame_))
            print(f"{current_thread().name} has proceesed frame :{idx}")           

        except Exception as e:
            print(f"[ERROR] Frame {idx} failed: {e}")
    
    print(f"totol time is : {time.time()- start_time}")

def encode_video(output_path, fps, width, height):
    out = av.open(output_path, mode='w')

    #stream = out.add_stream('h264_nvenc', rate=fps)
    stream = out.add_stream('libx264', rate=fps)
    stream.width = width
    stream.height = height
    stream.pix_fmt = 'yuv420p'

    stream.options = {'crf': '18', 'preset': 'veryfast'}
    #stream.options = {'rc': 'constqp', 'qp': '18','preset': 'p7','profile': 'high'}

    next_index = 0
    buffer = {}
    end_count = 0

    while True:
        idx, frame = processed_q.get()
        if idx is None:
            end_count += 1
            if end_count >= NUM_WORKERS:
                break
            continue
        buffer[idx] = frame

        # 只要有当前要写入的帧，就输出
        while next_index in buffer:
            frame = buffer.pop(next_index)
            video_frame = av.VideoFrame.from_ndarray(frame, format='rgb24')
            for packet in stream.encode(video_frame):
                out.mux(packet)
            next_index += 1

    # 清空编码缓存
    for packet in stream.encode():
        out.mux(packet)
    out.close()
# 启动任务
video_path = './temp/buhui.mp4'
output_path = './temp/buhui-wangrong-enhancer.mp4'
new_face_image_path = './temp/wangrong.jpg'
face_analyser = FaceAnalysis(name='buffalo_l', root='./', providers=['CUDAExecutionProvider'])
face_analyser.prepare(ctx_id=0, det_size=(640, 640))

faces = face_analyser.get(cv2.imread(new_face_image_path))
if faces:
    new_face = sorted(faces, key=lambda x: x.bbox[0])[0]
        
    #threading.Thread(target=decode_video, args=(video_path,)).start()
    #threading.Thread(target=process_frame_loop, args=(new_face,)).start()
    #threading.Thread(target=encode_video, args=(output_path, 25,576,1080)).start()
    threads = []
        
    # 解码线程
    t_decode = threading.Thread(target=decode_video,args=(video_path,))
    threads.append(t_decode)
    t_decode.start()
    
    # 多处理线程
    for i in range(NUM_WORKERS):
        t = threading.Thread(target=process_frame_loop, args=(new_face,),name=f"process_worker_{i}")
        threads.append(t)
        t.start()
    
    # 编码线程
    t_encode = threading.Thread(target=encode_video,args=(output_path, 25,576,1080))
    threads.append(t_encode)
    t_encode.start()
    
    for t in threads:
        t.join()
    
    print("Processing complete.")