import os
import torch

# 强制使用 CPU
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

import cv2
import torch
import numpy as np
from redis import Redis
import json
from kafka import KafkaConsumer
import threading
import redis 
from config import (
    KAFKA_BROKER, KAFKA_GROUP_ID_PREFIX, WORKER_CONFIGS,
    REDIS_HOST, REDIS_PORT, REDIS_PASSWORD, MAIN_REDIS_DB,
    UPLOAD_DIR, RESULT_DIR, DEEPFACE_MODEL_PATH
)
from deepface import DeepFace


KAFKA_TOPIC = WORKER_CONFIGS["compare"]["kafka_topic"]
KAFKA_GROUP_ID = f"compare_{KAFKA_GROUP_ID_PREFIX}"

# 初始化 Kafka
consumer = KafkaConsumer(
    KAFKA_TOPIC,
    bootstrap_servers=[KAFKA_BROKER],
    group_id=KAFKA_GROUP_ID,
    auto_offset_reset='earliest',
    enable_auto_commit=True,
    value_deserializer=lambda x: json.loads(x.decode('utf-8'))
)

# 初始化 Redis
redis_client = Redis(
    host=REDIS_HOST,
    port=REDIS_PORT,
    password=REDIS_PASSWORD,
    db=WORKER_CONFIGS["compare"]["redis_db"]
)

main_redis_client = Redis(
    host=REDIS_HOST,
    port=REDIS_PORT,
    password=REDIS_PASSWORD,
    db=MAIN_REDIS_DB
)


class FaceComparator:
    def __init__(self):
        print("初始化 DeepFace...")
        self.model_path = DEEPFACE_MODEL_PATH
        os.environ["DEEPFACE_HOME"] = self.model_path
        
    def detect(self, frame):
        """检测人脸并返回所有特征"""
        try:
            # 直接使用 represent 获取特征向量
            embeddings = DeepFace.represent(
                frame,
                model_name="Facenet512",
                detector_backend="retinaface",
                align=True,
                enforce_detection=False
            )
            
            if not isinstance(embeddings, list):
                embeddings = [embeddings]
            
            # 获取其他属性
            results = DeepFace.analyze(
                frame,
                actions=['age', 'gender', 'race', 'emotion'],
                detector_backend='retinaface',
                enforce_detection=False
            )
            
            if not isinstance(results, list):
                results = [results]
                
            # 合并结果
            for i in range(len(results)):
                if i < len(embeddings):
                    results[i]['embedding'] = embeddings[i]['embedding']  # 注意这里的变化
                    
            return results
        except Exception as e:
            print(f"人脸检测出错: {str(e)}")
            return []
        
    def format_results(self, faces):
        """格式化检测结果"""
        print("开始格式化检测结果...")
        try:
            formatted_results = []
            for face in faces:
                region = face.get('region', {})
                result = {
                    'bbox': [
                        region.get('x', 0),
                        region.get('y', 0),
                        region.get('w', 0),
                        region.get('h', 0)
                    ],
                    'gender': 1 if face.get('gender', '') == 'Man' else 0,
                    'age': float(face.get('age', 0)),
                    'det_score': 1.0,
                    'embedding': face.get('embedding', [])  # embedding 已经是列表格式
                }
                formatted_results.append(result)
            return formatted_results
        except Exception as e:
            print(f"格式化结果时出错: {str(e)}")
            return []
        
    def draw_results(self, frame, results):
        """在图像上绘制检测结果"""
        print("开始绘制检测结果...")
        try:
            annotated_frame = frame.copy()
            
            if not results:
                return annotated_frame
            
            for r in results:
                bbox = r['bbox']
                x, y, w, h = list(map(int, bbox))
                
                cv2.rectangle(annotated_frame,
                            (x, y),
                            (x + w, y + h),
                            (0, 255, 0), 2)
                
                gender_text = 'Male' if r['gender'] == 1 else 'Female'
                label = f"Age: {int(r['age'])} Gender: {gender_text}"
                cv2.putText(annotated_frame, 
                           label,
                           (x, y-10),
                           cv2.FONT_HERSHEY_SIMPLEX,
                           0.5,
                           (0, 255, 0),
                           2)
                           
            return annotated_frame
            
        except Exception as e:
            print(f"绘制结果时出错: {str(e)}")
            return frame.copy()

detector = FaceComparator()

def process_image(image_path):
    try:
        original_img = cv2.imread(image_path)
        if original_img is None:
            raise ValueError(f"无法读取图像文件: {image_path}")
            
        if original_img.size == 0:
            raise ValueError("图像数据为空")

        faces = detector.detect(original_img)
        
        if not faces:
            json_results = []
        else:
            json_results = detector.format_results(faces)
            if not json_results:  # 如果格式化失败，使用空列表
                json_results = []
        
        annotated_img = detector.draw_results(original_img, json_results)
        
        # 即使没有检测到人脸或格式化失败，也返回处理结果
        return json_results, annotated_img
        
    except Exception as e:
        return None, None

