#!/usr/bin/env python3
"""
提示词验证脚本

用于验证图像和视频生成的提示词质量。
"""

import re
import sys
from typing import Dict, List, Tuple


class PromptValidator:
    """提示词验证器"""
    
    # 质量关键词
    QUALITY_KEYWORDS = [
        "high quality", "best quality", "masterpiece",
        "detailed", "8K", "4K", "HD", "professional",
        "高质量", "精细", "专业"
    ]
    
    # 风格关键词
    STYLE_KEYWORDS = [
        "anime", "realistic", "cinematic", "cartoon",
        "动漫", "写实", "电影", "卡通"
    ]
    
    # 视频运动关键词
    MOTION_KEYWORDS = [
        "smooth", "natural", "fluid", "motion",
        "流畅", "自然", "运动", "动作"
    ]
    
    def __init__(self):
        self.errors = []
        self.warnings = []
        self.suggestions = []
    
    def validate_image_prompt(self, prompt: str) -> Dict[str, any]:
        """验证图像生成提示词"""
        self.errors = []
        self.warnings = []
        self.suggestions = []
        
        # 检查长度
        if len(prompt) < 10:
            self.errors.append("提示词过短（少于 10 个字符）")
        elif len(prompt) < 50:
            self.warnings.append("提示词较短，建议添加更多细节描述")
        
        if len(prompt) > 500:
            self.warnings.append("提示词过长，可能影响生成效果")
        
        # 检查是否包含质量关键词
        has_quality = any(kw in prompt.lower() for kw in self.QUALITY_KEYWORDS)
        if not has_quality:
            self.suggestions.append(
                "建议添加质量关键词，如：high quality, detailed, 高质量, 精细细节"
            )
        
        # 检查是否包含风格关键词
        has_style = any(kw in prompt.lower() for kw in self.STYLE_KEYWORDS)
        if not has_style:
            self.suggestions.append(
                "建议添加风格关键词，如：anime style, realistic, 动漫风格, 写实风格"
            )
        
        # 检查是否有主体描述
        if not self._has_subject(prompt):
            self.errors.append("缺少明确的主体描述（人物、场景或物品）")
        
        return self._build_result()
    
    def validate_video_prompt(self, prompt: str) -> Dict[str, any]:
        """验证视频生成提示词"""
        self.errors = []
        self.warnings = []
        self.suggestions = []
        
        # 检查长度
        if len(prompt) < 20:
            self.errors.append("提示词过短（少于 20 个字符）")
        elif len(prompt) < 50:
            self.warnings.append("提示词较短，建议添加更多动作和镜头描述")
        
        if len(prompt) > 800:
            self.warnings.append("提示词过长，可能影响生成效果")
        
        # 检查是否包含运动关键词
        has_motion = any(kw in prompt.lower() for kw in self.MOTION_KEYWORDS)
        if not has_motion:
            self.suggestions.append(
                "建议添加运动关键词，如：smooth motion, natural movement, 流畅运动, 自然过渡"
            )
        
        # 检查是否有动作描述
        action_words = ["动", "移动", "转", "走", "跑", "飞", "摇", "摆",
                       "move", "turn", "walk", "run", "fly", "swing"]
        has_action = any(word in prompt.lower() for word in action_words)
        if not has_action:
            self.warnings.append("缺少明确的动作描述")
        
        # 检查是否有镜头描述
        camera_words = ["镜头", "camera", "shot", "angle", "推", "拉", "摇", "移"]
        has_camera = any(word in prompt.lower() for word in camera_words)
        if not has_camera:
            self.suggestions.append(
                "建议添加镜头描述，如：镜头推进, 镜头跟随, camera follows"
            )
        
        return self._build_result()
    
    def _has_subject(self, prompt: str) -> bool:
        """检查是否有主体描述"""
        # 简单检查：是否包含名词性词汇
        subject_indicators = [
            "人", "女", "男", "孩", "角色", "场景", "森林", "城市", "房间",
            "物", "杯", "书", "车",
            "person", "woman", "man", "child", "character",
            "scene", "forest", "city", "room", "object"
        ]
        return any(word in prompt.lower() for word in subject_indicators)
    
    def _build_result(self) -> Dict[str, any]:
        """构建验证结果"""
        is_valid = len(self.errors) == 0
        
        return {
            "valid": is_valid,
            "errors": self.errors,
            "warnings": self.warnings,
            "suggestions": self.suggestions,
            "score": self._calculate_score()
        }
    
    def _calculate_score(self) -> int:
        """计算提示词质量分数（0-100）"""
        score = 100
        score -= len(self.errors) * 30
        score -= len(self.warnings) * 10
        score -= len(self.suggestions) * 5
        return max(0, score)


def validate_prompt(prompt: str, prompt_type: str = "image") -> Dict[str, any]:
    """
    验证提示词
    
    Args:
        prompt: 提示词文本
        prompt_type: 提示词类型（"image" 或 "video"）
    
    Returns:
        验证结果字典
    """
    validator = PromptValidator()
    
    if prompt_type == "image":
        return validator.validate_image_prompt(prompt)
    elif prompt_type == "video":
        return validator.validate_video_prompt(prompt)
    else:
        return {
            "valid": False,
            "errors": [f"未知的提示词类型: {prompt_type}"],
            "warnings": [],
            "suggestions": [],
            "score": 0
        }


def print_result(result: Dict[str, any]):
    """打印验证结果"""
    print("\n" + "="*60)
    print("提示词验证结果")
    print("="*60)
    
    # 状态
    status = "✅ 通过" if result["valid"] else "❌ 未通过"
    print(f"\n状态: {status}")
    print(f"质量分数: {result['score']}/100")
    
    # 错误
    if result["errors"]:
        print(f"\n❌ 错误 ({len(result['errors'])}):")
        for i, error in enumerate(result["errors"], 1):
            print(f"  {i}. {error}")
    
    # 警告
    if result["warnings"]:
        print(f"\n⚠️  警告 ({len(result['warnings'])}):")
        for i, warning in enumerate(result["warnings"], 1):
            print(f"  {i}. {warning}")
    
    # 建议
    if result["suggestions"]:
        print(f"\n💡 建议 ({len(result['suggestions'])}):")
        for i, suggestion in enumerate(result["suggestions"], 1):
            print(f"  {i}. {suggestion}")
    
    print("\n" + "="*60 + "\n")


def main():
    """主函数"""
    if len(sys.argv) < 2:
        print("用法: python validate_prompt.py <prompt> [type]")
        print("  prompt: 提示词文本")
        print("  type: 提示词类型（image 或 video，默认 image）")
        sys.exit(1)
    
    prompt = sys.argv[1]
    prompt_type = sys.argv[2] if len(sys.argv) > 2 else "image"
    
    print(f"\n验证 {prompt_type} 提示词:")
    print(f'"{prompt}"')
    
    result = validate_prompt(prompt, prompt_type)
    print_result(result)
    
    # 返回状态码
    sys.exit(0 if result["valid"] else 1)


if __name__ == "__main__":
    main()
