#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
上传人脸照片修复版 - 解决文件太大问题
自动压缩超过3MB的照片后再上传
"""

from requests import get, post
import os
from PIL import Image
import io

def compress_image_for_upload(image_path, max_size_mb=3):
    """
    压缩图片用于上传
    
    Args:
        image_path: 图片路径
        max_size_mb: 最大文件大小(MB)
    
    Returns:
        bytes: 压缩后的图片数据，如果失败返回None
    """
    try:
        # 检查文件大小
        original_size = os.path.getsize(image_path)
        original_size_mb = original_size / (1024 * 1024)
        
        print(f"原始文件大小: {original_size_mb:.2f}MB")
        
        # 如果文件已经小于目标大小，直接读取返回
        if original_size_mb <= max_size_mb:
            with open(image_path, 'rb') as f:
                return f.read()
        
        # 打开图片进行压缩
        with Image.open(image_path) as img:
            # 转换为RGB模式
            if img.mode != 'RGB':
                img = img.convert('RGB')
            
            width, height = img.size
            print(f"原始尺寸: {width}x{height}")
            
            # 如果图片很大，先调整尺寸
            if max(width, height) > 2000:
                scale = 2000 / max(width, height)
                new_width = int(width * scale)
                new_height = int(height * scale)
                img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
                print(f"调整尺寸至: {new_width}x{new_height}")
            
            # 尝试不同质量等级进行压缩
            for quality in [85, 75, 65, 55, 45, 35, 25]:
                output_buffer = io.BytesIO()
                img.save(output_buffer, format='JPEG', quality=quality, optimize=True)
                
                compressed_data = output_buffer.getvalue()
                compressed_size_mb = len(compressed_data) / (1024 * 1024)
                
                print(f"质量{quality}%: {compressed_size_mb:.2f}MB")
                
                if compressed_size_mb <= max_size_mb:
                    compression_ratio = (original_size - len(compressed_data)) / original_size * 100
                    print(f"✓ 压缩成功！压缩率: {compression_ratio:.1f}%")
                    return compressed_data
                
                output_buffer.close()
            
            print("✗ 无法压缩到目标大小")
            return None
            
    except Exception as e:
        print(f"压缩失败: {str(e)}")
        return None

def upload_face_photo(file_path, server_url="http://10.143.25.3:8090", max_size_mb=3):
    """
    上传单张人脸照片（带压缩）
    
    Args:
        file_path: 照片文件路径
        server_url: 服务器地址
        max_size_mb: 最大文件大小(MB)
    
    Returns:
        bool: 上传是否成功
        str: 结果信息
    """
    try:
        filename = os.path.basename(file_path)
        print(f"\n处理照片: {filename}")
        
        # 检查文件是否存在
        if not os.path.exists(file_path):
            return False, f"文件不存在: {file_path}"
        
        # 检查文件格式
        if not file_path.lower().endswith(('.jpg', '.jpeg', '.png')):
            return False, f"不支持的文件格式: {file_path}"
        
        # 压缩图片
        print("正在压缩图片...")
        compressed_data = compress_image_for_upload(file_path, max_size_mb)
        
        if compressed_data is None:
            return False, "图片压缩失败"
        
        # 准备上传数据
        files = {
            "file": (filename, compressed_data, "image/jpeg"),
        }
        
        # 获取认证token
        print("正在获取认证token...")
        param = {
            "username": "tangchao",
            "password": "123456a",
            "client_id": "client",
            "grant_type": "password",
            "client_secret": "123456",
        }
        
        try:
            token_response = get(url=f"{server_url}/oauth/token", params=param, timeout=10)
            token_data = token_response.json()
            token = token_data["access_token"]
        except Exception as e:
            return False, f"获取token失败: {e}"
        
        # 设置请求头
        headers = {
            "authorization": f"Bearer {token}",
        }
        
        # 上传文件
        print("正在上传文件...")
        try:
            response = post(
                url=f"{server_url}/school/stuface/upload",
                headers=headers,
                files=files,
                timeout=60  # 60秒超时
            )
            
            if response.status_code == 200:
                result = response.json()
                if result.get('success'):
                    return True, f"上传成功: {filename}"
                else:
                    error_msg = result.get('msg', '未知错误')
                    return False, f"上传失败: {error_msg}"
            else:
                return False, f"HTTP错误: {response.status_code}"
                
        except Exception as e:
            return False, f"上传请求失败: {e}"
            
    except Exception as e:
        return False, f"处理过程出错: {e}"

def batch_upload_faces(students_dir="./students/", max_size_mb=3):
    """
    批量上传学生人脸照片
    
    Args:
        students_dir: 学生照片目录路径
        max_size_mb: 最大文件大小(MB)
    """
    if not os.path.exists(students_dir):
        print(f"目录不存在: {students_dir}")
        return
    
    print("开始批量上传学生人脸照片...")
    print("=" * 50)
    
    success_count = 0
    failed_count = 0
    failed_files = []
    
    # 支持的图片格式
    valid_extensions = {'.jpg', '.jpeg', '.png'}
    
    # 获取所有图片文件
    image_files = []
    for root, dirs, files in os.walk(students_dir):
        for file in files:
            file_ext = os.path.splitext(file)[1].lower()
            if file_ext in valid_extensions:
                image_files.append(os.path.join(root, file))
    
    print(f"找到 {len(image_files)} 张照片")
    print("-" * 50)
    
    # 处理每张照片
    for i, file_path in enumerate(image_files, 1):
        print(f"\n[{i}/{len(image_files)}] 处理: {os.path.basename(file_path)}")
        
        success, message = upload_face_photo(file_path, max_size_mb=max_size_mb)
        
        if success:
            success_count += 1
            print(f"✓ {message}")
        else:
            failed_count += 1
            failed_files.append(os.path.basename(file_path))
            print(f"✗ {message}")
        
        # 短暂延迟，避免服务器压力过大
        import time
        time.sleep(0.5)
    
    # 打印统计结果
    print("\n" + "=" * 50)
    print("上传完成！")
    print(f"成功: {success_count}")
    print(f"失败: {failed_count}")
    print(f"总计: {len(image_files)}")
    
    if failed_files:
        print(f"\n失败的文件:")
        for failed_file in failed_files:
            print(f"  - {failed_file}")

def main():
    """
    主函数
    """
    import sys
    
    print("上传人脸照片修复版")
    print("=" * 30)
    print("自动压缩超过3MB的照片后再上传")
    print()
    
    if len(sys.argv) > 1:
        # 命令行模式
        path = sys.argv[1]
        if os.path.exists(path):
            if os.path.isfile(path):
                # 上传单张照片
                success, message = upload_face_photo(path)
                print(message)
            elif os.path.isdir(path):
                # 批量上传
                batch_upload_faces(path)
        else:
            print(f"路径不存在: {path}")
    else:
        # 交互模式
        print("选择操作:")
        print("1. 批量上传学生照片目录")
        print("2. 上传单张照片")
        print("3. 自定义目录上传")
        
        choice = input("\n请输入选择(1-3): ").strip()
        
        if choice == "1":
            batch_upload_faces()
        elif choice == "2":
            file_path = input("请输入照片文件路径: ").strip()
            if os.path.exists(file_path):
                success, message = upload_face_photo(file_path)
                print(message)
            else:
                print(f"文件不存在: {file_path}")
        elif choice == "3":
            dir_path = input("请输入照片目录路径: ").strip()
            if os.path.exists(dir_path):
                batch_upload_faces(dir_path)
            else:
                print(f"目录不存在: {dir_path}")
        else:
            print("无效选择")

if __name__ == "__main__":
    main()