#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
海报API通用模型定义
支持多种模板类型的海报生成
"""

from typing import List, Dict, Any, Optional, Union
from pydantic import BaseModel, Field


# 基础模板信息
class TemplateInfo(BaseModel):
    """模板信息模型"""
    id: str = Field(..., description="模板ID")
    name: str = Field(..., description="模板名称")
    description: str = Field(..., description="模板描述")
    size: List[int] = Field(..., description="模板尺寸 [宽, 高]")
    required_fields: List[str] = Field(..., description="必填字段列表")
    optional_fields: List[str] = Field(default=[], description="可选字段列表")
    
    class Config:
        schema_extra = {
            "example": {
                "id": "vibrant",
                "name": "Vibrant活力模板",
                "description": "适合旅游景点、娱乐活动的活力海报模板",
                "size": [900, 1200],
                "required_fields": ["title", "slogan", "price", "ticket_type", "content_items"],
                "optional_fields": ["remarks", "tag", "pagination"]
            }
        }


# 通用内容模型（支持任意字段）
class PosterContent(BaseModel):
    """通用海报内容模型，支持动态字段"""
    
    class Config:
        extra = "allow"  # 允许额外字段
        schema_extra = {
            "example": {
                "title": "海洋奇幻世界",
                "slogan": "探索深海秘境，感受蓝色奇迹的无限魅力",
                "price": "299",
                "ticket_type": "成人票",
                "content_items": [
                    "海洋馆门票1张（含所有展区）",
                    "海豚表演VIP座位"
                ]
            }
        }


# 保持向后兼容的Vibrant内容模型
class VibrantPosterContent(PosterContent):
    """Vibrant海报内容模型（向后兼容）"""
    title: Optional[str] = Field(None, description="主标题（8-12字符）")
    slogan: Optional[str] = Field(None, description="副标题/宣传语（10-20字符）")
    price: Optional[str] = Field(None, description="价格，如果没有价格则用'欢迎咨询'")
    ticket_type: Optional[str] = Field(None, description="票种类型（如成人票、套餐票等）")
    content_button: Optional[str] = Field(None, description="内容按钮文字")
    content_items: Optional[List[str]] = Field(None, description="套餐内容列表（3-5个项目）")
    remarks: Optional[List[str]] = Field(None, description="备注信息（1-3条）")
    tag: Optional[str] = Field(None, description="标签")
    pagination: Optional[str] = Field(None, description="分页信息")


# 通用海报生成请求
class PosterGenerationRequest(BaseModel):
    """通用海报生成请求模型"""
    # 模板相关
    template_id: str = Field(..., description="模板ID")
    
    # 内容相关（二选一）
    content: Optional[Dict[str, Any]] = Field(None, description="直接提供的海报内容")
    source_data: Optional[Dict[str, Any]] = Field(None, description="源数据，用于AI生成内容")
    
    # 生成参数
    topic_name: Optional[str] = Field(None, description="主题名称，用于文件命名")
    image_path: Optional[str] = Field(None, description="指定图片路径，如果不提供则随机选择")
    image_dir: Optional[str] = Field(None, description="图片目录，如果不提供使用默认目录")
    output_dir: Optional[str] = Field(None, description="输出目录，如果不提供使用默认目录")
    
    # AI参数
    temperature: Optional[float] = Field(default=0.7, description="AI生成温度参数")
    
    class Config:
        schema_extra = {
            "example": {
                "template_id": "vibrant",
                "source_data": {
                    "scenic_info": {
                        "name": "天津冒险湾",
                        "location": "天津市",
                        "type": "水上乐园"
                    },
                    "product_info": {
                        "name": "冒险湾门票",
                        "price": 299,
                        "type": "成人票"
                    },
                    "tweet_info": {
                        "title": "夏日清凉首选",
                        "content": "天津冒险湾水上乐园，多种刺激项目等你来挑战..."
                    }
                },
                "topic_name": "天津冒险湾",
                "temperature": 0.7
            }
        }


# 内容生成请求
class ContentGenerationRequest(BaseModel):
    """内容生成请求模型"""
    template_id: str = Field(..., description="模板ID")
    source_data: Dict[str, Any] = Field(..., description="源数据，用于AI生成内容")
    temperature: Optional[float] = Field(default=0.7, description="AI生成温度参数")
    
    class Config:
        schema_extra = {
            "example": {
                "template_id": "vibrant",
                "source_data": {
                    "scenic_info": {"name": "天津冒险湾", "type": "水上乐园"},
                    "product_info": {"name": "门票", "price": 299},
                    "tweet_info": {"title": "夏日清凉", "content": "水上乐园体验"}
                },
                "temperature": 0.7
            }
        }


# 保持向后兼容
class VibrantPosterRequest(PosterGenerationRequest):
    """Vibrant海报生成请求模型（向后兼容）"""
    template_id: str = Field(default="vibrant", description="模板ID，默认为vibrant")
    
    # 兼容旧的字段结构
    scenic_info: Optional[Dict[str, Any]] = Field(None, description="景区信息")
    product_info: Optional[Dict[str, Any]] = Field(None, description="产品信息")
    tweet_info: Optional[Dict[str, Any]] = Field(None, description="推文信息")
    
    def __init__(self, **data):
        # 转换旧格式到新格式
        if 'scenic_info' in data or 'product_info' in data or 'tweet_info' in data:
            source_data = {}
            for key in ['scenic_info', 'product_info', 'tweet_info']:
                if key in data and data[key] is not None:
                    source_data[key] = data.pop(key)
            if source_data and 'source_data' not in data:
                data['source_data'] = source_data
        
        super().__init__(**data)


# 标准API响应基础类
class BaseAPIResponse(BaseModel):
    """API响应基础模型"""
    success: bool = Field(..., description="操作是否成功")
    message: str = Field(default="", description="响应消息")
    request_id: str = Field(..., description="请求ID")
    timestamp: str = Field(..., description="响应时间戳")


# 通用海报生成响应
class PosterGenerationResponse(BaseAPIResponse):
    """通用海报生成响应模型"""
    data: Optional[Dict[str, Any]] = Field(None, description="响应数据")
    
    class Config:
        schema_extra = {
            "example": {
                "success": True,
                "message": "海报生成成功",
                "request_id": "poster-20240715-123456-a1b2c3d4",
                "timestamp": "2024-07-15T12:34:56Z",
                "data": {
                    "template_id": "vibrant",
                    "topic_name": "天津冒险湾",
                    "poster_path": "/result/posters/vibrant_海洋奇幻世界_20240715_123456.png",
                    "content": {
                        "title": "海洋奇幻世界",
                        "slogan": "探索深海秘境，感受蓝色奇迹的无限魅力",
                        "price": "299"
                    },
                    "metadata": {
                        "image_used": "/data/images/ocean_park.jpg",
                        "generation_method": "ai_generated",
                        "template_size": [900, 1200],
                        "processing_time": 3.2
                    }
                }
            }
        }


# 内容生成响应
class ContentGenerationResponse(BaseAPIResponse):
    """内容生成响应模型"""
    data: Optional[Dict[str, Any]] = Field(None, description="生成的内容")
    
    class Config:
        schema_extra = {
            "example": {
                "success": True,
                "message": "内容生成成功",
                "request_id": "content-20240715-123456-a1b2c3d4",
                "timestamp": "2024-07-15T12:34:56Z",
                "data": {
                    "template_id": "vibrant",
                    "content": {
                        "title": "冒险湾水世界",
                        "slogan": "夏日激情体验，清凉无限乐趣",
                        "price": "299",
                        "ticket_type": "成人票"
                    },
                    "metadata": {
                        "generation_method": "ai_generated",
                        "temperature": 0.7
                    }
                }
            }
        }


# 模板列表响应
class TemplateListResponse(BaseAPIResponse):
    """模板列表响应模型"""
    data: Optional[List[TemplateInfo]] = Field(None, description="模板列表")
    
    class Config:
        schema_extra = {
            "example": {
                "success": True,
                "message": "获取模板列表成功",
                "request_id": "templates-20240715-123456",
                "timestamp": "2024-07-15T12:34:56Z",
                "data": [
                    {
                        "id": "vibrant",
                        "name": "Vibrant活力模板",
                        "description": "适合旅游景点、娱乐活动的活力海报模板",
                        "size": [900, 1200],
                        "required_fields": ["title", "slogan", "price"],
                        "optional_fields": ["remarks", "tag"]
                    }
                ]
            }
        }


# 保持向后兼容的响应模型
class VibrantPosterResponse(BaseModel):
    """Vibrant海报生成响应模型（向后兼容）"""
    request_id: str = Field(..., description="请求ID")
    topic_name: str = Field(..., description="主题名称")
    poster_path: str = Field(..., description="生成的海报文件路径")
    generated_content: Dict[str, Any] = Field(..., description="生成或使用的海报内容")
    image_used: str = Field(..., description="使用的图片路径")
    generation_method: str = Field(..., description="生成方式：direct（直接使用提供内容）或ai_generated（AI生成）")
    
    class Config:
        schema_extra = {
            "example": {
                "request_id": "vibrant-20240715-123456-a1b2c3d4",
                "topic_name": "天津冒险湾",
                "poster_path": "/result/posters/vibrant_海洋奇幻世界_20240715_123456.png",
                "generated_content": {
                    "title": "海洋奇幻世界",
                    "slogan": "探索深海秘境，感受蓝色奇迹的无限魅力",
                    "price": "299",
                    "ticket_type": "成人票",
                    "content_items": ["海洋馆门票1张", "海豚表演VIP座位"]
                },
                "image_used": "/data/images/ocean_park.jpg",
                "generation_method": "ai_generated"
            }
        }


class BatchVibrantPosterRequest(BaseModel):
    """批量Vibrant海报生成请求模型"""
    base_path: str = Field(..., description="包含多个topic目录的基础路径")
    image_dir: Optional[str] = Field(None, description="图片目录")
    scenic_info_file: Optional[str] = Field(None, description="景区信息文件路径")
    product_info_file: Optional[str] = Field(None, description="产品信息文件路径")
    output_base: Optional[str] = Field(default="result/posters", description="输出基础目录")
    parallel_count: Optional[int] = Field(default=3, description="并发处理数量")
    temperature: Optional[float] = Field(default=0.7, description="AI生成温度参数")
    
    class Config:
        schema_extra = {
            "example": {
                "base_path": "/root/TravelContentCreator/result/run_20250710_165327",
                "image_dir": "/root/TravelContentCreator/data/images",
                "scenic_info_file": "/root/TravelContentCreator/resource/data/Object/天津冒险湾.txt",
                "product_info_file": "/root/TravelContentCreator/resource/data/Product/product.bak",
                "output_base": "result/posters",
                "parallel_count": 3,
                "temperature": 0.7
            }
        }


class BatchVibrantPosterResponse(BaseModel):
    """批量Vibrant海报生成响应模型"""
    request_id: str = Field(..., description="批量处理请求ID")
    total_topics: int = Field(..., description="总共处理的topic数量")
    successful_count: int = Field(..., description="成功生成的海报数量")
    failed_count: int = Field(..., description="失败的海报数量")
    output_base_dir: str = Field(..., description="输出基础目录")
    successful_topics: List[str] = Field(..., description="成功处理的topic列表")
    failed_topics: List[Dict[str, str]] = Field(..., description="失败的topic及错误信息")
    
    class Config:
        schema_extra = {
            "example": {
                "request_id": "batch-vibrant-20240715-123456",
                "total_topics": 5,
                "successful_count": 4,
                "failed_count": 1,
                "output_base_dir": "/result/posters/run_20250710_165327",
                "successful_topics": ["topic_1", "topic_2", "topic_3", "topic_4"],
                "failed_topics": [
                    {"topic": "topic_5", "error": "图片文件不存在"}
                ]
            }
        } 