import os
import time
import glob
import json
import numpy as np
import trimesh
import argparse
from scipy.spatial.transform import Rotation
import PIL.Image
from PIL import Image
import math
import trimesh.transformations as tf
from trimesh.exchange.gltf import export_glb

os.environ['PYOPENGL_PLATFORM'] = 'egl'

import gradio as gr

def parse_args():
    parser = argparse.ArgumentParser(description='Create animations for 3D models')
    
    parser.add_argument(
        '--input',
        type=str,
        default='./data/demo_glb/',
        help='Input file or directory path (default: ./data/demo_glb/)'
    )
    
    parser.add_argument(
        '--log_path',
        type=str,
        default='./results/demo',
        help='Output directory path (default: results/demo)'
    )
    
    parser.add_argument(
        '--animation_type',
        type=str,
        default='rotate',
        choices=['rotate', 'float', 'explode', 'assemble', 'pulse', 'swing'],
        help='Type of animation to apply'
    )
    
    parser.add_argument(
        '--animation_duration',
        type=float,
        default=3.0,
        help='Duration of animation in seconds'
    )
    
    parser.add_argument(
        '--fps',
        type=int,
        default=30,
        help='Frames per second for animation'
    )
    
    return parser.parse_args()

def get_input_files(input_path):
    if os.path.isfile(input_path):
        return [input_path]
    elif os.path.isdir(input_path):
        return glob.glob(os.path.join(input_path, '*'))
    else:
        raise ValueError(f"Input path {input_path} is neither a file nor a directory")

args = parse_args()

LOG_PATH = args.log_path
os.makedirs(LOG_PATH, exist_ok=True)

print(f"Output directory: {LOG_PATH}")

def normalize_mesh(mesh):
    """Normalize mesh to fit in a unit cube centered at origin"""
    vertices = mesh.vertices
    bounds = np.array([vertices.min(axis=0), vertices.max(axis=0)])
    center = (bounds[0] + bounds[1]) / 2
    scale = 1.0 / (bounds[1] - bounds[0]).max()
    
    # Create a copy to avoid modifying the original
    normalized_mesh = mesh.copy()
    normalized_mesh.vertices = (vertices - center) * scale
    
    return normalized_mesh, center, scale

def create_rotation_animation(mesh, duration=3.0, fps=30):
    """Create a rotation animation around the Y axis"""
    num_frames = int(duration * fps)
    frames = []
    
    # Normalize the mesh for consistent animation
    mesh, original_center, original_scale = normalize_mesh(mesh)
    
    for frame_idx in range(num_frames):
        t = frame_idx / (num_frames - 1)  # Normalized time [0, 1]
        angle = t * 2 * math.pi  # Full rotation
        
        # Create a copy of the mesh to animate
        animated_mesh = mesh.copy()
        
        # Apply rotation around Y axis
        rotation_matrix = tf.rotation_matrix(angle, [0, 1, 0])
        animated_mesh.apply_transform(rotation_matrix)
        
        # Add to frames
        frames.append(animated_mesh)
    
    return frames

def create_float_animation(mesh, duration=3.0, fps=30, amplitude=0.2):
    """Create a floating animation where the mesh moves up and down"""
    num_frames = int(duration * fps)
    frames = []
    
    # Normalize the mesh for consistent animation
    mesh, original_center, original_scale = normalize_mesh(mesh)
    
    for frame_idx in range(num_frames):
        t = frame_idx / (num_frames - 1)  # Normalized time [0, 1]
        
        # Create a copy of the mesh to animate
        animated_mesh = mesh.copy()
        
        # Apply floating motion (sinusoidal)
        y_offset = amplitude * math.sin(2 * math.pi * t)
        translation_matrix = tf.translation_matrix([0, y_offset, 0])
        animated_mesh.apply_transform(translation_matrix)
        
        # Add to frames
        frames.append(animated_mesh)
    
    return frames

