import gradio as gr
import pytexrecon  # 导入你的pytexrecon封装
import numpy as np
from PIL import Image
import io
from typing import Dict, Tuple, Optional


# -------------------------- 核心功能：拆分步骤，优先验证TexRecon实例 --------------------------
def init_texrecon(verbose: bool = True) -> pytexrecon.TexRecon:
    """单独初始化TexRecon实例（优先验证核心依赖）"""
    if verbose:
        print("=== 步骤1/3：初始化TexRecon实例（优先验证） ===")
    
    texrecon_instance = pytexrecon.TexRecon(verbose=verbose)
    texrecon_instance._check_texrecon()
    return texrecon_instance

        


def generate_test_data() -> Tuple[bytes, Dict[str, bytes], str, Dict[str, str]]:
    """生成TexRecon测试所需的模拟数据"""
    print("=== 步骤2/3：生成模拟测试数据 ===")
    # 1. 模拟PLY网格（3个顶点+1个三角形面）
    ply_content = """ply
format ascii 1.0
element vertex 3
property float x
property float y
property float z
element face 1
property list uchar int vertex_indices
end_header
0.0 0.0 0.0
1.0 0.0 0.0
0.5 1.0 0.0
3 0 1 2
"""
    mesh_bytes = ply_content.encode("utf-8")

    # 2. 模拟2张64x64纯色图像（红+蓝）
    images: Dict[str, bytes] = {}
    # 红色图像
    red_img = Image.new("RGB", (64, 64), color="red")
    red_byteio = io.BytesIO()
    red_img.save(red_byteio, format="PNG")
    images["test1.png"] = red_byteio.getvalue()
    # 蓝色图像
    blue_img = Image.new("RGB", (64, 64), color="blue")
    blue_byteio = io.BytesIO()
    blue_img.save(blue_byteio, format="PNG")
    images["test2.png"] = blue_byteio.getvalue()

    # 3. 模拟相机内参（针孔相机，适配64x64图像）
    intrinsics = "0 32.0 32.0 32.0 32.0 0.0 0.0 0.0 0.0"

    # 4. 模拟相机姿态（单位矩阵，相机在z轴2m处）
    pose_matrix = """1.0 0.0 0.0 0.0
0.0 1.0 0.0 0.0
0.0 0.0 1.0 2.0
0.0 0.0 0.0 1.0"""
    poses = {"test1.png": pose_matrix, "test2.png": pose_matrix}

    return mesh_bytes, images, intrinsics, poses


def run_texturing_with_test_data(texrecon_instance: pytexrecon.TexRecon) -> Tuple[bytes, bytes, Image.Image]:
    """用模拟数据调用TexRecon纹理重建"""
    print("=== 步骤3/3：执行纹理重建 ===")
    mesh_data, images, intrinsics, poses = generate_test_data()
    try:
        textured_mesh_bytes, texture_bytes = texrecon_instance.run_texturing(
            mesh_data=mesh_data,
            images=images,
            intrinsics=intrinsics,
            poses=poses,
            texture_resolution=256,
            fill_holes=True
        )
        texture_img = Image.open(io.BytesIO(texture_bytes))
        return textured_mesh_bytes, texture_bytes, texture_img
    except Exception as e:
        raise RuntimeError(f"纹理重建执行失败：\n{str(e)}")


# -------------------------- 测试流程：分步执行，优先验证实例 --------------------------
def full_test_flow() -> Tuple[str, Optional[pytexrecon.TexRecon], Optional[bytes], Optional[bytes], Optional[Image.Image]]:
    """完整测试流程（分步反馈）"""
    status = ""
    texrecon_instance = None
    textured_mesh_bytes = None
    texture_bytes = None
    texture_img = None

    try:
        # 步骤1：优先初始化TexRecon（核心验证）
        texrecon_instance = init_texrecon()
        status += "✅ 步骤1/3：TexRecon实例初始化成功！\n"
        status += f"- 可执行文件路径：{texrecon_instance.texrecon_path}\n"

        # 步骤2+3：生成数据并执行纹理重建
        textured_mesh_bytes, texture_bytes, texture_img = run_texturing_with_test_data(texrecon_instance)
        status += "✅ 步骤2/3：模拟测试数据生成成功！\n"
        status += "✅ 步骤3/3：纹理重建执行成功！\n"
        status += "- 已生成带纹理模型（.obj）和纹理图（.png）\n"

    except Exception as e:
        status = f"❌ 测试失败：\n{str(e)}"

    return status, texrecon_instance, textured_mesh_bytes, texture_bytes, texture_img


# -------------------------- Gradio界面：修复File.update()错误 --------------------------
with gr.Blocks(title="pytexrecon 分步测试工具") as demo:
    gr.Markdown("# pytexrecon 全流程测试（优先验证TexRecon）")
    gr.Markdown("🔍 已修复Gradio兼容性问题，支持最新版本")

    # 存储TexRecon实例
    texrecon_state = gr.State(None)

    with gr.Row():
        test_btn = gr.Button("🚀 开始全流程测试", variant="primary", size="lg")

    # 结果展示区域
    with gr.Group():
        gr.Markdown("### 测试状态日志")
        status_box = gr.Textbox(label="状态", lines=8, interactive=False, placeholder="点击测试按钮开始...")

        gr.Markdown("### TexRecon可执行文件路径")
        texrecon_path_box = gr.Textbox(label="路径", interactive=False, placeholder="初始化成功后显示...")

        gr.Markdown("### 输出结果")
        with gr.Row():
            texture_preview = gr.Image(label="纹理图预览", interactive=False)
            mesh_download = gr.File(label="带纹理模型（.obj）")
            texture_download = gr.File(label="纹理图（.png）")

    # 绑定按钮事件：使用新的更新方式（移除.update()）
    def on_test_click():
        status, texrecon_inst, mesh_bytes, tex_bytes, tex_img = full_test_flow()
        
        # 1. 处理TexRecon路径显示
        texrecon_path = texrecon_inst.texrecon_path if texrecon_inst else "未初始化"
        
        # 2. 处理文件下载组件（直接返回值，不使用.update()）
        # 对于文件组件，直接返回 (数据, 文件名) 元组或None
        mesh_file = (mesh_bytes, "test_textured_mesh.obj") if mesh_bytes else None
        texture_file = (tex_bytes, "test_texture.png") if tex_bytes else None

        return (
            status,
            texrecon_inst,
            texrecon_path,
            mesh_file,
            texture_file,
            tex_img
        )

    # 绑定输出映射
    test_btn.click(
        fn=on_test_click,
        outputs=[
            status_box,
            texrecon_state,
            texrecon_path_box,
            mesh_download,
            texture_download,
            texture_preview
        ]
    )


# 启动应用
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)