from fastapi import FastAPI, HTTPException, File, UploadFile, Depends, Security
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.security import APIKeyHeader
from ultralytics import YOLO
import cv2
import numpy as np
import json
import uvicorn
from kafka import KafkaProducer, KafkaConsumer
from redis import Redis
import uuid
import os
from datetime import datetime, timedelta, timezone
import threading
import torch
import string

torch.cuda.set_device(1) 

app = FastAPI()
pose_app = FastAPI()
app.mount("/pose", pose_app)

# CORS配置
ALLOWED_ORIGINS = 'https://beta.obscura.work'

# 只为主应用添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOWED_ORIGINS,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 配置
MODEL_PATH = "/home/zydi/models/yolov8x-pose.pt"  # 请替换为您的模型路径
KAFKA_BROKER = "222.186.10.253:9092"
KAFKA_TOPIC = "pose"  # 指定Kafka topic
KAFKA_GROUP_ID = "pose_group"  # 指定消费者组ID

REDIS_HOST = "222.186.10.253"
REDIS_PORT = 6379
REDIS_PASSWORD = "Obscura@2024" 
REDIS_DB = 3 
REDIS_API_DB = 12
REDIS_API_USAGE_DB = 13
UPLOAD_DIR = "/www/wwwroot/beta.obscura.work/upload_files/upload"
RESULT_DIR = "/www/wwwroot/beta.obscura.work/upload_files/result"
MAX_FILE_AGE = timedelta(hours=1)

# 确保目录存在
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(RESULT_DIR, exist_ok=True)

# 初始化 Kafka
producer = KafkaProducer(bootstrap_servers=[KAFKA_BROKER])
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=REDIS_DB
)

redis_api_client = Redis(
    host=REDIS_HOST,
    port=REDIS_PORT,
    password=REDIS_PASSWORD,
    db=REDIS_API_DB
)
redis_api_usage_client = Redis(
    host=REDIS_HOST,
    port=REDIS_PORT,
    password=REDIS_PASSWORD,
    db=REDIS_API_USAGE_DB
)

# 定义API密钥头部
api_key_header = APIKeyHeader(name="Authorization", auto_error=False)

# 定义 base62 字符集
BASE62 = string.digits + string.ascii_letters

# 验证API密钥的函数
async def get_api_key(api_key: str = Security(api_key_header)):
    if api_key and api_key.startswith("Bearer "):
        key = api_key.split(" ")[1]
        if key.startswith("obs-"):
            return key
    raise HTTPException(
        status_code=401,
        detail="无效的API密钥",
        headers={"WWW-Authenticate": "Bearer"},
    )

def calculate_tokens(file_path: str, file_type: str) -> int:
    base_tokens = 0
    
    try:
        file_size = os.path.getsize(file_path)  # 获取文件大小（字节）
        
        # 基础token：每MB文件大小消耗10个token
        base_tokens = int((file_size / (1024 * 1024)) * 10)
        
        if file_type == "image":
            img = cv2.imread(file_path)
            if img is None:
                raise ValueError("无法读取图片文件")
            height, width = img.shape[:2]
            pixel_count = height * width
            
            # 图片token：每100个像素额外消耗5个token
            image_tokens = int((pixel_count / 10000) * 5)
            
            base_tokens += image_tokens
        
        elif file_type == "video":
            cap = cv2.VideoCapture(file_path)
            if not cap.isOpened():
                raise ValueError("无法打开视频文件")
            fps = cap.get(cv2.CAP_PROP_FPS)
            frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            cap.release()
            
            pixel_count = width * height * frame_count
            duration = frame_count / fps  # 视频时长（秒）
            
            # 视频token：每100万像素每秒额外消耗1个token
            video_tokens = int((pixel_count / 10000) * (duration / 60))
            
            base_tokens += video_tokens
        
        return max(1, base_tokens)  # 确保至少返回1个token
    except Exception as e:
        print(f"计算token时出错: {str(e)}")
        return 1  # 出错时返回默认值1


async def verify_api_key(api_key: str = Depends(get_api_key)):
    redis_key = f"api_key:{api_key}"
    
    api_key_info = redis_api_client.hgetall(redis_key)
    
    if not api_key_info:
        raise HTTPException(status_code=401, detail="无效的API密钥")
    
    api_key_info = {k.decode('utf-8'): v.decode('utf-8') for k, v in api_key_info.items()}
    
    if api_key_info.get('is_active') != '1':
        raise HTTPException(status_code=401, detail="API密钥已停用")
    
    expires_at = datetime.fromisoformat(api_key_info.get('expires_at'))
    if datetime.now(timezone.utc) > expires_at:
        raise HTTPException(status_code=401, detail="API密钥已过期")
    
    usage_info = redis_api_usage_client.hgetall(redis_key)
    usage_info = {k.decode('utf-8'): v.decode('utf-8') for k, v in usage_info.items()}
    
    return {
        "api_key": api_key, 
        **api_key_info,
        **usage_info
    }


