import os
import requests
import json
from pathlib import Path
import uuid

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

def _build_callback_data(result, success, data=None, message=None):
    """构建callback数据的辅助函数"""
    callback_data = {
        "success": success,
        "original_url": result.get('original_url', ''),
        "index": result.get('index', 0),
        "total": result.get('total', 1)
    }
    if data is not None:
        callback_data["data"] = data
    if message is not None:
        callback_data["message"] = message
    return callback_data

def _download_image(image_url, timeout=60):
    """下载图片数据"""
    try:
        response = requests.get(image_url, timeout=timeout)
        response.raise_for_status()
        return response.content, None
    except requests.RequestException as e:
        return None, f"下载图片失败: {str(e)}"

# 调用去背景API
def call_jingrow_rmbg_api(image_urls, callback=None):
    api_url = f"{get_jingrow_api_url()}/rmbg/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密钥未设置"}
    
    try:
        response = requests.post(api_url, json={"urls": image_urls}, 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 not line:
                continue
                
            try:
                result = json.loads(line)
                total_processed += 1
                
                if result.get('status') == 'success' and 'image_url' in result:
                    image_data, error_msg = _download_image(result['image_url'])
                    if image_data:
                        if callback:
                            callback(_build_callback_data(result, True, data=image_data))
                        total_success += 1
                    elif callback:
                        callback(_build_callback_data(result, False, message=error_msg))
                elif callback:
                    callback(_build_callback_data(result, False, message=result.get('message', '未知错误')))
                    
            except json.JSONDecodeError:
                continue
        
        return {"success": True, "total_processed": total_processed, "total_success": total_success}
            
    except requests.RequestException as e:
        send_notification(f"调用API时发生错误: {str(e)}")
        return {"success": False, "message": str(e)}

# 处理去背景任务的主函数
def process_remove_background_task(pg_name, pagetype):
    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"]
        )    
            
        # 过滤出支持的图片格式
        allowed_extensions = {'.png', '.jpg', '.jpeg', '.webp'}
        
        # 获取reference_image字段的值
        reference_images = {
            getattr(pg, field) 
            for field in ["reference_image_1", "reference_image_2", "reference_image_3", "reference_image_4"]
            if hasattr(pg, field) and getattr(pg, field)
        }
        
        # 过滤图片，排除reference_image字段的值和已处理过的图片(_nobg)
        image_files = [
            f for f in file_pgs 
            if Path(f.file_name).suffix.lower() in allowed_extensions 
            and "_nobg" not in f.file_name
            and f.file_url not in reference_images
        ]
                
        if not image_files:
            update_status_progress(pagetype, pg_name, "任务未完成", 0)
            return False
        
        update_progress(pg_name, 40, pagetype)
        
        # 构建所有图片URL，反转顺序使第一张图片排在前面
        site_url = jingrow.utils.get_url()
        image_urls = [f"{site_url}{img.file_url}" for img in reversed(image_files)]
        
        # 用于跟踪处理进度
        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
                
            # 找到对应的原始文件信息
            original_url = result.get("original_url", "")
            original_file = next((f for f in image_files if original_url.endswith(f.file_url)), None)
            
            if not original_file:
                return
                
            # 生成文件名：保存为PNG格式，因为去背景通常会产生透明背景
            base_filename = os.path.splitext(os.path.basename(original_file.file_name))[0]
            file_name = f"{base_filename}_nobg_{uuid.uuid4().hex[:4]}.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_rmbg_api(image_urls, 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)
        
    except Exception:
        update_status_progress(pagetype, pg_name, "任务未完成", 0)
        return False
