import io
import os
import requests
import base64
import json
from pathlib import Path
from PIL import Image

import jingrow
from jingrow.ai.utils.utils import update_status_progress, update_progress, send_notification, upload_image, get_jingrow_api_headers, get_jingrow_api_url

# 调用添加背景API
def call_jingrow_add_bg_api(image_urls, config=None, callback=None):
    """
    调用添加背景API，支持实时处理返回的图片
    
    Args:
        image_urls: 要处理的图片URL列表
        config: 添加背景的配置参数
        callback: 回调函数，用于处理每张返回的图片
    """
    # 获取API URL
    api_url = f"{get_jingrow_api_url()}/add_bg/batch"
    
    headers = get_jingrow_api_headers()
    if not headers:
        send_notification("Jingrow API 未设置，请在 AI设置 中设置 Jingrow Api Key 和 Jingrow Api Secret")
        return {"success": False, "message": "API密钥未设置"}
        
    # 准备请求数据
    request_data = {"urls": image_urls}
    if config:
        request_data["config"] = config
    
    try:
        # 调用API并获取流式响应
        response = requests.post(api_url, json=request_data, headers=headers, stream=True, timeout=180)
        
        if response.status_code != 200:
            send_notification(f"API请求失败 (HTTP {response.status_code})")
            return {"success": False, "message": f"HTTP错误: {response.status_code}"}
            
        # 处理流式响应
        total_processed = 0
        total_success = 0
        
        for line in response.iter_lines():
            if line:
                try:
                    result = json.loads(line)
                    total_processed += 1
                    
                    if result.get('status') == 'success' and 'image_content' in result:
                        # 获取图片内容
                        image_data = base64.b64decode(result['image_content'].split(',')[1] if ',' in result['image_content'] else result['image_content'])
                        if callback:
                            callback({
                                "success": True,
                                "original_url": result.get('original_url', ''),
                                "data": image_data,
                                "index": result.get('index', 0),
                                "total": result.get('total', 1)
                            })
                        total_success += 1
                    elif callback:
                        callback({
                            "success": False,
                            "message": result.get('message', '未知错误'),
                            "original_url": result.get('original_url', ''),
                            "index": result.get('index', 0),
                            "total": result.get('total', 1)
                        })
                            
                except Exception as e:
                    continue
        
        return {
            "success": True,
            "total_processed": total_processed,
            "total_success": total_success
        }
            
    except Exception as e:
        send_notification(f"调用API时发生错误: {str(e)}")
        return {"success": False, "message": str(e)}

# 处理添加背景任务的主函数
def process_add_background_task(pg_name, pagetype, config=None):
    try:
        pg = jingrow.get_pg(pagetype, pg_name)
        if not pg:
            return False
     
        # 更新状态为"任务进行中"
        update_status_progress(pagetype, pg_name, "任务进行中", 20)
       
        # 获取关联的文件
        file_pgs = jingrow.db.get_all(
            "File",
            {"attached_to_pagetype": pagetype, "attached_to_name": pg_name},
            ["file_url", "file_name"]
        )
                   
        # 构建所有图片URL
        site_url = jingrow.utils.get_url()
        image_urls = []
        
        # 按照原始附件列表的顺序添加图片
        for img_file in file_pgs:
            # 检查文件扩展名
            if Path(img_file.file_name).suffix.lower() != '.png':
                continue
                
            # 只选择已去背的图片，且跳过已添加背景的图片
            if "_nobg" not in img_file.file_name or "_added_bg" in img_file.file_name:
                continue
                
            # 构建完整URL
            full_image_url = f"{site_url}{img_file.file_url}"
            image_urls.append(full_image_url)
            
        # 反转列表顺序，使第一张图片排在前面
        image_urls.reverse()
        
        if not image_urls:
            update_status_progress(pagetype, pg_name, "任务未完成", 0)
            return False
        
        update_progress(pg_name, 40, pagetype)
        
        # 用于跟踪处理进度
        total_images = len(image_urls)
        processed_images = 0
        success_count = 0
        
        def handle_image_result(result):
            nonlocal processed_images, success_count
            processed_images += 1
            
            if not result.get("success", False):
                return
                
            # 从原始URL中提取文件名
            original_url = result.get("original_url", "")
            # 找到对应的原始文件信息
            matching_files = [f for f in file_pgs if original_url.endswith(f.file_url)]
            
            if not matching_files:
                return
                
            original_file = matching_files[0]
            base_filename = os.path.splitext(os.path.basename(original_file.file_name))[0]
            # 保存为PNG格式
            file_name = f"{base_filename}_added_bg.png"
            
            # 上传处理后的图片
            if upload_image(pg_name, result.get("data"), file_name, pagetype):
                success_count += 1
            
            # 更新进度
            progress = 40 + int((processed_images / total_images) * 60)
            update_progress(pg_name, progress, pagetype)
        
        # 调用API并实时处理返回的图片
        api_result = call_jingrow_add_bg_api(image_urls, config, callback=handle_image_result)
        
        # 检查API调用结果
        if not api_result.get("success", False):
            update_status_progress(pagetype, pg_name, "任务未完成", 0)
            return False
        
        # 更新任务状态
        final_status = "已添加背景颜色" if success_count > 0 else "任务未完成"
        final_progress = 100 if success_count > 0 else 0
        update_status_progress(pagetype, pg_name, final_status, final_progress)
        
        return success_count > 0
        
    except Exception as e:
        jingrow.log_error("添加背景任务处理失败", f"处理添加背景任务时发生错误: {str(e)}")
        update_status_progress(pagetype, pg_name, "任务未完成", 0)
        return False
