from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS
import json
import os
import time
import requests
import base64
from io import BytesIO

from config import Config

app = Flask(__name__)
app.config.from_object(Config)
CORS(app)  # 启用CORS，允许跨域请求

# 加载发型数据
def load_hairstyles_from_file():
    """
    从hairId2url.txt文件加载发型数据
    解析文件中的发型ID和URL对应关系
    """
    hairstyles = []
    try:
        # 读取hairId2url.txt文件
        with open('../hairId2url.txt', 'r', encoding='utf-8') as f:
            lines = f.readlines()
        
        # 解析文件内容，提取发型ID和URL
        for line in lines:
            line = line.strip()
            
            # 查找包含 -> 分隔符的行，这些是发型ID和URL的映射
            if '->' in line and '.jpg' in line:
                # 分割行获取文件名和URL
                parts = line.split('->')
                if len(parts) == 2:
                    filename = parts[0].strip()
                    url = parts[1].strip()
                    
                    # 提取发型ID（去掉.jpg后缀）
                    if filename.endswith('.jpg'):
                        hair_id = filename[:-4]
                        
                        # 生成发型名称
                        hairstyle_name = f'发型{hair_id[-4:]}'
                        
                        # 添加到发型列表
                        hairstyles.append({
                            'id': hair_id,
                            'name': hairstyle_name,
                            'image_url': url
                        })
    except Exception as e:
        print(f'加载发型数据失败: {str(e)}')
    
    return hairstyles

# 加载发型数据
hairstyles_data = load_hairstyles_from_file()

@app.route('/api/hairstyles', methods=['GET'])
def get_hairstyles():
    """
    获取发型列表
    返回所有发型数据，不再区分男女
    """
    try:
        return jsonify({
            'success': True,
            'data': {
                'hairstyles': hairstyles_data
            }
        })
    except Exception as e:
        return jsonify({
            'success': False,
            'message': f'获取发型列表失败: {str(e)}'
        }), 500

@app.route('/api/change-hair', methods=['POST'])
def change_hair():
    """
    更换发型
    接收用户上传的图片和发型ID，调用远程API处理
    """
    try:
        # 检查是否有文件上传
        if 'image' not in request.files:
            return jsonify({
                'success': False,
                'message': '请上传图片'
            }), 400
        
        # 获取图片文件和发型ID
        image_file = request.files['image']
        hairstyle_id = request.form.get('hairstyle_id')
        
        if not hairstyle_id:
            return jsonify({
                'success': False,
                'message': '请选择发型'
            }), 400
        
        # 验证文件类型
        if not allowed_file(image_file.filename):
            return jsonify({
                'success': False,
                'message': '不支持的文件类型，请上传图片文件'
            }), 400
        
        # 读取图片数据
        image_data = image_file.read()
        
        # 调用远程API处理图片
        # 注意：这里是模拟实现，实际项目中需要根据远程API的要求进行调整
        result_image = call_remote_api(image_data, hairstyle_id)
        
        # 将处理后的图片转换为base64格式返回
        result_base64 = base64.b64encode(result_image).decode('utf-8')
        result_data_url = f'data:image/jpeg;base64,{result_base64}'
        
        return jsonify({
            'success': True,
            'data': {
                'result_image': result_data_url
            }
        })
    except Exception as e:
        return jsonify({
            'success': False,
            'message': f'处理图片失败: {str(e)}'
        }), 500

def allowed_file(filename):
    """
    验证文件类型是否允许
    """
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']

def call_remote_api(image_data, hairstyle_id):
    """
    调用远程API处理图片
    根据test_hair.py的实现方式
    """
    try:
        # 将图片数据转换为base64格式
        encoded_string = base64.b64encode(image_data).decode('utf-8')
        img64_str = f"data:image/jpeg;base64,{encoded_string}"
        
        # 构建请求数据（参考test_hair.py的格式）
        data = {
            "hair_id": hairstyle_id,
            "task_id": f"task_{int(time.time())}",  # 生成唯一的任务ID
            "user_img_path": img64_str,
            "is_hr": "false",
            "output_format": "base64"
        }
        
        # 打印传入的参数（只打印部分信息，避免日志过长）
        print("=== 调用远程API参数 ===")
        print(f"hair_id: {hairstyle_id}")
        print(f"task_id: {data['task_id']}")
        print(f"is_hr: {data['is_hr']}")
        print(f"output_format: {data['output_format']}")
        print(f"user_img_path 长度: {len(img64_str)}")
        print(f"user_img_path 前缀: {img64_str[:50]}...")  # 只打印前缀
        print("====================")
        
        # 设置请求头（参考test_hair.py）
        headers = {'Content-Type': 'application/json'}
        
        # 发送请求到远程API
        print(f"发送请求到: {app.config['REMOTE_API_URL']}")
        response = requests.post(
            app.config['REMOTE_API_URL'],
            headers=headers,
            json=data,
            timeout=app.config['API_TIMEOUT']
        )
        
        # 打印响应状态
        print(f"响应状态码: {response.status_code}")
        print(f"响应内容长度: {len(response.text)}")
        # 只打印响应内容的前100个字符，避免base64数据过长
        response_text = response.text
        if len(response_text) > 100:
            print(f"响应内容: {response_text[:100]}...")
        else:
            print(f"响应内容: {response_text}")
        
        # 解析响应数据
        response_data = response.json()
        
        # 检查响应状态
        if response.status_code != 200:
            raise Exception(f"远程API返回错误: {response_data.get('msg', '未知错误')}")
        
        # 参考test_hair.py的处理方式，直接检查是否有'data'字段
        # 不再依赖'state'字段，因为远程API可能使用不同的成功标识
        result_base64 = response_data.get('data', '')
        
        # 检查是否有数据返回，即使有msg字段也不视为错误
        # 因为远程API可能在成功时同时返回msg: "success"和data字段
        if not result_base64:
            # 如果没有'data'字段，检查是否有其他错误信息
            error_msg = response_data.get('msg', '未知错误')
            raise Exception(f"远程API返回错误: {error_msg}")
        
        # 如果有数据返回，即使msg字段有值也视为成功
        # 打印成功信息
        success_msg = response_data.get('msg', '成功')
        print(f"远程API调用成功: {success_msg}")
        
        # 移除可能的前缀
        if result_base64.startswith("data:image"):
            result_base64 = result_base64.split(",", 1)[1]
        
        # 解码base64数据为图片
        result_image = base64.b64decode(result_base64)
        
        print("远程API调用成功，返回处理后的图片")
        return result_image
    except Exception as e:
        print(f'调用远程API失败: {str(e)}')
        # 如果远程API调用失败，返回原始图片
        return image_data

@app.route('/')
def index():
    """
    提供前端静态文件
    开发环境下使用，生产环境应该使用Nginx等静态文件服务器
    """
    return send_from_directory('../frontend', 'index.html')

@app.route('/<path:path>')
def send_static(path):
    """
    提供前端静态文件
    """
    return send_from_directory('../frontend', path)

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)