async def update_token_usage(api_key: str, new_tokens_used: int, model_name: str):
    redis_key = f"api_key:{api_key}"
    current_time = datetime.now(timezone.utc).isoformat()
    
    pipe = redis_api_usage_client.pipeline()
    
    # 更新总的token使用量
    pipe.hincrby(redis_key, "tokens_used", new_tokens_used)
    pipe.hset(redis_key, "last_used_at", current_time)
    
    # 更新特定模型的token使用量
    model_tokens_field = f"{model_name}_tokens_used"
    model_last_used_field = f"{model_name}_last_used_at"
    
    pipe.hincrby(redis_key, model_tokens_field, new_tokens_used)
    pipe.hset(redis_key, model_last_used_field, current_time)
    
    pipe.execute()

class PoseDetector:
    def __init__(self, model_path):
        self.model = YOLO(model_path).to('cuda:1')

    def detect(self, frame):
        results = self.model(frame, device='cuda:1')
        return results

    def format_results(self, results):
        formatted_results = []
        for r in results:
            boxes = r.boxes
            keypoints = r.keypoints
            for i in range(len(boxes)):
                box = boxes[i]
                kpts = keypoints[i]
                formatted_results.append({
                    "bbox": box.xyxy.tolist()[0],
                    "confidence": box.conf.item(),
                    "keypoints": kpts.xy.tolist()[0]
                })
        return formatted_results

    def draw_results(self, frame, results, original_shape):
        for r in results:
            annotated_frame = r.plot(img=frame)
            # 调整坐标以适应原始图像大小
            h, w = annotated_frame.shape[:2]
            scale_x, scale_y = original_shape[1] / w, original_shape[0] / h
            annotated_frame = cv2.resize(annotated_frame, (original_shape[1], original_shape[0]))
        return annotated_frame

detector = PoseDetector(MODEL_PATH)

def process_image(image_data, filename):
    try:
        img = cv2.imdecode(np.frombuffer(image_data, np.uint8), cv2.IMREAD_COLOR)
        original_shape = img.shape
        # Convert BGR to RGB
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        # Resize image to fit model requirements (640x640)
        img_resized = cv2.resize(img, (640, 640))
        
        # Normalize and reshape to BCHW format
        img_tensor = torch.from_numpy(img_resized).float().permute(2, 0, 1).unsqueeze(0) / 255.0
        img_tensor = img_tensor.to('cuda:1')

        results = detector.detect(img_tensor)
        
        # Format results for JSON
        json_results = detector.format_results(results)

        # Draw results on original image
        annotated_img = detector.draw_results(img_resized, results, original_shape)
        
        # Save annotated image
        annotated_filename = f"pose_{filename}"
        annotated_path = os.path.join(RESULT_DIR, annotated_filename)
        cv2.imwrite(annotated_path, cv2.cvtColor(annotated_img, cv2.COLOR_RGB2BGR))

        return json_results, annotated_filename
    except Exception as e:
        print(f"Error processing image: {str(e)}")
        return None, None

def process_video(video_data, filename):
    try:
        temp_video_path = os.path.join(UPLOAD_DIR, f"pose_{filename}")
        with open(temp_video_path, 'wb') as temp_video:
            temp_video.write(video_data)

        cap = cv2.VideoCapture(temp_video_path)
        frame_count = 0
        json_results = []
        
        # Get video properties
        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)
        
        # Create output video file
        annotated_filename = f"pose_{filename}"
        output_path = os.path.join(RESULT_DIR, annotated_filename)
        out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))

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

            # Process one frame per second
            if frame_count % fps == 0:
                # Convert BGR to RGB
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                
                # Resize frame to fit model requirements (640x640)
                frame_resized = cv2.resize(frame_rgb, (640, 640))
                
                # Normalize and reshape to BCHW format
                frame_tensor = torch.from_numpy(frame_resized).float().permute(2, 0, 1).unsqueeze(0) / 255.0
                frame_tensor = frame_tensor.to('cuda:1')

                results = detector.detect(frame_tensor)
                frame_json_results = detector.format_results(results)
                json_results.append({"frame": frame_count, "detections": frame_json_results})
                
                # Draw results on original frame
                annotated_frame = detector.draw_results(frame_resized, results, original_shape)
                # Convert RGB back to BGR for OpenCV
                annotated_frame = cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR)
            else:
                annotated_frame = frame

            out.write(annotated_frame)
            frame_count += 1

        cap.release()
        out.release()

        # Clean up temporary input video file
        os.remove(temp_video_path)

        return json_results, annotated_filename
    except Exception as e:
        print(f"Error processing video: {str(e)}")
        return None, None
    
