import json
import sys
import os
import io
import requests
import time
import random
import re
import uuid
import urllib.request
import urllib3
from pathlib import Path
from urllib.parse import urlparse
from PIL import Image
from typing import Dict, List, Any, Optional

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

# 默认配置变量
api_url = ""
application_id = ""
data_id = ""
data_version = ""
session_id = ""
midjourney_suffix = "mj"
add_title_to_image_name = 0
http_proxy = ""
https_proxy = ""
midjourney_options = {"ar": "1:1", "v": "6.1"}

# 禁用不安全请求警告
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

class Midjourney:
    def __init__(self, channel_id, oauth_token):
        self.API_URL = api_url
        self.APPLICATION_ID = application_id
        self.DATA_ID = data_id
        self.DATA_VERSION = data_version
        self.SESSION_ID = session_id
        
        self.channel_id = channel_id
        self.oauth_token = oauth_token

        # 设置代理 - 仅当代理配置存在且不为空时才使用
        self.proxies = {}
        if http_proxy and http_proxy.strip():
            self.proxies['http'] = http_proxy
        if https_proxy and https_proxy.strip():
            self.proxies['https'] = https_proxy

        self.client = requests.Session()
        self.client.headers.update({
            'Authorization': oauth_token
        })
        # 为会话设置代理 - 仅当代理不为空时
        if self.proxies:
            self.client.proxies.update(self.proxies)

        try:
            response = self.client.get(f'{self.API_URL}/channels/{self.channel_id}')
            data = response.json()
            
            self.guild_id = data['guild_id']
            response = self.client.get(f'{self.API_URL}/users/@me')
            data = response.json()
            
            self.user_id = data['id']
        except Exception as e:
            print(f"初始化Midjourney客户端失败: {str(e)}")
            raise Exception(f"初始化Midjourney客户端失败: {str(e)}")

    @staticmethod
    def first_where(array, key, value=None):
        for item in array:
            if callable(key) and key(item):
                return item
            if isinstance(key, str) and item[key].startswith(value):
                return item
        return None

    # 使用给定的提示词生成图像请求
    def imagine(self, prompt, seed=None, pg_name=None, pagetype=None):
        params = {
            'type': 2,
            'application_id': self.APPLICATION_ID,
            'guild_id': self.guild_id,
            'channel_id': self.channel_id,
            'session_id': self.SESSION_ID,
            'data': {
                'version': self.DATA_VERSION,
                'id': self.DATA_ID,
                'name': 'imagine',
                'type': 1,
                'options': [{
                    'type': 3,
                    'name': 'prompt',
                    'value': prompt
                }],
                'application_command': {
                    'id': self.DATA_ID,
                    'application_id': self.APPLICATION_ID,
                    'version': self.DATA_VERSION,
                    'default_member_permissions': None,
                    'type': 1,
                    'nsfw': False,
                    'name': 'imagine',
                    'description': 'Create images with Midjourney',
                    'dm_permission': True,
                    'options': [{
                        'type': 3,
                        'name': 'prompt',
                        'description': 'The prompt to imagine',
                        'required': True
                    }]
                },
                'attachments': []
            }
        }

        try:
            r = self.client.post(f'{self.API_URL}/interactions', json=params)
            time.sleep(20)  # 初始等待时间

            imagine_message = None

            count = 0
            while imagine_message is None:
                imagine_message = self.get_imagine(prompt, count, seed, pg_name, pagetype)
                if imagine_message is None:
                    time.sleep(3)  # 每次循环等待时间
                    count += 1
                    if count > 60:
                        break

            return imagine_message
        except Exception as e:
            print(f"发送Imagine请求失败: {str(e)}")
            raise Exception(f"发送Imagine请求失败: {str(e)}")

    # 获取生成图像的消息
    def get_imagine(self, prompt, count=0, seed=None, pg_name=None, pagetype=None):
        try:
            # 获取最近的5条消息
            response = self.client.get(f'{self.API_URL}/channels/{self.channel_id}/messages?limit=5')
            data = response.json()
            
            def criteria(item):
                if seed is not None and f"--seed {seed}" in item['content'] and pg_name and pagetype:
                    # 匹配百分比，支持多种格式：(93%) 或 93%
                    progress_match = re.search(r'(?:\(|（)?(\d+)%(?:\)|）)?', item['content'])
                    if progress_match:
                        progress_value = int(progress_match.group(1))
                        # 如果是进行中的消息，返回False继续查找完成的消息
                        if "%" in item['content']:
                            return False
                
                # 排除进行中消息
                if "%" in item['content']:
                    return False
                    
                # seed 匹配
                if seed is not None:
                    seed_pattern = f"--seed {seed}"
                    if seed_pattern not in item['content']:
                        return False
                return True

            raw_message = self.first_where(data, criteria)

            if raw_message is None:
                return None
                
            r_final = {
                'id': raw_message['id'],
                'prompt': prompt,
                'raw_message': raw_message
            }
            
            return r_final
        except Exception as e:
            print(f"获取通道消息失败: {str(e)}")
            return None

    # 生成图像并返回生成的消息ID
    def generate(self, prompt, options={}, sources={}, pg_name=None, pagetype=None):
        try:
            prompt = prompt.strip()

            # 只在没有设置时才使用默认值
            if "ar" not in options or options['ar'] == "":
                options['ar'] = '3:2'

            if "seed" not in options or options['seed'] == "":
                options['seed'] = random.randint(0, 4294967295)

            # Force seed as last parameter in the dict
            seed = options.pop('seed')
            options['seed'] = seed

            parameter = ""
            no_value_key = ['relax', 'fast', 'turbo', 'raw', 'tile']
            for key, value in options.items():
                if key in no_value_key:
                    parameter += f" --{key}"
                else:
                    parameter += f" --{key} {value}"

            # Image sources
            if 'images' in sources:
                prompt = " ".join(sources['images']) + " " + prompt

            # Add omni reference
            if 'omni' in sources and len(sources['omni']) > 0:
                prompt = prompt + " --oref " + " ".join(sources['omni'])

            # Add caracters
            if 'caracters' in sources and len(sources['caracters']) > 0:
                prompt = prompt + " --cref " + " ".join(sources['caracters'])

            # Add styles
            if 'styles' in sources and len(sources['styles']) > 0:
                prompt = prompt + " --sref " + " ".join(sources['styles'])

            # On ajoute une seed pour éviter les doublons
            prompt = prompt + parameter
            
            imagine = self.imagine(prompt, seed, pg_name, pagetype)

            return {
                'imagine_message_id': imagine['id'],
                'raw_message': imagine['raw_message'],
                'seed': seed
            }
        except Exception as e:
            print(f"生成图像失败: {str(e)}")
            raise Exception(f"生成图像失败: {str(e)}")

    # 从提示词中提取参数和来源
    def get_parameter_from_prompt(self, prompt):
        sources = {}
        # Get all sources from prompt
        # Get --sref http://xxxx, http://xxxx, http://xxxx
        sources['styles'] = []
        styles_match = re.search(r'--sref\s+((?:https?://\S+(?:,\s*)?)+)', prompt)
        if styles_match:
            sources['styles'] = re.findall(r'https?://\S+', styles_match.group(0))
            prompt = prompt[:styles_match.start()] + prompt[styles_match.end():]

        # Get --cref http://xxxx, http://xxxx, http://xxxx
        sources['caracters'] = []
        caracters_match = re.search(r'--cref\s+((?:https?://\S+(?:,\s*)?)+)', prompt)
        if caracters_match:
            sources['caracters'] = re.findall(r'https?://\S+', caracters_match.group(0))
            prompt = prompt[:caracters_match.start()] + prompt[caracters_match.end():]

        # Get all parameters from prompt
        parameters = re.findall(r'--(\w+\s[\w.:]+)', prompt)
        # Clean prompt
        for parameter in parameters:
            prompt = prompt.replace(f'--{parameter}', '')

        # Delete multiple spaces
        prompt = re.sub(' {2,}', ' ', prompt)
        prompt = prompt.strip()
        return {parameter.split(' ')[0].replace('--', ''): parameter.split(' ')[1] for parameter in
                parameters}, sources, prompt


