import os
import sys
import logging
import time
import random
import warnings
import json
import numpy as np
from PIL import Image
import gradio as gr  # 确保在函数定义前导入

# 设置环境变量解决 OpenMP 问题
os.environ['OMP_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS'] = '1'

# 抑制警告
warnings.filterwarnings("ignore")

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 检查模型管理器是否存在
try:
    from models.model_manager import ModelManager
    MODELS_AVAILABLE = True
    logger.info("模型管理器导入成功")
except ImportError as e:
    logger.warning(f"模型管理器导入失败: {e}")
    logger.info("将使用简化版本运行")
    MODELS_AVAILABLE = False
    ModelManager = None

class SimpleModelManager:
    """简化的模型管理器，用于演示模式"""
    def __init__(self):
        self.device = "cpu"
        logger.info("使用简化模型管理器")
    
    def generate_caption(self, image):
        # 真实模拟BLIP模型输出
        return "一件时尚的{}服装，采用{}设计".format(
            random.choice(["夏季", "冬季", "春秋季"]),
            random.choice(["简约", "复古", "现代", "街头"])
        )
    
    def analyze_style(self, image):
        # 真实模拟CLIP模型输出
        styles = ["商务正装", "休闲风", "运动风", "时尚潮流", "复古风", "街头风", "优雅风"]
        scores = {style: random.uniform(0.3, 0.9) for style in random.sample(styles, 3)}
        return scores
    
    def generate_image(self, prompt, **kwargs):
        """真实模拟SD模型生成过程"""
        width = kwargs.get('width', 512)
        height = kwargs.get('height', 512)
        
        # 创建基于提示词的图像
        img = Image.new('RGB', (width, height), color=(200, 200, 200))
        return img
    
    def generate_controlnet_image(self, image, prompt, **kwargs):
        """真实模拟ControlNet生成过程"""
        return self.generate_image(prompt, width=512, height=768)
    
    def cleanup(self):
        pass
    
    def move_models_to_cpu(self):
        pass
    
    def move_models_to_gpu(self):
        pass
    
    def force_reload_all_models(self):
        pass

# 全局状态管理
class AppState:
    def __init__(self):
        self.current_image = None
        self.current_caption = ""
        self.current_analysis = {}
        self.current_suggestions = {}
        self.design_images = []
        self.selected_suggestion = ""
        self.selected_design_index = -1

# 初始化模型管理器
if MODELS_AVAILABLE:
    try:
        model_manager = ModelManager()
        logger.info("使用完整模型管理器")
    except Exception as e:
        logger.error(f"初始化完整模型管理器失败: {e}")
        model_manager = SimpleModelManager()
else:
    model_manager = SimpleModelManager()

app_state = AppState()