def create_explode_animation(mesh, duration=3.0, fps=30):
    """Create an explode animation where parts of the mesh move outward"""
    num_frames = int(duration * fps)
    frames = []
    
    # Normalize the mesh for consistent animation
    mesh, original_center, original_scale = normalize_mesh(mesh)
    
    # Split the mesh into components
    # If the mesh can't be split, we'll just move vertices outward
    try:
        components = mesh.split(only_watertight=False)
        if len(components) <= 1:
            raise ValueError("Mesh cannot be split into components")
    except:
        # If splitting fails, work with the original mesh
        components = None
    
    for frame_idx in range(num_frames):
        t = frame_idx / (num_frames - 1)  # Normalized time [0, 1]
        
        if components:
            # Create a scene to hold all components
            scene = trimesh.Scene()
            
            # Move each component outward from center
            for component in components:
                # Create a copy of the component
                animated_component = component.copy()
                
                # Calculate direction from center to component centroid
                direction = animated_component.centroid
                if np.linalg.norm(direction) < 1e-10:
                    # If component is at center, choose random direction
                    direction = np.random.rand(3) - 0.5
                
                direction = direction / np.linalg.norm(direction)
                
                # Apply explosion movement
                translation = direction * t * 0.5  # Scale factor for explosion
                translation_matrix = tf.translation_matrix(translation)
                animated_component.apply_transform(translation_matrix)
                
                # Add to scene
                scene.add_geometry(animated_component)
            
            # Convert scene to mesh (approximation)
            animated_mesh = trimesh.util.concatenate(scene.dump())
        else:
            # Work with vertices directly if components approach failed
            animated_mesh = mesh.copy()
            vertices = animated_mesh.vertices.copy()
            
            # Calculate directions from center (0,0,0) to each vertex
            directions = vertices.copy()
            norms = np.linalg.norm(directions, axis=1, keepdims=True)
            mask = norms > 1e-10
            directions[mask] = directions[mask] / norms[mask]
            directions[~mask] = np.random.rand(np.sum(~mask), 3) - 0.5
            
            # Apply explosion factor
            vertices += directions * t * 0.3
            animated_mesh.vertices = vertices
        
        # Add to frames
        frames.append(animated_mesh)
    
    return frames

def create_assemble_animation(mesh, duration=3.0, fps=30):
    """Create an assembly animation (reverse of explode)"""
    # Get explode animation and reverse it
    explode_frames = create_explode_animation(mesh, duration, fps)
    return list(reversed(explode_frames))

def create_pulse_animation(mesh, duration=3.0, fps=30, min_scale=0.8, max_scale=1.2):
    """Create a pulsing animation where the mesh scales up and down"""
    num_frames = int(duration * fps)
    frames = []
    
    # Normalize the mesh for consistent animation
    mesh, original_center, original_scale = normalize_mesh(mesh)
    
    for frame_idx in range(num_frames):
        t = frame_idx / (num_frames - 1)  # Normalized time [0, 1]
        
        # Create a copy of the mesh to animate
        animated_mesh = mesh.copy()
        
        # Apply pulsing motion (sinusoidal scale)
        scale_factor = min_scale + (max_scale - min_scale) * (0.5 + 0.5 * math.sin(2 * math.pi * t))
        scale_matrix = tf.scale_matrix(scale_factor)
        animated_mesh.apply_transform(scale_matrix)
        
        # Add to frames
        frames.append(animated_mesh)
    
    return frames

def create_swing_animation(mesh, duration=3.0, fps=30, max_angle=math.pi/6):
    """Create a swinging animation where the mesh rotates back and forth"""
    num_frames = int(duration * fps)
    frames = []
    
    # Normalize the mesh for consistent animation
    mesh, original_center, original_scale = normalize_mesh(mesh)
    
    for frame_idx in range(num_frames):
        t = frame_idx / (num_frames - 1)  # Normalized time [0, 1]
        
        # Create a copy of the mesh to animate
        animated_mesh = mesh.copy()
        
        # Apply swinging motion (sinusoidal rotation)
        angle = max_angle * math.sin(2 * math.pi * t)
        rotation_matrix = tf.rotation_matrix(angle, [0, 1, 0])
        animated_mesh.apply_transform(rotation_matrix)
        
        # Add to frames
        frames.append(animated_mesh)
    
    return frames

