# modules/visualization/routes.py
from flask import Blueprint, request, jsonify, send_file
import os
import time
import re
import requests
import plotly.graph_objects as go
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from io import BytesIO
import matplotlib
import json
matplotlib.use('Agg')  # 非交互式后端

visualization_bp = Blueprint('visualization', __name__)

# 确保存储目录存在
os.makedirs('static', exist_ok=True)

@visualization_bp.route('/mindmap', methods=['POST'])
def generate_mindmap():
    """生成思维导图"""
    try:
        data = request.json
        content = data.get('content', '')
        
        if not content:
            return jsonify({
                "success": False,
                "message": "思维导图内容不能为空"
            }), 400
        
        # 提取@startmindmap和@endmindmap之间的内容
        pattern = r'@startmindmap\n([\s\S]*?)@endmindmap'
        match = re.search(pattern, content)
        
        if match:
            processed_content = f"@startmindmap\n{match.group(1)}\n@endmindmap"
        else:
            # 如果内容中没有正确格式，直接使用原始内容
            processed_content = content
            
            # 确保内容有开始和结束标记
            if '@startmindmap' not in processed_content:
                processed_content = "@startmindmap\n" + processed_content
            if '@endmindmap' not in processed_content:
                processed_content += "\n@endmindmap"
        
        # 使用HuggingFace Space API生成图像
        response = requests.post(
            'https://mistpe-flask.hf.space/v1/images/generations',
            headers={
                'Authorization': 'Bearer sk-xxx',  # 实际应用中应从配置中获取
                'Content-Type': 'application/json'
            },
            json={
                'model': 'dall-e-3',
                'prompt': processed_content,
                'n': 1,
                'size': '1024x1024'
            }
        )
        
        if response.status_code != 200:
            return jsonify({
                "success": False,
                "message": "生成思维导图图像失败",
                "status_code": response.status_code
            }), 500
            
        response_data = response.json()
        
        if not response_data.get('data') or not response_data['data'][0].get('url'):
            return jsonify({
                "success": False,
                "message": "API返回的数据中没有图像URL"
            }), 500
            
        image_url = response_data['data'][0]['url']
        
        # 下载图像并保存到本地
        img_response = requests.get(image_url)
        if img_response.status_code != 200:
            return jsonify({
                "success": False,
                "message": "下载生成的图像失败"
            }), 500
            
        # 保存图像到本地
        img_filename = f'mindmap_{int(time.time())}.png'
        img_path = os.path.join('static', img_filename)
        
        with open(img_path, 'wb') as f:
            f.write(img_response.content)
        
        # 构建URL
        local_img_url = f'/static/{img_filename}'
        
        return jsonify({
            "success": True,
            "url": local_img_url,
            "original_url": image_url,
            "message": "思维导图生成成功"
        })
        
    except Exception as e:
        import traceback
        traceback.print_exc()
        return jsonify({
            "success": False,
            "message": str(e)
        }), 500
@visualization_bp.route('/3d-surface', methods=['POST'])
def generate_3d_surface():
    """生成3D表面图并返回嵌入式HTML"""
    try:
        data = request.json
        code = data.get('code', '')
        
        if not code:
            return jsonify({
                "success": False,
                "message": "请提供函数代码"
            }), 400
        
        # 清理代码（移除可能的Markdown标记）
        code = re.sub(r'```python\n', '', code)
        code = re.sub(r'```', '', code)
        
        # 在安全环境中执行代码
        local_vars = {}
        try:
            # 仅提供numpy库
            exec(code, {"__builtins__": __builtins__, "np": np, "numpy": np}, local_vars)
            
            # 检查是否有create_3d_plot函数
            if 'create_3d_plot' not in local_vars:
                return jsonify({
                    "success": False,
                    "message": "提供的代码未包含create_3d_plot函数"
                }), 500
            
            # 执行函数获取数据
            plot_data = local_vars['create_3d_plot']()
            
            if not isinstance(plot_data, dict) or not all(k in plot_data for k in ['x', 'y', 'z']):
                return jsonify({
                    "success": False,
                    "message": "函数未返回有效的3D数据"
                }), 500
            
            # 创建Plotly图形
            fig = go.Figure()
            
            # 根据类型添加跟踪
            if plot_data.get('type') == 'scatter3d':
                fig.add_trace(go.Scatter3d(
                    x=plot_data['x'],
                    y=plot_data['y'],
                    z=plot_data['z'],
                    mode='markers',
                    marker=dict(
                        size=4,
                        color=plot_data.get('color', plot_data['z']),
                        colorscale='Viridis',
                        opacity=0.8
                    )
                ))
            else:  # 默认为surface
                # 处理数据格式
                try:
                    x = np.array(plot_data['x'])
                    y = np.array(plot_data['y'])
                    z = np.array(plot_data['z'])
                    
                    if len(x.shape) == 1 and len(y.shape) == 1:
                        X, Y = np.meshgrid(x, y)
                        if len(z.shape) == 1:
                            Z = z.reshape(len(y), len(x))
                        else:
                            Z = z
                    else:
                        X = x
                        Y = y
                        Z = z
                        
                    fig.add_trace(go.Surface(
                        z=Z,
                        x=X,
                        y=Y,
                        colorscale='Viridis'
                    ))
                except Exception as data_error:
                    return jsonify({
                        "success": False,
                        "message": f"处理3D数据时出错: {str(data_error)}"
                    }), 500
            
            # 设置图形布局
            fig.update_layout(
                title=plot_data.get('title', '3D 可视化'),
                scene=dict(
                    xaxis_title='X',
                    yaxis_title='Y',
                    zaxis_title='Z'
                ),
                width=800,
                height=600,
                margin=dict(l=0, r=0, b=0, t=30)  # 减小边距，使图形更紧凑
            )
            
            # 直接生成HTML
            html_content = fig.to_html(include_plotlyjs=True, full_html=True)
            
            # 保存为文件
            html_filename = f'3d_plot_{int(time.time())}.html'
            html_path = os.path.join('static', html_filename)
            
            with open(html_path, 'w', encoding='utf-8') as f:
                f.write(html_content)
            
            return jsonify({
                "success": True,
                "html_url": f'/static/{html_filename}',
                "message": "3D图形生成成功"
            })
            
        except Exception as e:
            import traceback
            error_traceback = traceback.format_exc()
            print(f"3D可视化执行错误: {error_traceback}")
            return jsonify({
                "success": False,
                "message": f"执行代码时出错: {str(e)}"
            }), 500
            
    except Exception as e:
        import traceback
        traceback.print_exc()
        return jsonify({
            "success": False,
            "message": str(e)
        }), 500