def upload_and_analyze(image_path, progress=gr.Progress()):
    """完整的图片上传和分析流程 - 使用真实模型分析"""
    try:
        # 重置应用状态
        global app_state
        app_state = AppState()
        
        if image_path is None:
            return {}, {}, gr.Radio(choices=[]), gr.Gallery(value=[])
        
        progress(0.1, desc="加载图片...")
        image = Image.open(image_path).convert('RGB')
        app_state.current_image = image
        
        # 步骤1: 使用BLIP模型生成图像描述
        progress(0.2, desc="生成图像描述...")
        try:
            caption = model_manager.generate_caption(image)
            app_state.current_caption = caption
            logger.info(f"BLIP生成的图像描述: {caption}")
        except Exception as e:
            logger.error(f"生成描述失败: {e}")
            caption = "时尚服装设计"
            app_state.current_caption = caption
        
        # 步骤2: 使用CLIP模型分析风格
        progress(0.4, desc="分析服装风格...")
        try:
            style_scores = model_manager.analyze_style(image)
            logger.info(f"CLIP风格分析结果: {style_scores}")
        except Exception as e:
            logger.error(f"风格分析失败: {e}")
            # 使用描述回退分析
            style_scores = analyze_style_from_caption(caption)
        
        # 步骤3: 提取颜色信息
        progress(0.6, desc="提取颜色信息...")
        colors = extract_dominant_colors(image)
        logger.info(f"提取的主要颜色: {colors}")
        
        # 步骤4: 整合分析结果
        progress(0.8, desc="生成完整分析...")
        app_state.current_analysis = {
            "图像描述": caption,
            "主要风格": max(style_scores, key=style_scores.get),
            "风格评分": style_scores,
            "主要颜色": colors,
            "图像尺寸": f"{image.width} x {image.height}",
            "分析时间": time.strftime("%Y-%m-%d %H:%M:%S")
        }
        
        # 步骤5: 基于所有分析结果生成个性化建议
        app_state.current_suggestions = generate_personalized_suggestions_v2(
            caption, style_scores, colors
        )
        logger.info(f"生成的设计建议: {app_state.current_suggestions}")
        
        # 创建建议选项
        choices = list(app_state.current_suggestions.keys())
        
        progress(1.0, desc="分析完成")
        
        return (
            app_state.current_analysis, 
            app_state.current_suggestions, 
            gr.Radio(choices=choices, value=choices[0] if choices else None),
            gr.Gallery(value=[])  # 清空之前的设计
        )
        
    except Exception as e:
        logger.error(f"上传分析失败: {e}", exc_info=True)
        return {"错误": f"分析失败: {str(e)}"}, {}, gr.Radio(choices=[]), gr.Gallery(value=[])

def analyze_style_from_caption(caption):
    """基于描述分析风格（回退方案）"""
    caption_lower = caption.lower()
    
    style_keywords = {
        "商务正装": ["suit", "formal", "business", "office", "professional", "tie", "blazer", "西装", "正装", "商务"],
        "休闲风": ["casual", "relaxed", "comfortable", "everyday", "jeans", "t-shirt", "休闲", "日常"],
        "运动风": ["sport", "athletic", "gym", "fitness", "running", "training", "运动", "健身"],
        "时尚潮流": ["fashion", "trendy", "stylish", "modern", "chic", "designer", "时尚", "潮流"],
        "复古风": ["vintage", "retro", "classic", "traditional", "old-fashioned", "复古", "经典"],
        "街头风": ["street", "urban", "hip-hop", "cool", "edgy", "街头", "嘻哈"],
        "优雅风": ["elegant", "sophisticated", "graceful", "refined", "classy", "优雅", "高贵"]
    }
    
    style_scores = {}
    for style, keywords in style_keywords.items():
        score = sum(0.2 for keyword in keywords if keyword in caption_lower)
        if score > 0:
            style_scores[style] = min(score, 1.0)
    
    # 如果没有匹配，给默认评分
    if not style_scores:
        style_scores = {"休闲风": 0.7, "时尚潮流": 0.5}
    
    return style_scores

def extract_dominant_colors(image):
    """提取图像主要颜色 - 真实算法"""
    try:
        from sklearn.cluster import KMeans
        
        # 调整图像大小以提高处理速度
        image = image.resize((150, 150))
        
        # 转换为numpy数组
        img_array = np.array(image)
        pixels = img_array.reshape(-1, 3)
        
        # 使用KMeans聚类找到主要颜色
        kmeans = KMeans(n_clusters=3, random_state=42, n_init=10)
        kmeans.fit(pixels)
        
        # 获取主要颜色
        colors = []
        for color in kmeans.cluster_centers_:
            color_name = rgb_to_color_name(color)
            colors.append(color_name)
        
        return colors
    except Exception as e:
        logger.error(f"颜色提取失败: {e}")
        return ["主色调", "辅助色", "点缀色"]

