# Copyright (c) 2025, JINGROW and contributors
# For license information, please see license.txt

import requests
from typing import Dict, Any
import logging
import io

from PIL import Image, ImageOps

import jingrow
from jingrow.utils.auth import get_jingrow_cloud_api_headers, get_jingrow_cloud_api_url

logger = logging.getLogger(__name__)


@jingrow.whitelist(allow_guest=True)
def remove_background_from_file_free(file) -> Dict[str, Any]:
    """
    免费图片去背景工具（无需认证）
    参考 remove_background_from_file 实现，调用 Jingrow Cloud API 实现图片背景移除
    
    Args:
        file: 上传的文件对象
    
    Returns:
        dict: 处理结果，格式与 remove_background_from_file 一致
    """
    try:
        # 获取认证头（使用系统配置的 API Key）
        headers = get_jingrow_cloud_api_headers()
        
        if not headers or not headers.get('Authorization'):
            return {
                "success": False,
                "error": "API密钥未设置，请联系管理员配置"
            }
        
        # 读取文件内容
        # whitelist 服务传递的 file 是 FastAPI 的 UploadFile 对象
        # UploadFile.file 是底层的 SpooledTemporaryFile，可以同步读取
        # 需要先重置文件指针到开头（因为可能已经被读取过）
        file.file.seek(0)
        file_content = file.file.read()
        
        if not file_content or len(file_content) == 0:
            return {
                "success": False,
                "error": "文件内容为空"
            }
        
        # 获取文件扩展名和内容类型
        filename = getattr(file, 'filename', None) or 'image.png'
        ext = filename.split('.')[-1].lower() if '.' in filename else 'png'
        if ext == 'jpg':
            ext = 'jpg'
            content_type = 'image/jpeg'
        elif ext == 'jpeg':
            ext = 'jpg'
            content_type = 'image/jpeg'
        elif ext == 'png':
            ext = 'png'
            content_type = 'image/png'
        elif ext == 'webp':
            ext = 'webp'
            content_type = 'image/webp'
        else:
            ext = 'png'
            content_type = 'image/png'
        
        # 处理EXIF orientation（参考 remove_background_from_file）
        image = Image.open(io.BytesIO(file_content))
        image = ImageOps.exif_transpose(image)
        output = io.BytesIO()
        save_format = image.format or content_type.split('/')[-1].upper()
        image.save(output, format=save_format)
        image_bytes = output.getvalue()
        
        # 调用文件上传端点（与 remove_background_from_file 使用相同的接口）
        api_url = f"{get_jingrow_cloud_api_url()}/rmbg/file/free"
        
        # 移除 Content-Type，让 requests 自动设置 multipart/form-data
        upload_headers = {k: v for k, v in headers.items() if k.lower() != 'content-type'}
        files = {
            'file': (f'image.{ext}', image_bytes, content_type)
        }
        
        # 转发请求
        response = requests.post(
            api_url,
            files=files,
            headers=upload_headers,
            timeout=180
        )
        
        if response.status_code == 200:
            result_data = response.json()
            
            # 使用 image_url 字段，返回格式与 remove_background_from_file 一致
            if 'image_url' in result_data:
                image_url = result_data.get('image_url', '')
                return {
                    "success": True,
                    "data": [{
                        "success": True,
                        "image_url": image_url,
                        "index": 1,
                        "total": 1
                    }],
                    "total_processed": 1,
                    "total_success": 1
                }
            
            return {
                "success": False,
                "error": result_data.get('error', '未找到图片URL')
            }
        else:
            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})"
            
            return {
                "success": False,
                "error": error_message
            }
            
    except Exception as e:
        logger.error(f"调用图片去背景API异常：{str(e)}", exc_info=True)
        return {
            "success": False,
            "error": f"调用图片去背景API异常：{str(e)}"
        }


@jingrow.whitelist()
def remove_background_from_file(file) -> Dict[str, Any]:
    """
    图片去背景工具
    调用 Jingrow Cloud API 实现图片背景移除
    
    Args:
        file: 上传的文件对象（multipart/form-data）
    
    Returns:
        dict: 处理结果
    """
    try:
        # 获取认证头
        headers = get_jingrow_cloud_api_headers()
        
        if not headers or not headers.get('Authorization'):
            return {
                "success": False,
                "error": "API密钥未设置，请在设置中配置 Jingrow Cloud Api Key 和 Jingrow Cloud Api Secret"
            }
        
        # 读取文件内容
        file.file.seek(0)
        file_content = file.file.read()
        
        if not file_content or len(file_content) == 0:
            return {
                "success": False,
                "error": "文件内容为空"
            }
        
        # 获取文件扩展名和内容类型
        filename = getattr(file, 'filename', None) or 'image.png'
        ext = filename.split('.')[-1].lower() if '.' in filename else 'png'
        if ext == 'jpg':
            ext = 'jpg'
            content_type = 'image/jpeg'
        elif ext == 'jpeg':
            ext = 'jpg'
            content_type = 'image/jpeg'
        elif ext == 'png':
            ext = 'png'
            content_type = 'image/png'
        elif ext == 'webp':
            ext = 'webp'
            content_type = 'image/webp'
        else:
            ext = 'png'
            content_type = 'image/png'
        
        # 处理EXIF orientation
        image = Image.open(io.BytesIO(file_content))
        image = ImageOps.exif_transpose(image)
        output = io.BytesIO()
        save_format = image.format or content_type.split('/')[-1].upper()
        image.save(output, format=save_format)
        image_bytes = output.getvalue()
        
        # 调用文件上传端点
        api_url = f"{get_jingrow_cloud_api_url()}/rmbg/file"
        
        # 移除 Content-Type，让 requests 自动设置 multipart/form-data
        upload_headers = {k: v for k, v in headers.items() if k.lower() != 'content-type'}
        files = {
            'file': (f'image.{ext}', image_bytes, content_type)
        }
        
        # 转发请求
        response = requests.post(
            api_url,
            files=files,
            headers=upload_headers,
            timeout=180
        )
        
        if response.status_code == 200:
            result_data = response.json()
            
            # 使用 image_url 字段
            if 'image_url' in result_data:
                image_url = result_data.get('image_url', '')
                return {
                    "success": True,
                    "data": [{
                        "success": True,
                        "image_url": image_url,
                        "index": 1,
                        "total": 1
                    }],
                    "total_processed": 1,
                    "total_success": 1
                }
            
            return {
                "success": False,
                "error": result_data.get('error', '未找到图片URL')
            }
        else:
            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})"
            
            return {
                "success": False,
                "error": error_message
            }
            
    except Exception as e:
        logger.error(f"调用图片去背景API异常：{str(e)}", exc_info=True)
        return {
            "success": False,
            "error": f"调用图片去背景API异常：{str(e)}"
        }
