import json
import base64
import requests
import uuid
from typing import Dict, List, Any, Optional

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

# 定义配置变量
images_per_prompt = int(jingrow.db.get_single_value("Ai Settings", "images_per_prompt") or 4)
add_title_to_image_name = jingrow.db.get_single_value("Ai Settings", "add_title_to_image_name") or 0

@jingrow.whitelist()
def call_jingrow_jart_v1_api(prompts: List[str], config: Dict = None, callback = None) -> Dict:
    try:
        # 获取API URL
        api_url = f"{get_jingrow_api_url()}/jart_v1/batch"
        
        headers = get_jingrow_api_headers()
        if not headers:
            error_message = "Jingrow API 未设置，请在 AI设置 中设置 Jingrow Api Key 和 Jingrow Api Secret"
            send_notification(error_message)
            return {"success": False, "message": error_message}
        
        default_config = {"images_per_prompt": images_per_prompt}
        if config:
            default_config.update(config)
        
        request_data = {
            "prompts": prompts,
            "config": default_config
        }
        
        images_data = []
        total_processed = 0
        total_success = 0
        
        response = requests.post(
            api_url, 
            json=request_data, 
            headers=headers, 
            stream=True, 
            timeout=60
        )
        
        if response.status_code != 200:
            try:
                error_data = response.json()
                error_message = error_data.get("message") or error_data.get("error") or f"API请求失败 (HTTP {response.status_code})"
            except:
                error_message = f"API请求失败 (HTTP {response.status_code})"
            
            if response.status_code in [401, 403]:
                auth_error_message = "API认证失败，请检查 Jingrow Api Key 和 Jingrow Api Secret 是否正确"
                send_notification(auth_error_message)
                return {"success": False, "message": auth_error_message}
            
            send_notification(f"API请求失败: {error_message}")
            return {"success": False, "message": error_message}
        
        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_base64 = result['image_content'].split(',')[1]
                        img_data = base64.b64decode(image_base64)
                        
                        if callback:
                            callback({
                                "success": True,
                                "image_data": img_data
                            })
                        else:
                            images_data.append({
                                "image_data": img_data
                            })
                        total_success += 1
                    else:
                        error_message = result.get('message', '未知错误')
                        if "余额不足" in error_message:
                            send_notification(f"{error_message}")
                            return {"success": False, "message": error_message}
                        elif callback:
                            callback({
                                "success": False,
                                "message": error_message
                            })
                except Exception:
                    continue
        
        if callback:
            return {
                "success": True,
                "total_processed": total_processed,
                "total_success": total_success
            }
            
        if not images_data and not callback:
            return {"success": False, "message": "未能获取到任何图像数据"}
        
        return {"success": True, "data": images_data}
        
    except Exception as e:
        return {"success": False, "message": f"调用批量图像生成API异常：{str(e)}"}

def process_jart_task(pg_name: str, pagetype: str) -> bool:
    try:
        pg = jingrow.get_pg(pagetype, pg_name)
        if not pg:
            return False
        
        prompt = None
        
        if hasattr(pg, "prompts") and pg.prompts:
            prompt = pg.prompts
        elif hasattr(pg, "prompts_zh") and pg.prompts_zh:
            prompt = get_translated_prompts(pg.prompts_zh)
        
        if not prompt:
            update_status_progress(pagetype, pg_name, "任务未完成", 0)
            send_notification("提示词未设置，请检查配置")
            return False
        
        update_status_progress(pagetype, pg_name, "任务进行中", 20)
        
        processed_count = 0
        success_count = 0
        expected_total = images_per_prompt
        
        def handle_image_result(result):
            nonlocal processed_count, success_count
            
            if not result.get("success", False):
                processed_count += 1
                return
            
            unique_id = str(uuid.uuid4())[:8]
            fields_to_include = ""
            
            if hasattr(pg, "serial_number") and pg.serial_number:
                fields_to_include += f"{pg.serial_number}_"
                
            if int(add_title_to_image_name) and hasattr(pg, "title") and pg.title:
                fields_to_include += f"{pg.title}_"
                
            index_info = str(success_count + 1)
            image_filename = f"{fields_to_include}{pg_name}_{unique_id}_{index_info}.png"
            
            if upload_image(pg_name, result.get("image_data"), image_filename, pagetype):
                success_count += 1
            
            processed_count += 1
            
            current_progress = 20 + (processed_count * 80) // expected_total
            if hasattr(pg, "progress"):
                update_progress(pg_name, min(current_progress, 100), pagetype)
        
        api_result = call_jingrow_jart_v1_api([prompt], callback=handle_image_result)
        
        if not api_result.get("success", False):
            error_message = api_result.get("message", "未知错误")
            
            if "余额不足" in error_message:
                update_status_progress(pagetype, pg_name, "任务未完成", 0)
                if hasattr(pg, "api_error"):
                    pg.api_error = "credit_insufficient"
                    pg.save()
                    jingrow.db.commit()
                
                send_notification(f"任务未完成：{error_message}")
                return False
            
            update_status_progress(pagetype, pg_name, "任务未完成", 0)
            send_notification(f"API调用失败: {error_message}")
            return False
        
        if success_count == 0:
            update_status_progress(pagetype, pg_name, "任务未完成", 0)
            send_notification("未成功生成任何图像")
            return False
        
        update_status_progress(pagetype, pg_name, "已出图", 100)
        return True
        
    except Exception as e:
        update_status_progress(pagetype, pg_name, "任务未完成", 0)
        return False