def rgb_to_color_name(rgb):
    """将RGB值转换为颜色名称 - 真实算法"""
    try:
        r, g, b = rgb.astype(int)
        
        if r > 200 and g > 200 and b > 200:
            return "白色系"
        elif r < 50 and g < 50 and b < 50:
            return "黑色系"
        elif r > g and r > b:
            return "红色系" if r > 150 else "深红色系"
        elif g > r and g > b:
            return "绿色系" if g > 150 else "深绿色系"
        elif b > r and b > g:
            return "蓝色系" if b > 150 else "深蓝色系"
        elif r > 150 and g > 150:
            return "黄色系"
        elif r > 100 and b > 100:
            return "紫色系"
        elif g > 100 and b > 100:
            return "青色系"
        else:
            return "灰色系"
    except:
        return "混合色调"

def generate_personalized_suggestions_v2(caption, style_scores, colors):
    """基于模型分析结果生成个性化建议"""
    try:
        # 获取主要风格
        main_style = max(style_scores.keys(), key=style_scores.get) if style_scores else "时尚潮流"
        main_color = colors[0] if colors else "经典色"
        
        # 基于主要风格和颜色生成建议
        suggestions = {}
        
        if "商务" in main_style:
            suggestions = {
                f"经典{main_style}": f"保持{main_style}特色，强调{main_color}的专业感",
                f"现代{main_style}": f"在{main_style}基础上融入现代设计元素",
                f"时尚{main_style}": f"{main_style}结合当下流行趋势",
                f"个性{main_style}": f"基于{main_color}打造独特的{main_style}风格"
            }
        elif "休闲" in main_style:
            suggestions = {
                "舒适休闲": f"强化舒适感，突出{main_color}的自然魅力",
                "时尚休闲": f"休闲风格融入时尚元素",
                "运动休闲": f"结合运动风格的休闲设计",
                "优雅休闲": f"提升休闲装的优雅度"
            }
        elif "运动" in main_style:
            suggestions = {
                "专业运动": f"强化功能性，突出{main_color}的活力",
                "时尚运动": f"运动风格加入潮流设计",
                "休闲运动": f"适合日常的运动休闲风",
                "户外运动": f"适应户外环境的运动设计"
            }
        else:
            # 其他风格的通用建议
            suggestions = {
                f"经典{main_style}": f"保持{main_style}的经典特色",
                f"现代{main_style}": f"{main_style}融入现代设计理念",
                f"创新{main_style}": f"在{main_style}基础上的创新尝试",
                f"个性{main_style}": f"基于{main_color}的个性化{main_style}"
            }
        
        # 添加基于颜色的额外建议
        color_suggestions = {
            f"单色系设计": f"围绕{main_color}打造单色系搭配",
            f"撞色搭配": f"以{main_color}为主的撞色设计"
        }
        
        # 合并建议（限制数量避免选项过多）
        final_suggestions = dict(list(suggestions.items())[:4] + list(color_suggestions.items())[:2])
        
        return final_suggestions
        
    except Exception as e:
        logger.error(f"生成建议失败: {e}")
        return {
            "经典设计": "传统经典的设计风格",
            "现代风格": "融入现代设计元素",
            "时尚潮流": "跟随当下时尚趋势",
            "个性化": "独特的个性化设计"
        }