def process_video(video_path):
    try:
        cap = cv2.VideoCapture(video_path)
        frame_count = 0
        json_results = []
        
        fps = int(cap.get(cv2.CAP_PROP_FPS))
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        original_shape = (height, width)
        
        out = cv2.VideoWriter(video_path.replace(UPLOAD_DIR, RESULT_DIR), cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))

        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            if frame_count % fps == 0:
                preprocessed_frame = preprocess_frame(frame)
                
                results = detector.detect(preprocessed_frame)
                frame_json_results = detector.format_results(results, original_shape)
                json_results.append({"frame": frame_count, "detections": frame_json_results})
                
                annotated_frame = detector.draw_results(frame, frame_json_results)
                out.write(annotated_frame)

            frame_count += 1

        cap.release()
        out.release()

        return json_results
    except Exception as e:
        print(f"处理视频时出错: {str(e)}")
        return None

def preprocess_frame(frame):
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    frame_resized = cv2.resize(frame_rgb, (640, 640))
    frame_transposed = frame_resized.transpose((2, 0, 1))
    frame_contiguous = np.ascontiguousarray(frame_transposed)
    frame_tensor = torch.from_numpy(frame_contiguous).float()
    frame_normalized = frame_tensor / 255.0
    frame_batched = frame_normalized.unsqueeze(0)
    return frame_batched
    
def process_task():
    print("开始处理任务,等待Kafka消息...")
    for message in consumer:
        print(f"收到Kafka消息: topic={message.topic}, partition={message.partition}, offset={message.offset}")
        task = message.value
        task_id = task['task_id']
        filename = task['filename']
        file_type = task['file_type']
        
        print(f"解析任务信息: ID={task_id}, 文件名={filename}, 类型={file_type}")
        
        file_path = os.path.join(UPLOAD_DIR, filename)
        # 检查键类型并更新状态
        task_key = f"task:{task_id}"
        try:
            key_type = main_redis_client.type(task_key)
            if key_type != b'hash':
                main_redis_client.delete(task_key)
            main_redis_client.hset(task_key, "status", "processing")
            print(f"任务 {task_id} 状态更新为 'processing'")
        except redis.exceptions.ResponseError as e:
            print(f"更新任务 {task_id} 状态时出错: {str(e)}")
            continue  # 跳过这个任务，继续处理下一个
        
        try:
            if file_type == "image":
                print(f"开始处理图像: {filename}")
                json_results, annotated_img = process_image(file_path)
                if json_results and annotated_img is not None:
                    result_filename = f"compare_{filename}"
                    result_path = os.path.join(RESULT_DIR, result_filename)
                    cv2.imwrite(result_path, annotated_img)
                    
                    redis_client.hset(f"compare_result:{task_id}", 
                                    "result", json.dumps(json_results))
                    redis_client.hset(f"compare_result:{task_id}", 
                                    "result_file", result_filename)
                    
                    main_redis_client.hset(f"task:{task_id}", "status", "completed")
                    main_redis_client.hset(f"task:{task_id}", "result_type", "compare")
                    main_redis_client.hset(f"task:{task_id}", "result_key", f"compare_result:{task_id}")
                    
                    print(f"图像 {filename} 处理完成,结果已保存")
                else:
                    print(f"图像 {filename} 处理失败")
                    main_redis_client.hset(f"task:{task_id}", "status", "failed")
            else:  # video
                print(f"开始处理视频: {filename}")
                json_results = process_video(file_path)
                if json_results:
                    result_filename = f"compare_{filename}"
                    redis_client.hset(f"compare_result:{task_id}", 
                                    "result", json.dumps(json_results))
                    redis_client.hset(f"compare_result:{task_id}", 
                                    "result_file", result_filename)
                    
                    main_redis_client.hset(f"task:{task_id}", "status", "completed")
                    main_redis_client.hset(f"task:{task_id}", "result_type", "compare")
                    main_redis_client.hset(f"task:{task_id}", "result_key", f"compare_result:{task_id}")
                    
                    print(f"视频 {filename} 处理完成,结果已保存")
                else:
                    print(f"视频 {filename} 处理失败")
                    main_redis_client.hset(f"task:{task_id}", "status", "failed")
        except Exception as e:
            print(f"处理任务 {task_id} 时出错: {str(e)}")
            main_redis_client.hset(f"task:{task_id}", "status", "failed")
            main_redis_client.hset(f"task:{task_id}", "error", str(e))
        
        print(f"任务 {task_id} 处理完毕,等待下一个Kafka消息...")
def listen_redis_changes():
    pubsub = redis_client.pubsub()
    pubsub.psubscribe('__keyspace@3__:compare_result:*')  # 监听所有compare_result键的变化
    
    for message in pubsub.listen():
        if message['type'] == 'pmessage':
            key = message['channel'].decode('utf-8').split(':')[-1]
            operation = message['data'].decode('utf-8')
            
            if operation == 'hmset':
                value = redis_client.hgetall(f"compare_result:{key}")
                if value:
                    result = {k.decode(): v.decode() for k, v in value.items()}
                    print(f"Result update for task {key}: {result}")
                    

if __name__ == "__main__":
    print("compare处理程序启动...")
    # 启动处理任务的线程
    task_thread = threading.Thread(target=process_task, daemon=True)
    task_thread.start()
    print("任务处理线程已启动")
    
    # 启动Redis监听线程
    redis_thread = threading.Thread(target=listen_redis_changes, daemon=True)
    redis_thread.start()
    print("Redis监听线程已启动")
    
    print("主程序进入等待状态...")
    # 保持主线程运行
    task_thread.join()
    redis_thread.join()