import json
import sys
import os
import io
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
import random
import re
import uuid
import urllib.request
import urllib3
import logging
from pathlib import Path
from urllib.parse import urlparse
from PIL import Image
import asyncio
from typing import Dict, Any, List, AsyncGenerator, Optional
from settings import settings
from utils import get_new_image_url, is_valid_image_url


# 设置日志记录器
logger = logging.getLogger("midjourney_service")

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

class MidjourneyService:
    def __init__(self):
        """初始化MidjourneyService"""
        self.API_URL = settings.midjourney_api_url
        self.APPLICATION_ID = settings.midjourney_application_id
        self.DATA_ID = settings.midjourney_data_id
        self.DATA_VERSION = settings.midjourney_data_version
        self.SESSION_ID = settings.midjourney_session_id
        
        # 设置代理
        self.proxies = {}
        if settings.http_proxy and settings.http_proxy.strip():
            self.proxies['http'] = settings.http_proxy
        if settings.https_proxy and settings.https_proxy.strip():
            self.proxies['https'] = settings.https_proxy

        # 确保保存目录存在
        os.makedirs(settings.save_dir, exist_ok=True)

    @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

    async def initialize_client(self, oauth_token, channel_id):
        """初始化Discord客户端会话"""
        client = requests.Session()
        
        retry = Retry(total=3, backoff_factor=0.5)
        adapter = HTTPAdapter(max_retries=retry)
        client.mount('http://', adapter)
        client.mount('https://', adapter)
        
        client.headers.update({
            'Authorization': oauth_token
        })
        
        if self.proxies:
            client.proxies.update(self.proxies)

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

    async def imagine(self, client, guild_id, channel_id, prompt, seed=None):
        """发送imagine请求到Discord"""
        params = {
            'type': 2,
            'application_id': self.APPLICATION_ID,
            'guild_id': guild_id,
            'channel_id': 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:
            client.post(f'{self.API_URL}/interactions', json=params)
            await asyncio.sleep(20)
            
            imagine_message = None
            count = 0
            last_progress = 0
            
            while count < settings.max_polling_attempts:
                progress_msg = await self.get_progress_message(client, channel_id, prompt, seed)
                if progress_msg:
                    content = progress_msg.get('content', '')
                    progress_match = re.search(r'(?:\(|（)?(\d+)%(?:\)|）)?', content)
                    if progress_match:
                        progress_value = int(progress_match.group(1))
                        if progress_value > last_progress:
                            last_progress = progress_value
                            logger.info(f"生成进度: {progress_value}%")
                            yield {
                                "status": "progress", 
                                "progress": progress_value, 
                                "message_id": progress_msg.get('id'),
                                "content": content
                            }
                
                imagine_message = await self.get_imagine(client, channel_id, prompt, count, seed)
                if imagine_message:
                    break
                
                await asyncio.sleep(settings.polling_interval)
                count += 1

            if count >= settings.max_polling_attempts:
                logger.error(f"轮询超过最大尝试次数: {settings.max_polling_attempts}")
                yield {"status": "error", "message": "获取结果超时，超过最大轮询次数"}
                return
            
            if not imagine_message:
                logger.error("轮询结束但没有获取到有效结果")
                yield {"status": "error", "message": "没有获取到有效结果"}
                return

            logger.info(f"成功获取图像结果，消息ID: {imagine_message.get('id')}")
            yield {
                "status": "success", 
                "message_id": imagine_message.get('id'),
                "content": imagine_message.get('content', ''),
                "images": self.extract_image_urls(imagine_message)
            }
            
        except Exception as e:
            logger.error(f"发送Imagine请求失败: {str(e)}")
            yield {"status": "error", "message": f"发送Imagine请求失败: {str(e)}"}

    async def get_imagine(self, client, channel_id, prompt, count=0, seed=None):
        """获取生成图像的消息"""
        try:
            response = client.get(f'{self.API_URL}/channels/{channel_id}/messages?limit=10')
            data = response.json()
            
            def criteria(item):
                content = item.get('content', '')
                
                if seed is not None and f"--seed {seed}" in content:
                    progress_match = re.search(r'(?:\(|（)?(\d+)%(?:\)|）)?', content)
                    if progress_match:
                        progress_value = int(progress_match.group(1))
                        logger.info(f"任务进度: {progress_value}%")
                        if "%" in content:
                            return False
                
                if "%" in content:
                    return False
                    
                if seed is not None:
                    seed_pattern = f"--seed {seed}"
                    if seed_pattern not in content:
                        return False
                    else:
                        if 'attachments' in item and len(item.get('attachments', [])) > 0:
                            return True
                return False

            return self.first_where(data, criteria)
            
        except Exception as e:
            logger.error(f"获取通道消息失败: {str(e)}")
            return None

    async def get_progress_message(self, client, channel_id, prompt, seed=None):
        """获取包含进度信息的消息"""
        try:
            response = client.get(f'{self.API_URL}/channels/{channel_id}/messages?limit=10')
            data = response.json()
            
            for item in data:
                content = item.get('content', '')
                if seed is not None and f"--seed {seed}" in content and "%" in content:
                    progress_match = re.search(r'(?:\(|（)?(\d+)%(?:\)|）)?', content)
                    if progress_match:
                        return item
            
            return None
            
        except Exception as e:
            logger.error(f"获取进度消息失败: {str(e)}")
            return None

    def extract_image_urls(self, message):
        """从消息中提取图像URL"""
        image_urls = []
        if 'attachments' in message and message['attachments']:
            for attachment in message['attachments']:
                if 'url' in attachment:
                    image_urls.append(attachment['url'])
        return image_urls

    async def split_image(self, image_url):
        """将一张大图切割成四张子图，保存到本地并返回URLs"""
        try:
            response = requests.get(image_url, proxies=self.proxies if self.proxies else None, timeout=30)
            if response.status_code != 200:
                logger.error(f"下载图像失败，状态码: {response.status_code}")
                return None
                
            image_data = response.content
            img = Image.open(io.BytesIO(image_data))
            width, height = img.size
            
            original_format = img.format
            if not original_format:
                parsed_url = urlparse(image_url)
                original_format = os.path.splitext(parsed_url.path)[1][1:].upper()
                if not original_format:
                    original_format = 'PNG'
            
            if width < 500 or height < 500:
                logger.error(f"图像尺寸不符合预期: {width}x{height}, 应该接近2048x2048")
                return None
                
            half_width = width // 2
            half_height = height // 2
            
            quadrants = [
                img.crop((0, 0, half_width, half_height)),
                img.crop((half_width, 0, width, half_height)),
                img.crop((0, half_height, half_width, height)),
                img.crop((half_width, half_height, width, height))
            ]
            
            image_id = uuid.uuid4().hex[:10]
            save_dir = os.path.abspath(settings.save_dir)
            os.makedirs(save_dir, exist_ok=True)
            
            image_urls = []
            for i, quadrant in enumerate(quadrants, 1):
                try:
                    filename = f"split_{image_id}_{i}.{original_format.lower()}"
                    file_path = os.path.join(save_dir, filename)
                    
                    save_params = {"format": original_format}
                    if original_format in ['PNG', 'JPEG', 'JPG']:
                        save_params["optimize"] = True
                    if original_format in ['JPEG', 'JPG']:
                        save_params["quality"] = 95
                    
                    quadrant.save(file_path, **save_params)
                    
                    if not os.path.exists(file_path):
                        raise Exception(f"文件保存失败: {file_path}")
                        
                    file_size = os.path.getsize(file_path)
                    if file_size == 0:
                        raise Exception(f"保存的文件大小为0: {file_path}")
                        
                    image_url = f"{settings.download_url}/{filename}"
                    image_urls.append(image_url)
                    
                    logger.info(f"成功保存分割图片 {i}/4: {filename}")
                except Exception as e:
                    logger.error(f"保存分割图片 {i}/4 失败: {str(e)}")
                    for url in image_urls:
                        try:
                            file_path = os.path.join(save_dir, os.path.basename(url))
                            if os.path.exists(file_path):
                                os.remove(file_path)
                        except Exception as del_e:
                            logger.error(f"删除失败的图片文件时出错: {str(del_e)}")
                    return None
                
            if len(image_urls) != 4:
                logger.error(f"分割图片数量不正确: 期望4张，实际{len(image_urls)}张")
                return None
                
            logger.info("成功完成图片分割，生成4张子图")
            return image_urls
            
        except Exception as e:
            logger.error(f"分割图像失败: {str(e)}")
            return None

    async def split_images(self, image_urls: List[str], config: Optional[Dict] = None):
        """批量处理多个图像URL"""
        if not config:
            config = {}
            
        total = len(image_urls)
        success_count = 0
        
        for i, image_url in enumerate(image_urls, 1):
            try:
                if not is_valid_image_url(image_url):
                    yield {
                        "status": "error",
                        "index": i,
                        "total": total,
                        "success_count": success_count,
                        "message": "无效的图片URL"
                    }
                    continue
                
                split_urls = await self.split_image(image_url)
                if split_urls and len(split_urls) == 4:
                    success_count += 1
                    yield {
                        "status": "success",
                        "index": i,
                        "total": total,
                        "success_count": success_count,
                        "images": split_urls
                    }
                else:
                    yield {
                        "status": "error",
                        "index": i,
                        "total": total,
                        "success_count": success_count,
                        "message": "分割图片失败"
                    }
            except Exception as e:
                yield {
                    "status": "error",
                    "index": i,
                    "total": total,
                    "success_count": success_count,
                    "message": str(e)
                }

    async def generate_image(self, prompt, config=None):
        """生成图像并以流式方式返回结果"""
        if not config:
            config = {}
            
        oauth_token = config.get('oauth_token', settings.midjourney_oauth_token)
        channel_id = config.get('channel_id', settings.midjourney_channel_id)
        
        if not oauth_token or not channel_id:
            yield {
                "status": "error",
                "message": "缺少Discord配置",
                "success_count": 0
            }
            return

        try:
            client, guild_id, user_id = await self.initialize_client(oauth_token, channel_id)
        except Exception as e:
            yield {
                "status": "error",
                "message": str(e),
                "success_count": 0
            }
            return

        prompt = prompt.strip()
        options = settings.midjourney_default_options.copy()
        if 'options' in config:
            options.update(config.get('options', {}))
            
        if "seed" not in options or not options.get('seed'):
            options['seed'] = random.randint(0, 4294967295)
            
        if 'image_weight' in config and isinstance(config['image_weight'], (int, float)):
            iw = max(0.1, min(config['image_weight'], 3))
            options['iw'] = iw
            
        seed = options.pop('seed', None)
        
        parameter = ""
        no_value_key = ['relax', 'fast', 'turbo', 'tile']
        for key, value in options.items():
            if key in no_value_key:
                parameter += f" --{key}"
            else:
                parameter += f" --{key} {value}"
        
        if seed:
            parameter += f" --seed {seed}"
            logger.info(f"使用的seed值: {seed}")
            
        if 'image_urls' in config and config['image_urls']:
            image_urls = config['image_urls']
            if not isinstance(image_urls, list):
                image_urls = [image_urls]
                
            new_image_urls = []
            for image_url in image_urls:
                if is_valid_image_url(image_url):
                    try:
                        new_url = get_new_image_url(image_url)
                        if new_url:
                            new_image_urls.append(new_url)
                    except Exception as e:
                        logger.warning(f"转换图片URL失败: {str(e)}")
                        
            if new_image_urls:
                prompt = " ".join(new_image_urls) + " " + prompt
            
        if 'characters' in config and config['characters']:
            char_urls = []
            characters = config['characters'] if isinstance(config['characters'], list) else [config['characters']]
            
            for char_url in characters:
                if is_valid_image_url(char_url):
                    try:
                        new_url = get_new_image_url(char_url)
                        if new_url:
                            char_urls.append(new_url)
                    except Exception as e:
                        logger.warning(f"转换角色图片URL失败: {str(e)}")
            
            if char_urls:
                prompt = prompt + " --cref " + " ".join(char_urls)
            
        if 'styles' in config and config['styles']:
            style_urls = []
            styles = config['styles'] if isinstance(config['styles'], list) else [config['styles']]
            
            for style_url in styles:
                if is_valid_image_url(style_url):
                    try:
                        new_url = get_new_image_url(style_url)
                        if new_url:
                            style_urls.append(new_url)
                    except Exception as e:
                        logger.warning(f"转换风格图片URL失败: {str(e)}")
            
            if style_urls:
                prompt = prompt + " --sref " + " ".join(style_urls)
            
        prompt = prompt + parameter
        success_count = 0
        
        async for result in self.imagine(client, guild_id, channel_id, prompt, seed):
            if result.get("status") == "progress":
                yield {
                    "status": "progress",
                    "progress": result.get("progress", 0),
                    "seed": seed
                }
            elif result.get("status") == "success":
                success_count += 1
                response = {
                    "status": "success",
                    "success_count": success_count,
                    "image_urls": result.get("images", [])
                }
                
                if config.get("split_image", True) and result.get("images"):
                    try:
                        orig_image_url = result["images"][0]
                        split_urls = await self.split_image(orig_image_url)
                        if split_urls:
                            response["image_urls"] = split_urls
                    except Exception as e:
                        logger.error(f"分割图像失败: {str(e)}")
                        response["status"] = "error"
                        response["message"] = f"分割图像失败: {str(e)}"
                
                yield response
            else:
                yield {
                    "status": "error",
                    "message": result.get("message", "未知错误"),
                    "success_count": success_count
                }