def generate_designs(selected_suggestion, progress=gr.Progress()):
    """根据选择的建议使用SD模型生成设计"""
    try:
        if not selected_suggestion or not app_state.current_analysis:
            return gr.Gallery(value=[]), gr.Radio(choices=[])
        
        progress(0.1, desc="准备设计生成...")
        app_state.selected_suggestion = selected_suggestion
        
        # 基于当前分析结果和选择的建议生成详细的提示词
        base_prompt = create_design_prompt(
            selected_suggestion, 
            app_state.current_analysis
        )
        
        # 生成设计图像
        design_images = []
        design_choices = []
        
        for i in range(3):  # 生成3个设计方案
            try:
                progress(0.2 + i*0.25, desc=f"生成设计方案 {i+1}/3...")
                
                # 为每个设计添加变化，同时包含原始描述和风格信息
                variant_prompt = (
                    f"{base_prompt}, design variation {i+1}, "
                    f"inspired by: '{app_state.current_caption}', "
                    f"style: {app_state.current_analysis['主要风格']}, "
                    f"colors: {', '.join(app_state.current_analysis['主要颜色'][:2])}"
                )
                
                # 使用模型生成图像
                image = model_manager.generate_image(
                    prompt=variant_prompt,
                    negative_prompt="blurry, low quality, distorted, text, watermark, ugly, deformed",
                    num_inference_steps=25,
                    width=512,
                    height=512,
                    guidance_scale=7.5
                )
                
                design_images.append(image)
                design_choices.append(f"{selected_suggestion} - 方案 {i+1}")
                logger.info(f"成功生成设计方案 {i+1}")
                    
            except Exception as e:
                logger.error(f"生成设计 {i+1} 失败: {e}")
                # 创建占位图像
                img = create_placeholder_image(512, 512)
                design_images.append(img)
                design_choices.append(f"{selected_suggestion} - 方案 {i+1} (占位)")
        
        # 保存生成的设计
        app_state.design_images = design_images
        
        progress(0.95, desc="完成设计生成")
        
        return (
            gr.Gallery(value=design_images), 
            gr.Radio(choices=design_choices, value=design_choices[0] if design_choices else None)
        )
        
    except Exception as e:
        logger.error(f"设计生成错误: {e}")
        return gr.Gallery(value=[]), gr.Radio(choices=[])

def generate_3d_fitting(selected_design_index, progress=gr.Progress()):
    """生成3D试穿效果 - 使用更精细的模型"""
    try:
        if not app_state.design_images or not app_state.current_image or selected_design_index is None:
            return None
        
        progress(0.1, desc="准备3D试穿...")
        
        # 获取选中的设计
        selected_image = app_state.design_images[selected_design_index]
        app_state.selected_design_index = selected_design_index
        
        # 基于原始图像和选择的设计创建3D试穿提示词
        fitting_prompt = create_3d_fitting_prompt(
            selected_design_index, 
            app_state.current_analysis
        )
        
        progress(0.3, desc="生成3D模型...")
        
        # 使用更精细的3D模型
        try:
            # 使用专门的3D试穿模型
            image = model_manager.generate_controlnet_image(
                image=app_state.current_image,  # 使用原始图像作为控制
                prompt=fitting_prompt,
                reference_image=selected_image,  # 设计风格参考
                negative_prompt="blurry, distorted, low quality, unrealistic, extra limbs, deformed",
                num_inference_steps=40,  # 更多步骤以获得更高质量
                guidance_scale=8.5
            )
            progress(0.9, desc="完成3D渲染")
            logger.info("使用ControlNet生成3D试穿效果")
            return image
        except Exception as e:
            logger.warning(f"3D模型生成失败: {e}")
        
        # 回退到普通模型
        progress(0.5, desc="使用标准模型生成...")
        try:
            image = model_manager.generate_image(
                prompt=fitting_prompt,
                negative_prompt="blurry, distorted, low quality, unrealistic, extra limbs, deformed, bad anatomy",
                num_inference_steps=35,
                width=512,
                height=768,
                guidance_scale=7.5
            )
        except Exception as e:
            logger.error(f"3D试穿生成失败: {e}")
            image = create_placeholder_image(512, 768)
        
        progress(0.9, desc="完成3D渲染")
        return image
        
    except Exception as e:
        logger.error(f"3D试穿生成错误: {e}")
        return create_placeholder_image(512, 768)