# 将一张大图切割成四张子图
def split_image_into_four(image_data):
    """
    将Midjourney生成的拼接图片分割成四张单独的图片
    简单的四象限分割：上下各一半，左右各一半
    
    Args:
        image_data: 图像二进制数据
        
    Returns:
        包含四张子图像数据的列表
    """
    try:
        # 从二进制数据创建图像对象
        img = Image.open(io.BytesIO(image_data))
        width, height = img.size
        
        # 确认图像尺寸足够大可以分割
        if width < 500 or height < 500:
            print(f"图像尺寸太小无法分割: {width}x{height}, 需要宽度和高度都大于500px")
            return None
        
        # 计算每个象限的尺寸 - 简单的四象限分割
        half_width = width // 2
        half_height = height // 2
        
        # 分割图像为四个象限
        top_left = img.crop((0, 0, half_width, half_height))
        top_right = img.crop((half_width, 0, width, half_height))
        bottom_left = img.crop((0, half_height, half_width, height))
        bottom_right = img.crop((half_width, half_height, width, height))
        
        # 将每个象限转换为二进制数据
        quadrants = []
        for quadrant in [top_left, top_right, bottom_left, bottom_right]:
            buffer = io.BytesIO()
            quadrant.save(buffer, format="PNG")
            quadrants.append(buffer.getvalue())
            
        return quadrants
    except Exception as e:
        print(f"分割图像失败: {str(e)}")
        return None

