#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
通义万相文生图V2版API核心功能模块

本模块提供通义万相文生图V2版API的核心功能实现，包括参数校验、API调用等。
可被MCP服务器和客户端模块共同使用，避免代码重复。

作者: John Chen
"""

import os
from typing import List, Optional, Dict, Any
from dashscope import ImageSynthesis


class BailianImageCore:
    """
    通义万相文生图V2版API核心功能类

    提供文生图的核心功能实现，包括参数校验、API调用等。
    """

    # 支持的模型列表
    SUPPORTED_MODELS = [
        "wanx-v1",
        "wanx-lite-v1",
        "wanx-sketch-v1",
        "wanx-background-generation-v2",
    ]

    # 支持的图像尺寸
    SUPPORTED_SIZES = [
        "1024*1024",
        "720*1280",
        "1280*720",
        "768*1344",
        "1344*768",
        "1536*1024",
        "1024*1536",
    ]

    def __init__(self, api_key: Optional[str] = None):
        """
        初始化核心功能类

        Args:
            api_key: DashScope API密钥，如果不提供则从环境变量DASHSCOPE_API_KEY读取
        """
        self.api_key = api_key or os.getenv("DASHSCOPE_API_KEY")
        if not self.api_key:
            raise ValueError("DASHSCOPE_API_KEY未设置，请设置环境变量或传入api_key参数")

    def validate_model(self, model: str) -> None:
        """
        校验模型参数

        Args:
            model: 模型名称

        Raises:
            ValueError: 当模型不支持时抛出异常
        """
        if model not in self.SUPPORTED_MODELS:
            raise ValueError(
                f"不支持的模型: {model}，支持的模型: {', '.join(self.SUPPORTED_MODELS)}"
            )

    def validate_size(self, size: str) -> None:
        """
        校验图像尺寸参数

        Args:
            size: 图像尺寸

        Raises:
            ValueError: 当尺寸不支持时抛出异常
        """
        if size not in self.SUPPORTED_SIZES:
            raise ValueError(
                f"不支持的尺寸: {size}，支持的尺寸: {', '.join(self.SUPPORTED_SIZES)}"
            )

    def validate_n(self, n: int) -> None:
        """
        校验生成数量参数

        Args:
            n: 生成图像数量

        Raises:
            ValueError: 当数量不在有效范围时抛出异常
        """
        if not isinstance(n, int) or n < 1 or n > 4:
            raise ValueError("生成数量必须是1-4之间的整数")

    def text2imagev2(
        self,
        prompt: str,
        negative_prompt: Optional[str] = None,
        model: str = "wanx-v1",
        size: str = "1024*1024",
        n: int = 1,
        seed: Optional[int] = None,
        ref_img: Optional[str] = None,
        ref_strength: Optional[float] = None,
        ref_mode: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        通义万相文生图V2版API调用

        支持文生图、支持正向和反向提示词、多种模型选择、自定义图像尺寸和生成数量。

        Args:
            prompt: 正向提示词，描述希望生成的图像内容
            negative_prompt: 反向提示词，描述不希望出现在图像中的内容（可选）
            model: 使用的模型，默认为'wanx-v1'
            size: 生成图像的尺寸，默认为'1024*1024'
            n: 生成图像的数量，范围1-4，默认为1
            seed: 随机种子，用于生成可重现的结果（可选）
            ref_img: 参考图像URL，用于图像参考生成（可选）
            ref_strength: 参考强度，范围0.1-1.0（可选）
            ref_mode: 参考模式，可选值：'repaint'（可选）

        Returns:
            Dict[str, Any]: 包含生成结果的字典，包含以下字段：
                - success: 是否成功
                - images: 生成的图像URL列表（成功时）
                - error: 错误信息（失败时）
                - request_id: 请求ID

        Raises:
            ValueError: 当参数不合法时抛出异常
        """
        # 参数校验
        self.validate_model(model)
        self.validate_size(size)
        self.validate_n(n)

        if ref_strength is not None and (ref_strength < 0.1 or ref_strength > 1.0):
            raise ValueError("参考强度必须在0.1-1.0之间")

        if ref_mode is not None and ref_mode not in ["repaint"]:
            raise ValueError("参考模式只支持'repaint'")

        try:
            # 构建请求参数
            params = {
                "model": model,
                "input": {"prompt": prompt},
                "parameters": {"size": size, "n": n},
            }

            # 添加可选参数
            if negative_prompt:
                params["input"]["negative_prompt"] = negative_prompt

            if seed is not None:
                params["parameters"]["seed"] = seed

            if ref_img:
                params["input"]["ref_img"] = ref_img
                if ref_strength is not None:
                    params["parameters"]["ref_strength"] = ref_strength
                if ref_mode is not None:
                    params["parameters"]["ref_mode"] = ref_mode

            # 调用API
            response = ImageSynthesis.call(**params)

            if response.status_code == 200:
                # 提取图像URL
                images = []
                if hasattr(response.output, "results"):
                    for result in response.output.results:
                        if hasattr(result, "url"):
                            images.append(result.url)

                return {
                    "success": True,
                    "images": images,
                    "request_id": getattr(response, "request_id", None),
                }
            else:
                return {
                    "success": False,
                    "error": f"API调用失败，状态码: {response.status_code}，错误信息: {response.message}",
                    "request_id": getattr(response, "request_id", None),
                }

        except Exception as e:
            return {
                "success": False,
                "error": f"调用过程中发生异常: {str(e)}",
                "request_id": None,
            }
