# tongyi_image_generator.py
import dashscope
import requests
import base64
from PIL import Image
from io import BytesIO
import os

class AIImageGenerator:
    def __init__(self, api_key):
        """
        初始化图片生成器
        :param api_key: 阿里云DashScope API密钥
        """
        self.api_key = api_key
        dashscope.api_key = api_key
        
    def generate_poster_image(self, prompt, size="1024*1024", n=1):
        """
        使用通义万相生成海报图片
        :param prompt: 图片描述提示词
        :param size: 图片尺寸，可选：512*512, 1024*1024, 1024*768, 768*1024等
        :param n: 生成图片数量
        :return: 图片URL列表或None
        """
        try:
            # 使用dashscope的image_synthesis模块
            from dashscope import ImageSynthesis
            
            # 构建增强的提示词
            enhanced_prompt = f"高质量海报设计，{prompt}，专业平面设计，色彩鲜艳，适合印刷"
            negative_prompt = "模糊，低质量，文字错误，变形，构图混乱"
            
            # 调用通义万相API
            response = ImageSynthesis.call(
                model="wanx-v1",
                prompt=enhanced_prompt,
                negative_prompt=negative_prompt,
                size=size,
                n=n
            )
            
            # 检查响应状态
            if response.status_code == 200:
                results = response.output.results
                image_urls = [result.url for result in results]
                print(f"✓ 成功生成 {len(image_urls)} 张图片")
                return image_urls
            else:
                print(f"✗ 生成失败: {response.code} - {response.message}")
                return None
                
        except Exception as e:
            print(f"✗ 生成过程中发生异常: {e}")
            return None
    
    def generate_poster_image_http(self, prompt, size="1024*1024"):
        """
        使用HTTP请求方式生成图片（备用方案）
        :param prompt: 图片描述提示词
        :param size: 图片尺寸
        :return: 图片URL或None
        """
        url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis"
        
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        enhanced_prompt = f"高质量海报设计，{prompt}，专业平面设计，色彩鲜艳，适合印刷"
        
        payload = {
            "model": "wanx-v1",
            "input": {
                "prompt": enhanced_prompt,
                "negative_prompt": "模糊，低质量，文字错误，变形，构图混乱"
            },
            "parameters": {
                "size": size,
                "n": 1
            }
        }
        
        try:
            response = requests.post(url, headers=headers, json=payload, timeout=30)
            result = response.json()
            
            # 检查响应
            if response.status_code == 200 and "output" in result:
                if "results" in result["output"] and len(result["output"]["results"]) > 0:
                    image_url = result["output"]["results"][0]["url"]
                    print(f"✓ 成功生成图片")
                    return image_url
                else:
                    print(f"✗ 响应中没有图片信息")
                    return None
            else:
                print(f"✗ API调用失败: {result}")
                return None
                
        except requests.exceptions.RequestException as e:
            print(f"✗ 网络请求异常: {e}")
            return None
        except Exception as e:
            print(f"✗ 其他异常: {e}")
            return None
    
    def save_image_from_url(self, image_url, filename, save_dir="generated_images"):
        """
        从URL保存图片
        :param image_url: 图片URL
        :param filename: 文件名
        :param save_dir: 保存目录
        """
        try:
            # 创建保存目录
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            
            # 完整文件路径
            filepath = os.path.join(save_dir, filename)
            
            # 下载图片
            response = requests.get(image_url, timeout=30)
            response.raise_for_status()
            
            # 打开并保存图片
            img = Image.open(BytesIO(response.content))
            img.save(filepath)
            print(f"✓ 图片已保存: {filepath}")
            return filepath
            
        except requests.exceptions.RequestException as e:
            print(f"✗ 下载图片失败: {e}")
            return None
        except Exception as e:
            print(f"✗ 保存图片失败: {e}")
            return None
    
    def save_image_from_base64(self, base64_data, filename, save_dir="generated_images"):
        """
        从base64数据保存图片
        :param base64_data: base64编码的图片数据
        :param filename: 文件名
        :param save_dir: 保存目录
        """
        try:
            # 创建保存目录
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            
            # 完整文件路径
            filepath = os.path.join(save_dir, filename)
            
            # 解码base64并保存
            image_data = base64.b64decode(base64_data)
            img = Image.open(BytesIO(image_data))
            img.save(filepath)
            print(f"✓ 图片已保存: {filepath}")
            return filepath
            
        except Exception as e:
            print(f"✗ 保存base64图片失败: {e}")
            return None