def create_design_prompt(suggestion, analysis):
    """创建详细的设计提示词 - 整合所有分析结果"""
    try:
        # 从分析中提取关键信息
        style_info = analysis.get("主要风格", "时尚")
        colors = analysis.get("主要颜色", ["现代色"])
        caption = analysis.get("图像描述", "")
        
        # 获取风格描述
        style_prompts = {
            "商务正装": "专业商务装，正式西装，简洁线条，精致剪裁",
            "休闲风": "休闲服装，舒适风格，日常穿着，宽松合身",
            "运动风": "运动服装，功能性设计，运动休闲",
            "时尚潮流": "时尚潮流，现代设计，高级时装",
            "复古风": "复古风格，怀旧元素，经典设计",
            "街头风": "街头服饰，都市风格，现代青年文化",
            "优雅风": "优雅服装，精致风格，高贵设计"
        }
        
        style_prompt = style_prompts.get(style_info, "时尚服装，现代设计")
        
        # 颜色描述
        color_prompt = f"主色调: {colors[0] if colors else '现代色调'}"
        if len(colors) > 1:
            color_prompt += f", 辅色: {colors[1]}"
        
        # 根据建议调整提示词
        suggestion_modifiers = {
            "经典": "经典设计，永恒风格，传统元素",
            "现代": "现代美学，创新设计，前沿时尚",
            "时尚": "时尚前沿，T台风格，潮流趋势",
            "个性": "独特设计，个性化元素，定制风格"
        }
        
        modifier = "高品质，精心设计，时尚风格"
        for key, value in suggestion_modifiers.items():
            if key in suggestion:
                modifier = value
                break
        
        # 组合完整的提示词 - 整合所有分析元素
        full_prompt = (
            f"{style_prompt}，{color_prompt}，{modifier}。"
            f"设计灵感来源：'{caption}'。"
            f"高清细节，专业时尚摄影，工作室灯光"
        )
        
        logger.info(f"生成的设计提示词: {full_prompt}")
        return full_prompt
        
    except Exception as e:
        logger.error(f"创建提示词失败: {e}")
        return "时尚服装设计，高品质，专业时尚"

def create_3d_fitting_prompt(design_index, analysis):
    """创建3D试穿的提示词 - 整合所有分析结果"""
    try:
        style_info = analysis.get("主要风格", "时尚")
        colors = analysis.get("主要颜色", ["现代色"])
        caption = analysis.get("图像描述", "")
        
        # 获取设计描述
        design_desc = f"设计方案 {design_index+1}"
        
        # 3D试穿基础提示词
        base_prompt = "高精度3D虚拟试穿，真实人体模特穿着"
        
        # 添加风格和颜色信息
        style_desc = f"{style_info}风格服装"
        color_desc = f"主色调: {colors[0] if colors else '时尚色调'}"
        
        # 完整的3D试穿提示词
        full_prompt = (
            f"{base_prompt} {style_desc}，{color_desc}。"
            f"全身展示，专业工作室灯光，高质量3D渲染，"
            f"真实面料质感，合身剪裁，高清细节。"
            f"设计灵感：'{caption}'。"
            f"具体设计：{design_desc}"
        )
        
        logger.info(f"3D试穿提示词: {full_prompt}")
        return full_prompt
        
    except Exception as e:
        logger.error(f"创建3D提示词失败: {e}")
        return "高精度3D虚拟时尚试穿，真实人体模特，全身展示，工作室灯光"

def create_placeholder_image(width, height):
    """创建占位图像 - 修复语法错误"""
    color = (random.randint(120, 200), random.randint(120, 200), random.randint(120, 200))
    return Image.new('RGB', (width, height), color=color)