def generate_gif_from_frames(frames, output_path, fps=30, resolution=(640, 480), background_color=(255, 255, 255, 255)):
    """Generate a GIF from animation frames"""
    gif_frames = []
    
    for frame in frames:
        # Create a scene with the frame
        scene = trimesh.Scene(frame)
        
        # Set camera and rendering parameters
        try:
            # Try to get a good view of the object
            scene.camera_transform = scene.camera_transform
        except:
            # If that fails, use a default camera position
            scene.camera_transform = tf.translation_matrix([0, 0, 2])
        
        # Render the frame
        try:
            img = scene.save_image(resolution=resolution, background=background_color)
            gif_frames.append(Image.open(img))
        except Exception as e:
            print(f"Error rendering frame: {str(e)}")
            # Create a blank image if rendering fails
            gif_frames.append(Image.new('RGB', resolution, (255, 255, 255)))
    
    # Save as GIF
    if gif_frames:
        gif_frames[0].save(
            output_path,
            save_all=True,
            append_images=gif_frames[1:],
            optimize=False,
            duration=int(1000 / fps),
            loop=0
        )
        return output_path
    else:
        return None

def create_animation_mesh(input_mesh_path, animation_type='rotate', duration=3.0, fps=30):
    """Create animation from input mesh based on animation type"""
    # Load the mesh
    try:
        mesh = trimesh.load(input_mesh_path)
    except Exception as e:
        print(f"Error loading mesh: {str(e)}")
        return None, None
    
    # Generate animation frames based on animation type
    if animation_type == 'rotate':
        frames = create_rotation_animation(mesh, duration, fps)
    elif animation_type == 'float':
        frames = create_float_animation(mesh, duration, fps)
    elif animation_type == 'explode':
        frames = create_explode_animation(mesh, duration, fps)
    elif animation_type == 'assemble':
        frames = create_assemble_animation(mesh, duration, fps)
    elif animation_type == 'pulse':
        frames = create_pulse_animation(mesh, duration, fps)
    elif animation_type == 'swing':
        frames = create_swing_animation(mesh, duration, fps)
    else:
        print(f"Unknown animation type: {animation_type}")
        return None, None
    
    base_filename = os.path.basename(input_mesh_path).rsplit('.', 1)[0]
    
    # Save animated mesh as GLB
    try:
        animated_glb_path = os.path.join(LOG_PATH, f'animated_{base_filename}.glb')
        
        # For GLB output, we'll use the first frame for now
        # In a production environment, you'd want to use proper animation keyframes
        if frames and len(frames) > 0:
            # First frame for static GLB
            first_frame = frames[0]
            # Export as GLB
            scene = trimesh.Scene(first_frame)
            scene.export(animated_glb_path)
        else:
            return None, None
    except Exception as e:
        print(f"Error exporting GLB: {str(e)}")
        animated_glb_path = None
    
    # Create GIF for preview
    try:
        animated_gif_path = os.path.join(LOG_PATH, f'animated_{base_filename}.gif')
        generate_gif_from_frames(frames, animated_gif_path, fps)
    except Exception as e:
        print(f"Error creating GIF: {str(e)}")
        animated_gif_path = None
    
    return animated_glb_path, animated_gif_path

def process_3d_model(input_3d, animation_type, animation_duration, fps):
    """Process a 3D model and apply animation"""
    print(f"Processing: {input_3d} with animation type: {animation_type}")
    
    try:
        # Create animation
        animated_glb_path, animated_gif_path = create_animation_mesh(
            input_3d,
            animation_type=animation_type,
            duration=animation_duration,
            fps=fps
        )
        
        if not animated_glb_path or not animated_gif_path:
            return "Error creating animation", None
        
        # Create a simple JSON metadata file
        metadata = {
            "animation_type": animation_type,
            "duration": animation_duration,
            "fps": fps,
            "original_model": os.path.basename(input_3d),
            "created_at": time.strftime("%Y-%m-%d %H:%M:%S")
        }
        
        json_path = os.path.join(LOG_PATH, f'metadata_{os.path.basename(input_3d).rsplit(".", 1)[0]}.json')
        with open(json_path, 'w') as f:
            json.dump(metadata, f, indent=4)
        
        return animated_glb_path, animated_gif_path, json_path
    except Exception as e:
        error_msg = f"Error processing file: {str(e)}"
        print(error_msg)
        return error_msg, None, None