# 海报生成示例
def create_event_poster():
    """创建民族团结活动海报"""
    
    # 初始化生成器（请替换为您的实际API Key）
    api_key = "your-deepseek-api-key-here"
    generator = AIImageGenerator(api_key)
    
    # 定义多个海报风格的提示词
    poster_prompts = [
        {
            "name": "民族团结活动海报_主视觉",
            "prompt": """
                新疆风格校园活动海报，民族团结主题，
                各族大学生手拉手跳舞，穿着中华民族传统服饰，
                背景有天山山脉和现代建筑元素，
                色彩温暖明亮，有杏花飘落的场景，
                现代插画风格，画面充满活力和欢乐氛围
            """
        },
        {
            "name": "民族团结活动海报_文字版",
            "prompt": """
                民族团结主题海报，包含"民族团结一家亲"艺术字体设计，
                各族学生一起学习交流的场景，
                背景有现代特色建筑和自然风光，
                色彩和谐温馨，具有浓郁的中华民族特色
            """
        }
    ]
    
    # 生成并保存海报
    for poster in poster_prompts:
        print(f"\n正在生成: {poster['name']}")
        
        # 方法1：使用dashscope SDK（推荐）
        image_urls = generator.generate_poster_image(poster['prompt'], size="1024*1024")
        
        if image_urls:
            # 保存第一张生成的图片
            generator.save_image_from_url(
                image_urls[0], 
                f"{poster['name']}.png"
            )
        else:
            print(f"✗ 使用SDK生成失败，尝试HTTP方式...")
            
            # 方法2：使用HTTP请求（备用）
            image_url = generator.generate_poster_image_http(poster['prompt'])
            if image_url:
                generator.save_image_from_url(
                    image_url, 
                    f"{poster['name']}_http.png"
                )
            else:
                print(f"✗ 两种方式均生成失败")


def create_custom_poster():
    """创建自定义海报"""
    
    api_key = "your-tongyi-api-key-here"
    generator = AIImageGenerator(api_key)
    
    # 自定义提示词
    custom_prompt = input("请输入海报描述（按Enter使用默认描述）: ").strip()
    if not custom_prompt:
        custom_prompt = """
            现代教育主题海报，学生们在科技感十足的教室里学习，
            AI辅助教学场景，未来科技风格，
            蓝色和白色为主色调，明亮清新
        """
    
    print(f"\n生成提示词: {custom_prompt}")
    
    # 生成图片
    image_urls = generator.generate_poster_image(custom_prompt, size="1024*1024")
    
    if image_urls:
        # 保存图片
        filename = "custom_poster.png"
        generator.save_image_from_url(image_urls[0], filename)
        print(f"\n✓ 海报生成完成！")
        print(f"  保存位置: generated_images/{filename}")
    else:
        print("\n✗ 海报生成失败，请检查API密钥和网络连接")


if __name__ == "__main__":
    print("=" * 50)
    print("通义万相海报生成器")
    print("=" * 50)
    
    while True:
        print("\n请选择功能:")
        print("1. 生成民族团结活动海报")
        print("2. 生成自定义海报")
        print("3. 退出")
        
        choice = input("请输入选择 (1/2/3): ").strip()
        
        if choice == "1":
            create_event_poster()
        elif choice == "2":
            create_custom_poster()
        elif choice == "3":
            print("感谢使用，再见！")
            break
        else:
            print("无效选择，请重新输入")

运行结果：

(ai_env) $ python3 tongyi_image_generator.py
==================================================
通义万相海报生成器
==================================================

请选择功能:
1. 生成民族团结活动海报
2. 生成自定义海报
3. 退出
请输入选择 (1/2/3): 1

正在生成: 民族团结活动海报_主视觉
✓ 成功生成 1 张图片
✓ 图片已保存: generated_images/民族团结活动海报_主视觉.png

正在生成: 民族团结活动海报_文字版
✓ 成功生成 1 张图片
✓ 图片已保存: generated_images/民族团结活动海报_文字版.png

请选择功能:
1. 生成民族团结活动海报
2. 生成自定义海报
3. 退出
请输入选择 (1/2/3): 3
感谢使用，再见！
(ai_env) $