def call_japi_midjourney(prompt: str, config: Dict = None, callback = None) -> Dict:

    try:
        api_url = f"{get_jingrow_cloud_api_url()}/midjourney/generate"
        
        headers = get_jingrow_cloud_api_headers()
        if not headers:
            error_message = "Jingrow API 未设置，请在 AI设置 中设置 Jingrow Api Key 和 Jingrow Api Secret"
            return {"success": False, "message": error_message}
        
        # 准备请求数据
        request_data = {
            "prompt": prompt,
            "config": config or {}
        }
        
        images_data = []
        total_processed = 0
        total_success = 0
        
        response = requests.post(
            api_url, 
            json=request_data, 
            headers=headers, 
            stream=True
        )
        
        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认证失败，请检查认证信息"
                return {"success": False, "message": auth_error_message}
            
            return {"success": False, "message": error_message}
        
        for line in response.iter_lines():
            if line:
                try:
                    result = json.loads(line.decode('utf-8'))
                    
                    # 处理进度更新
                    if result.get('status') == 'progress':
                        progress = result.get('progress', 0)
                        if callback:
                            callback({
                                "status": "progress",
                                "progress": progress,
                                "seed": result.get('seed')
                            })
                        continue
                    
                    # 处理最终结果
                    if result.get('status') == 'success' and 'image_urls' in result:
                        for image_url in result['image_urls']:
                            # 下载图片
                            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,
                                        "image_url": image_url
                                    })
                                else:
                                    images_data.append({
                                        "image_data": img_data,
                                        "image_url": image_url
                                    })
                                total_success += 1
                            else:
                                error_message = f"下载图片失败: HTTP {img_resp.status_code}"
                                if callback:
                                    callback({
                                        "success": False,
                                        "message": error_message
                                    })
                        total_processed += 1
                    else:
                        error_message = result.get('message', '未知错误')
                        if "余额不足" in error_message:
                            return {"success": False, "message": error_message}
                        elif callback:
                            callback({
                                "success": False,
                                "message": error_message
                            })
                except Exception as 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:
        return {"success": False, "message": f"调用Midjourney API异常：{str(e)}"}

