import io
import os
import json
import requests
import base64
import uuid
from pathlib import Path
from typing import Dict, List, Any, Optional

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

# 调用Jingrow图像放大API
def call_jingrow_upscale_api(image_urls: List[str], callback=None, config: Dict = None) -> Dict:
    try:
        # 获取API URL
        api_url = f"{get_jingrow_api_url()}/jupscale/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}
        
        # 准备请求数据
        request_data = {"image_urls": image_urls}
        if config:
            request_data["config"] = config
        
        images_data = []
        total_processed = 0
        total_success = 0
        
        response = requests.post(
            api_url, 
            json=request_data, 
            headers=headers, 
            stream=True, 
            timeout=180
        )
        
        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_url' in result:
                        image_url = result.get('image_url')
                        
                        # 下载图片
                        img_resp = requests.get(image_url, timeout=30)
                        if img_resp.status_code == 200:
                            img_data = img_resp.content
                            
                            if callback:
                                callback({
                                    "success": True,
                                    "image_data": img_data,
                                    "original_url": result.get('original_image_url', '')
                                })
                            else:
                                images_data.append({
                                    "image_data": img_data,
                                    "original_url": result.get('original_image_url', '')
                                })
                            total_success += 1
                        else:
                            error_message = f"下载处理后的图片失败: HTTP {img_resp.status_code}"
                            if callback:
                                callback({
                                    "success": False,
                                    "message": error_message
                                })
                    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 as e:
                    jingrow.log_error("解析API响应失败", str(e))
                    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:
        jingrow.log_error("API异常", f"调用批量图像放大API异常：{str(e)}")
        return {"success": False, "message": f"调用批量图像放大API异常：{str(e)}"}

# 处理图像放大任务的主函数
def process_upscale_image_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"]
        )   
            
        # 过滤出 .png 和 .jpg 格式的文件
        allowed_extensions = ['.png', '.jpg', '.jpeg']
        image_files = [f for f in file_pgs if Path(f.file_name).suffix.lower() in allowed_extensions]
                
        if not image_files:
            update_status_progress(pagetype, pg_name, "任务未完成", 0)
            jingrow.log_error(f"未找到图片附件", f"{pagetype}:{pg_name}")
            return False
        
        update_progress(pg_name, 40, pagetype)
        
        # 获取自定义配置
        config = None
        if hasattr(pg, "upscale_config") and pg.upscale_config:
            try:
                config = json.loads(pg.upscale_config)
            except Exception:
                pass
        
        # 构建图片URL列表
        site_url = jingrow.utils.get_url()
        image_urls = []
        
        for img_file in image_files:
            # 跳过已放大的图片
            if "_upscaled" 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)
            send_notification("没有可放大的图片")
            return False
            
        processed_count = 0
        success_count = 0
        expected_total = len(image_urls)
        
        # 处理每个图片的回调函数
        def handle_image_result(result):
            nonlocal processed_count, success_count
            
            if not result.get("success", False):
                processed_count += 1
                return
            
            # 获取原始图片URL
            original_url = result.get("original_url", "")
            
            # 如果能找到对应的文件记录，使用其文件名
            base_filename = None
            extension = ".png"  # 默认扩展名
            
            for img_file in image_files:
                if site_url + img_file.file_url == original_url:
                    base_filename = os.path.splitext(os.path.basename(img_file.file_name))[0]
                    extension = os.path.splitext(img_file.file_name)[1]
                    break
                    
            # 如果找不到对应的文件记录，生成随机文件名
            if not base_filename:
                unique_id = str(uuid.uuid4())[:8]
                base_filename = f"upscaled_{unique_id}"
            
            # 生成文件名
            file_name = f"{base_filename}_upscaled{extension}"
            
            # 上传图片
            if upload_image(pg_name, result.get("image_data"), file_name, pagetype):
                success_count += 1
            
            processed_count += 1
            
            # 更新进度
            current_progress = 40 + (processed_count * 60) // expected_total
            if hasattr(pg, "progress"):
                update_progress(pg_name, min(current_progress, 100), pagetype)
        
        # 调用API处理图片
        api_result = call_jingrow_upscale_api(image_urls, callback=handle_image_result, config=config)
        
        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:
        jingrow.log_error("处理图像放大任务时发生异常", f"{pagetype}:{pg_name}, 错误: {str(e)}")
        update_status_progress(pagetype, pg_name, "任务未完成", 0)
        return False