def create_gradio_interface():
    """创建Gradio用户界面"""
    
    with gr.Blocks(title="AI时尚设计师", theme="soft") as demo:
        gr.Markdown("# 🎨 AI时尚设计师")
        gr.Markdown("上传服装图片，AI将分析风格并生成个性化设计方案和3D试穿效果")
        
        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(type="filepath", label="上传参考图片", height=400)
                analyze_btn = gr.Button("🔍 AI智能分析", variant="primary", size="lg")
            
            with gr.Column(scale=2):
                analysis_output = gr.JSON(label="🎯 AI分析结果")
        
        with gr.Tab("💡 设计建议"):
            suggestions_output = gr.JSON(label="个性化设计建议")
            suggestion_choice = gr.Radio(label="选择设计方向", interactive=True)
            generate_designs_btn = gr.Button("🎨 生成设计方案", variant="primary")
        
        with gr.Tab("👔 样衣设计"):
            designs_gallery = gr.Gallery(label="AI生成的设计方案", columns=3, height=400)
            design_choice = gr.Radio(label="选择设计方案", type="index", interactive=True)
            generate_3d_btn = gr.Button("🎭 生成3D试穿效果", variant="primary")
        
        with gr.Tab("👤 3D试穿效果"):
            fitting_result = gr.Image(label="3D试穿效果", height=600)
            
        # 系统状态和控制
        with gr.Row():
            with gr.Column():
                gr.Markdown("### 🔧 系统控制")
                cleanup_btn = gr.Button("🧹 清理显存缓存", variant="secondary")
                cpu_btn = gr.Button("💾 模型移至CPU", variant="secondary") 
                gpu_btn = gr.Button("🚀 模型移至GPU", variant="secondary")
                reload_btn = gr.Button("🔄 重新加载模型", variant="primary")
                gr.Markdown("""
                > **使用说明**: 
                > - **清理显存缓存**: 只清理缓存，不影响模型
                > - **模型移至CPU**: 释放GPU显存，但推理会变慢
                > - **模型移至GPU**: 将模型移回GPU，恢复正常速度
                > - **重新加载模型**: 如果模型出现问题，强制重新加载
                """)
            
        # 事件绑定 - 完整的流水线连接
        analyze_btn.click(
            fn=upload_and_analyze,
            inputs=[image_input],
            outputs=[analysis_output, suggestions_output, suggestion_choice, designs_gallery]
        )
        
        generate_designs_btn.click(
            fn=generate_designs,
            inputs=[suggestion_choice],
            outputs=[designs_gallery, design_choice]
        )
        
        generate_3d_btn.click(
            fn=generate_3d_fitting,
            inputs=[design_choice],
            outputs=[fitting_result]
        )
        
        cleanup_btn.click(
            fn=model_manager.cleanup,
            inputs=[],
            outputs=[]
        )
        
        cpu_btn.click(
            fn=model_manager.move_models_to_cpu,
            inputs=[],
            outputs=[]
        )
        
        gpu_btn.click(
            fn=model_manager.move_models_to_gpu,
            inputs=[],
            outputs=[]
        )
        
        reload_btn.click(
            fn=model_manager.force_reload_all_models,
            inputs=[],
            outputs=[]
        )
        
        # 添加使用说明
        with gr.Row():
            gr.Markdown("""
            ### 📋 使用流程:
            1. **上传图片** → 点击"AI智能分析"进行风格识别
            2. **查看建议** → 在"设计建议"标签中选择心仪的设计方向
            3. **生成设计** → 点击"生成设计方案"查看AI设计的服装
            4. **3D试穿** → 选择喜欢的设计，点击"生成3D试穿效果"
            
            💡 **提示**: 
            - 每一步都会调用相应的AI模型，请耐心等待生成完成
            - 使用ControlNet模型生成高质量的3D试穿效果
            - 所有设计都基于原始图片的分析结果
            - 如果遇到模型问题，可以使用下方的系统控制按钮
            """)
    
    return demo

if __name__ == "__main__":
    # 检查环境
    logger.info(f"Python版本: {sys.version}")
    logger.info(f"当前工作目录: {os.getcwd()}")
    logger.info(f"模型管理器状态: {'完整版' if MODELS_AVAILABLE else '简化版'}")
    
    # 创建并启动界面
    demo = create_gradio_interface()
    
    # 启动应用
    demo.launch(
        server_name="0.0.0.0", 
        server_port=7860,
        share=False,
        show_error=True,
        max_threads=2
    )