def process_midjourney_node_task(prompt: str, config: Dict = None) -> Dict:
    """
    专门为AI绘画节点设计的Midjourney处理函数
    根据jingrow_midjourney配置自动选择使用JAPI还是Discord API
    
    Args:
        prompt: 提示词
        config: 配置参数，包含options等
        
    Returns:
        Dict: 包含success状态和图像数据或错误信息的字典
    """
    try:
        if not prompt:
            return {"success": False, "message": "提示词不能为空"}
            
        config = config or {}
        # 重新获取AI配置（确保使用最新配置）
        current_ai_config = get_ai_settings_from_jingrow()
        if current_ai_config is None:
            return {"success": False, "message": "获取AI Settings配置失败，请检查网络连接和登录状态"}
        
        jingrow_midjourney = current_ai_config.get('jingrow_midjourney', True)
        
        # 更新全局配置变量
        global api_url, application_id, data_id, data_version, session_id, midjourney_suffix, add_title_to_image_name, http_proxy, https_proxy, midjourney_options
        api_url = current_ai_config.get('midjourney_api_url') or "https://api.midjourney.com"
        application_id = current_ai_config.get('midjourney_application_id') or ""
        data_id = current_ai_config.get('midjourney_data_id') or ""
        data_version = current_ai_config.get('midjourney_data_version') or ""
        session_id = current_ai_config.get('midjourney_session_id') or ""
        midjourney_suffix = current_ai_config.get('midjourney_suffix') or "mj"
        add_title_to_image_name = current_ai_config.get('add_title_to_image_name') or 0
        http_proxy = current_ai_config.get('http_proxy') or ""
        https_proxy = current_ai_config.get('https_proxy') or ""
        
        # 获取Midjourney选项配置
        midjourney_options_str = current_ai_config.get('midjourney_options')
        if midjourney_options_str:
            try:
                midjourney_options = json.loads(midjourney_options_str)
            except:
                midjourney_options = {"ar": "1:1", "v": "6.1"}
        else:
            midjourney_options = {"ar": "1:1", "v": "6.1"}
        
        # 从全局配置获取midjourney_options
        options = midjourney_options.copy()
        
        if "options" in config:
            options.update(config.get("options", {}))
        
        # 设置默认选项（如果全局配置中也没有）
        if "ar" not in options or options['ar'] == "":
            options['ar'] = '1:1'
        if "seed" not in options or options['seed'] == "":
            options['seed'] = random.randint(0, 4294967295)
        
        # 根据配置选择使用JAPI还是Discord API
        if jingrow_midjourney:
            # 使用JAPI
            images_data = []
            
            def handle_image_result(result):
                nonlocal images_data
                
                if result.get("status") == "progress":
                    # 进度更新，AI节点中暂时不需要处理
                    return
                
                if not result.get("success", False):
                    return
                
                # 收集图像数据
                if result.get("image_data"):
                    images_data.append({
                        "image_data": result.get("image_data"),
                        "image_url": result.get("image_url")
                    })
            
            # 调用JAPI处理图片
            api_result = call_japi_midjourney(prompt, config=config, callback=handle_image_result)
            
            if not api_result.get("success", False):
                error_message = api_result.get("message", "未知错误")
                return {"success": False, "message": error_message}
            
            if not images_data:
                return {"success": False, "message": "未成功生成任何图像"}
                
            return {"success": True, "data": images_data}
            
        else:
            # 使用Discord API
            discord_channel_id = current_ai_config.get('discord_channel_id')
            discord_user_token = current_ai_config.get('discord_user_token')

            if not discord_channel_id or not discord_user_token:
                return {"success": False, "message": "Discord配置不完整，请在AI设置中配置discord_channel_id和discord_user_token"}

            # 初始化Midjourney客户端
            midjourney = Midjourney(discord_channel_id, discord_user_token)

        # 处理sources参数（参考图片）
        sources = config.get("sources", {})
        
        # 处理负面提示词
        if config.get("negative_prompt"):
            prompt = f"{prompt} --no {config.get('negative_prompt')}"
                
        # 生成图像并返回生成的消息ID
        result = midjourney.generate(prompt, options, sources, None, None)
        
        if not result:
            return {"success": False, "message": "Midjourney生成失败"}
        
        imagine_message_id = result.get('imagine_message_id')
        raw_message = result.get('raw_message')
        seed = result.get('seed')

        # 确保生成的消息包含正确的seed
        if f"--seed {seed}" not in raw_message['content']:
            return {"success": False, "message": "Seed不匹配！该消息不属于当前任务。"}

        # 设置代理用于下载Discord图片
        proxies = {}
        if http_proxy and http_proxy.strip():
            proxies['http'] = http_proxy
        if https_proxy and https_proxy.strip():
            proxies['https'] = https_proxy
            
        # 从原始消息中获取合并图像URL
        original_image_url = None
        if 'attachments' in raw_message and len(raw_message['attachments']) > 0:
            original_image_url = raw_message['attachments'][0]['url']
        
        if not original_image_url:
            return {"success": False, "message": "无法获取原始图像URL"}
                
        # 下载原始合并图像
        try:
            if proxies:
                response = requests.get(original_image_url, proxies=proxies, timeout=30)
            else:
                response = requests.get(original_image_url, timeout=30)
                
            if response.status_code != 200:
                raise Exception(f"下载原始图像失败: 状态码 {response.status_code}")
                
            original_image_data = response.content
        except Exception as e:
            raise Exception(f"下载原始图像失败: {str(e)}")
              
        # 分割原始图像为四个象限
        image_quadrants = split_image_into_four(original_image_data)
        if not image_quadrants:
            raise Exception("分割图像失败")

        # 准备返回的图像数据
        images_data = []
        for i, quadrant_data in enumerate(image_quadrants):
            images_data.append({
                "image_data": quadrant_data,
                "image_url": None  # Discord模式下没有URL
            })
            
        return {"success": True, "data": images_data}
        
    except Exception as e:
        return {"success": False, "message": f"处理Midjourney任务异常：{str(e)}"} 