@pose_app.post("/upload")
async def upload_file(file: UploadFile = File(...), api_key_info: dict = Depends(verify_api_key)):

    content = await file.read()
    file_extension = os.path.splitext(file.filename)[1].lower()
    new_filename = f"{uuid.uuid4()}{file_extension}"
    
    # Save the original file
    original_file_path = os.path.join(UPLOAD_DIR, new_filename)
    with open(original_file_path, "wb") as f:
        f.write(content)

    # 计算 token
    file_type = "image" if file_extension in ['.jpg', '.jpeg', '.png'] else "video"
    tokens_required = calculate_tokens(original_file_path, file_type)
    if tokens_required is None or tokens_required <= 0:
        raise HTTPException(status_code=500, detail="无法计算所需的token数量")
    
    # 检查并更新 token 使用量
    api_key = api_key_info['api_key']
    usage_key = f"api_key:{api_key}"
    total_tokens = int(redis_api_usage_client.hget(usage_key, "total_tokens") or 0)
    tokens_used = int(redis_api_usage_client.hget(usage_key, "tokens_used") or 0)

    if tokens_used + tokens_required > total_tokens:
        raise HTTPException(status_code=403, detail="Token 余额不足")

    # 更新 token 使用量
    model_name = "yolov8x-pose"
    await update_token_usage(api_key, tokens_required, model_name)


    # Send processing task to Kafka
    producer.send(KAFKA_TOPIC, json.dumps({
        "filename": new_filename,
        "file_type": "image" if file_extension in ['.jpg', '.jpeg', '.png'] else "video"
    }).encode('utf-8'))

    # Set initial status in Redis
    redis_key = f"pose_result:{new_filename}"
    redis_client.set(redis_key, json.dumps({"status": "queued"}))

    # 获取更新后的 token 使用情况
    updated_api_key_info = await verify_api_key(api_key)
    new_tokens_used = int(updated_api_key_info.get("tokens_used", 0))
    model_tokens_used = int(updated_api_key_info.get(f"{model_name}_tokens_used", 0))

    return JSONResponse(content={
        "message": "文件已上传并排队等待处理",
        "filename": new_filename,
        "tokens_used": tokens_required,
        "total_tokens_used": new_tokens_used,
        f"{model_name}_tokens_used": model_tokens_used,
        "tokens_remaining": total_tokens - new_tokens_used
    })


@pose_app.get("/result/{filename}", dependencies=[Depends(verify_api_key)])
async def get_pose_result(filename: str):
    redis_key = f"pose_result:{filename}"
    result = redis_client.get(redis_key)
    if result:
        result_data = json.loads(result)
        return JSONResponse(content=result_data)  # 直接返回整个结果，包括 status
    else:
        raise HTTPException(status_code=404, detail="Result not found")

@pose_app.get("/annotated/{filename}", dependencies=[Depends(verify_api_key)])
async def get_annotated_file(filename: str):
    redis_key = f"pose_result:{filename}"
    result = redis_client.get(redis_key)
    if result:
        result_data = json.loads(result)
        if result_data["status"] == "completed":
            annotated_filename = result_data["annotated_filename"]
            file_path = os.path.join(RESULT_DIR, annotated_filename)
            if os.path.exists(file_path):
                def iterfile():
                    with open(file_path, mode="rb") as file_like:
                        yield from file_like
                file_extension = os.path.splitext(annotated_filename)[1].lower()
                return StreamingResponse(iterfile(), media_type=f"image/{file_extension[1:]}" if file_extension in ['.jpg', '.jpeg', '.png'] else "video/mp4")
    
    raise HTTPException(status_code=404, detail="Annotated file not found")

def process_task():
    for message in consumer:
        task = message.value
        filename = task['filename']
        file_type = task['file_type']
        
        file_path = os.path.join(UPLOAD_DIR, filename)
        
        # Update status to "processing"
        redis_key = f"pose_result:{filename}"
        redis_client.set(redis_key, json.dumps({"status": "processing"}))
        
        try:
            if file_type == "image":
                with open(file_path, 'rb') as f:
                    content = f.read()
                json_results, annotated_filename = process_image(content, filename)
            else:
                with open(file_path, 'rb') as f:
                    content = f.read()
                json_results, annotated_filename = process_video(content, filename)
            
            if json_results and annotated_filename:
                redis_client.set(redis_key, json.dumps({
                    "json_results": json_results,
                    "status": "completed",
                    "annotated_filename": annotated_filename
                }))
            else:
                redis_client.set(redis_key, json.dumps({"status": "failed"}))
        except Exception as e:
            print(f"Error processing task: {str(e)}")
            redis_client.set(redis_key, json.dumps({"status": "failed", "error": str(e)}))

def listen_redis_changes():
    pubsub = redis_client.pubsub()
    pubsub.psubscribe('__keyspace@3__:pose_result:*')  # 监听所有pose_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 == 'set':
                value = redis_client.get(f"pose_result:{key}")
                if value:
                    result = json.loads(value)
                    print(f"Status update for {key}: {result['status']}")
                    
                    # 这里可以添加其他处理逻辑，比如发送通知等

if __name__ == "__main__":
    # 启动处理任务的线程
    threading.Thread(target=process_task, daemon=True).start()
    
    # 启动Redis监听线程
    threading.Thread(target=listen_redis_changes, daemon=True).start()
    
    uvicorn.run(app, host="0.0.0.0", port=7001)