_HEADER_ = '''
<h2><b>GLB 애니메이션 생성기 - 3D 모델 움직임 효과</b></h2>

이 데모를 통해 정적인 3D 모델(GLB 파일)에 다양한 애니메이션 효과를 적용할 수 있습니다.

❗️❗️❗️**중요사항:**
- 이 데모는 업로드된 GLB 파일에 애니메이션을 적용합니다.
- 다양한 애니메이션 스타일 중에서 선택하세요: 회전, 부유, 폭발, 조립, 펄스, 스윙.
- 결과는 애니메이션된 GLB 파일과 미리보기용 GIF 파일로 제공됩니다.
'''

_INFO_ = r"""
### 애니메이션 유형 설명
- **회전(rotate)**: 모델이 Y축을 중심으로 회전합니다.
- **부유(float)**: 모델이 위아래로 부드럽게 떠다닙니다.
- **폭발(explode)**: 모델의 각 부분이 중심에서 바깥쪽으로 퍼져나갑니다.
- **조립(assemble)**: 폭발 애니메이션의 반대 - 부품들이 함께 모입니다.
- **펄스(pulse)**: 모델이 크기가 커졌다 작아졌다를 반복합니다.
- **스윙(swing)**: 모델이 좌우로 부드럽게 흔들립니다.

### 팁
- 애니메이션 길이와 FPS를 조절하여 움직임의 속도와 부드러움을 조절할 수 있습니다.
- 복잡한 모델은 처리 시간이 더 오래 걸릴 수 있습니다.
- GIF 미리보기는 빠른 참조용이며, 고품질 결과를 위해서는 애니메이션된 GLB 파일을 다운로드하세요.
"""

# Gradio 인터페이스 설정
def create_gradio_interface():
    with gr.Blocks(title="GLB 애니메이션 생성기") as demo:
        # 제목 섹션
        gr.Markdown(_HEADER_)
        
        with gr.Row():
            with gr.Column():
                # 입력 컴포넌트
                input_3d = gr.Model3D(label="3D 모델 파일 업로드 (GLB 포맷)")
                
                with gr.Row():
                    animation_type = gr.Dropdown(
                        label="애니메이션 유형", 
                        choices=["rotate", "float", "explode", "assemble", "pulse", "swing"],
                        value="rotate"
                    )
                
                with gr.Row():
                    animation_duration = gr.Slider(
                        label="애니메이션 길이 (초)",
                        minimum=1.0,
                        maximum=10.0,
                        value=3.0,
                        step=0.5
                    )
                    fps = gr.Slider(
                        label="초당 프레임 수",
                        minimum=15,
                        maximum=60,
                        value=30,
                        step=1
                    )
                
                submit_btn = gr.Button("모델 처리 및 애니메이션 생성")
                
            with gr.Column():
                # 출력 컴포넌트
                output_3d = gr.Model3D(label="애니메이션 적용된 3D 모델")
                output_gif = gr.Image(label="애니메이션 미리보기 (GIF)")
                output_json = gr.File(label="메타데이터 파일 다운로드")
        
        # 애니메이션 유형 설명
        gr.Markdown(_INFO_)
        
        # 버튼 동작 설정
        submit_btn.click(
            fn=process_3d_model,
            inputs=[input_3d, animation_type, animation_duration, fps],
            outputs=[output_3d, output_gif, output_json]
        )
        
        # 예제 준비
        example_files = [ [f] for f in glob.glob('./data/demo_glb/*.glb') ]
        
        if example_files:
            example = gr.Examples(
                examples=example_files,
                inputs=[input_3d],
                examples_per_page=10,
            )
    
    return demo

# 메인 실행 부분
if __name__ == "__main__":
    demo = create_gradio_interface()
    demo.launch(share=True)