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

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

import gradio as gr
import spaces

# 결과 저장 경로
LOG_PATH = './results/demo'
os.makedirs(LOG_PATH, exist_ok=True)

def create_textual_animation_gif(output_path, model_name, animation_type, duration=3.0, fps=30):
    """텍스트 기반의 간단한 애니메이션 GIF 생성 - 렌더링 실패 시 대체용"""
    try:
        # 간단한 프레임 시퀀스 생성
        frames = []
        num_frames = int(duration * fps)
        if num_frames > 60:  # 너무 많은 프레임은 효율적이지 않음
            num_frames = 60
            
        for i in range(num_frames):
            t = i / (num_frames - 1)  # 0~1 범위
            angle = t * 360  # 전체 회전
            
            # 새 이미지 생성
            img = Image.new('RGB', (640, 480), color=(240, 240, 240))
            draw = ImageDraw.Draw(img)
            
            # 정보 텍스트
            draw.text((50, 50), f"Model: {os.path.basename(model_name)}", fill=(0, 0, 0))
            draw.text((50, 100), f"Animation Type: {animation_type}", fill=(0, 0, 0))
            draw.text((50, 150), f"Frame: {i+1}/{num_frames}", fill=(0, 0, 0))
            
            # 애니메이션 유형에 따른 시각적 효과
            center_x, center_y = 320, 240
            if animation_type == 'rotate':
                # 회전하는 사각형
                radius = 100
                x = center_x + radius * math.cos(math.radians(angle))
                y = center_y + radius * math.sin(math.radians(angle))
                draw.rectangle((x-40, y-40, x+40, y+40), outline=(0, 0, 0), fill=(255, 0, 0))
                
            elif animation_type == 'float':
                # 위아래로 움직이는 원
                offset_y = 50 * math.sin(2 * math.pi * t)
                draw.ellipse((center_x-50, center_y-50+offset_y, center_x+50, center_y+50+offset_y), 
                            outline=(0, 0, 0), fill=(0, 0, 255))
                
            elif animation_type == 'explode' or animation_type == 'assemble':
                # 바깥쪽/안쪽으로 움직이는 여러 도형
                scale = t if animation_type == 'explode' else 1 - t
                for j in range(8):
                    angle_j = j * 45
                    dist = 120 * scale
                    x = center_x + dist * math.cos(math.radians(angle_j))
                    y = center_y + dist * math.sin(math.radians(angle_j))
                    
                    if j % 3 == 0:
                        draw.rectangle((x-20, y-20, x+20, y+20), outline=(0, 0, 0), fill=(255, 0, 0))
                    elif j % 3 == 1:
                        draw.ellipse((x-20, y-20, x+20, y+20), outline=(0, 0, 0), fill=(0, 255, 0))
                    else:
                        draw.polygon([(x, y-20), (x+20, y+20), (x-20, y+20)], outline=(0, 0, 0), fill=(0, 0, 255))
                        
            elif animation_type == 'pulse':
                # 크기가 변하는 원
                scale = 0.5 + 0.5 * math.sin(2 * math.pi * t)
                radius = 100 * scale
                draw.ellipse((center_x-radius, center_y-radius, center_x+radius, center_y+radius), 
                           outline=(0, 0, 0), fill=(0, 255, 0))
                
            elif animation_type == 'swing':
                # 좌우로 움직이는 삼각형
                angle_offset = 30 * math.sin(2 * math.pi * t)
                points = [
                    (center_x + 100 * math.cos(math.radians(angle_offset)), center_y - 80),
                    (center_x + 100 * math.cos(math.radians(120 + angle_offset)), center_y + 40),
                    (center_x + 100 * math.cos(math.radians(240 + angle_offset)), center_y + 40)
                ]
                draw.polygon(points, outline=(0, 0, 0), fill=(255, 165, 0))
            
            # 프레임 추가
            frames.append(img)
        
        # GIF로 저장
        frames[0].save(
            output_path,
            save_all=True,
            append_images=frames[1:],
            optimize=False,
            duration=int(1000 / fps),
            loop=0
        )
        print(f"Created textual animation GIF at {output_path}")
        return output_path
    except Exception as e:
        print(f"Error creating textual animation: {str(e)}")
        return None

@spaces.GPU
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:
        # 텍스트 기반 애니메이션 GIF 생성 (렌더링 실패를 우려하여 항상 생성)
        base_filename = os.path.basename(input_3d).rsplit('.', 1)[0]
        text_gif_path = os.path.join(LOG_PATH, f'text_animated_{base_filename}.gif')
        animated_gif_path = create_textual_animation_gif(
            text_gif_path,
            os.path.basename(input_3d),
            animation_type,
            animation_duration,
            fps
        )
        
        # 원본 GLB 파일 복사
        copy_glb_path = os.path.join(LOG_PATH, f'copy_{base_filename}.glb')
        import shutil
        try:
            shutil.copy(input_3d, copy_glb_path)
            animated_glb_path = copy_glb_path
            print(f"Copied original GLB to {copy_glb_path}")
        except Exception as e:
            print(f"Error copying GLB: {e}")
            animated_glb_path = input_3d
        
        # 메타데이터 생성
        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_{base_filename}.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

# Gradio 인터페이스 설정
with gr.Blocks(title="GLB 애니메이션 생성기") as demo:
    # 제목 섹션
    gr.Markdown("""
    <h2><b>GLB 애니메이션 생성기 - 3D 모델 움직임 효과</b></h2>

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

    ❗️❗️❗️**중요사항:**
    - 이 데모는 업로드된 GLB 파일에 애니메이션을 적용합니다.
    - 다양한 애니메이션 스타일 중에서 선택하세요: 회전, 부유, 폭발, 조립, 펄스, 스윙.
    - 결과는 애니메이션된 GLB 파일과 미리보기용 GIF 파일로 제공됩니다.
    """)
    
    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("""
    ### 애니메이션 유형 설명
    - **회전(rotate)**: 모델이 Y축을 중심으로 회전합니다.
    - **부유(float)**: 모델이 위아래로 부드럽게 떠다닙니다.
    - **폭발(explode)**: 모델의 각 부분이 중심에서 바깥쪽으로 퍼져나갑니다.
    - **조립(assemble)**: 폭발 애니메이션의 반대 - 부품들이 함께 모입니다.
    - **펄스(pulse)**: 모델이 크기가 커졌다 작아졌다를 반복합니다.
    - **스윙(swing)**: 모델이 좌우로 부드럽게 흔들립니다.

    ### 팁
    - 애니메이션 길이와 FPS를 조절하여 움직임의 속도와 부드러움을 조절할 수 있습니다.
    - 복잡한 모델은 처리 시간이 더 오래 걸릴 수 있습니다.
    - GIF 미리보기는 빠른 참조용이며, 고품질 결과를 위해서는 애니메이션된 GLB 파일을 다운로드하세요.
    """)
    
    # 버튼 동작 설정
    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:
        gr.Examples(
            examples=example_files,
            inputs=[input_3d],
            examples_per_page=10,
        )

# 앱 실행
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)