import gradio as gr
import numpy as np
import trimesh
import tempfile
import os
import struct
from pathlib import Path
from typing import Tuple

class VoxParser:
    """MagicaVoxel .vox file parser"""
    
    def __init__(self, file_path):
        self.file_path = file_path
        
    def parse(self) -> dict:
        """Parse the .vox file structure"""
        try:
            with open(self.file_path, 'rb') as f:
                data = f.read()
                
            offset = 0
            
            # Read header
            header = data[offset:offset+4].decode('ascii', errors='ignore')
            offset += 4
            if header != 'VOX ':
                raise ValueError("Invalid VOX file header")
            
            offset += 4  # Skip version
            
            voxels = []
            palette = []
            size = {}
            
            # Parse chunks
            while offset < len(data):
                if offset + 12 > len(data):
                    break
                    
                chunk_id = data[offset:offset+4].decode('ascii', errors='ignore')
                offset += 4
                chunk_size = struct.unpack('<I', data[offset:offset+4])[0]
                offset += 4
                child_size = struct.unpack('<I', data[offset:offset+4])[0]
                offset += 4
                
                if chunk_id == 'SIZE':
                    if offset + 12 <= len(data):
                        size = {
                            'x': struct.unpack('<I', data[offset:offset+4])[0],
                            'y': struct.unpack('<I', data[offset+4:offset+8])[0],
                            'z': struct.unpack('<I', data[offset+8:offset+12])[0]
                        }
                    offset += chunk_size
                elif chunk_id == 'XYZI':
                    if offset + 4 <= len(data):
                        num_voxels = struct.unpack('<I', data[offset:offset+4])[0]
                        offset += 4
                        
                        for i in range(num_voxels):
                            if offset + 4 <= len(data):
                                x, y, z, color_index = struct.unpack('BBBB', data[offset:offset+4])
                                voxels.append({
                                    'x': x,
                                    'y': y,
                                    'z': z,
                                    'color_index': color_index
                                })
                                offset += 4
                elif chunk_id == 'RGBA':
                    for i in range(256):
                        if offset + 4 <= len(data):
                            r, g, b, a = struct.unpack('BBBB', data[offset:offset+4])
                            palette.append({'r': r, 'g': g, 'b': b, 'a': a})
                            offset += 4
                else:
                    offset += chunk_size
                
                if offset >= len(data):
                    break
            
            return {
                'voxels': voxels,
                'palette': palette or self._default_palette(),
                'size': size
            }
        except Exception as e:
            return {'voxels': [], 'palette': self._default_palette(), 'size': {'x': 0, 'y': 0, 'z': 0}}
    
    def _default_palette(self) -> list:
        colors = []
        for i in range(256):
            intensity = i / 255.0
            colors.append({
                'r': int(intensity * 255),
                'g': int(intensity * 255),
                'b': int(intensity * 255),
                'a': 255
            })
        return colors

class VoxToGlbConverter:
    """MagicaVoxel to GLB converter"""
    
    def __init__(self):
        self.voxel_size = 1.0
        
    def vox_to_glb(self, vox_file_path: str) -> Tuple[str, str]:
        """Convert .vox file to .glb file"""
        try:
            parser = VoxParser(vox_file_path)
            voxel_data = parser.parse()
            
            if not voxel_data['voxels']:
                return "", "No voxels found in the file"
            
            mesh = self.create_mesh_from_voxels(voxel_data)
            
            output_path = str(Path(tempfile.gettempdir()) / f"converted_model.glb")
            mesh.export(output_path)
            
            voxel_count = len(voxel_data['voxels'])
            return output_path, f"Converted {voxel_count} voxels to GLB format"
                
        except Exception as e:
            return "", f"Error converting file: {str(e)}"
    
    def create_mesh_from_voxels(self, voxel_data: dict) -> trimesh.Trimesh:
        """Create mesh from voxel data"""
        voxels = voxel_data['voxels']
        palette = voxel_data['palette']
        
        color_groups = {}
        for voxel in voxels:
            color_idx = voxel['color_index']
            if color_idx not in color_groups:
                color_groups[color_idx] = []
            color_groups[color_idx].append(voxel)
        
        meshes = []
        
        for color_idx, voxels in color_groups.items():
            color = palette[color_idx] if color_idx < len(palette) else {'r': 255, 'g': 255, 'b': 255, 'a': 255}
            
            cube = trimesh.creation.box(extents=[self.voxel_size, self.voxel_size, self.voxel_size])
            
            for voxel in voxels:
                translation = trimesh.transformations.translation_matrix([
                    voxel['x'] * self.voxel_size,
                    voxel['z'] * self.voxel_size,
                    voxel['y'] * self.voxel_size
                ])
                
                transformed_cube = cube.copy()
                transformed_cube.apply_transform(translation)
                
                vertex_colors = np.tile([color['r']/255, color['g']/255, color['b']/255, color['a']/255], 
                                      (len(transformed_cube.vertices), 1))
                transformed_cube.visual.vertex_colors = vertex_colors
                
                meshes.append(transformed_cube)
        
        if meshes:
            combined = trimesh.util.concatenate(meshes)
            return combined
        else:
            return trimesh.creation.box(extents=[self.voxel_size, self.voxel_size, self.voxel_size])

def process_vox_file(vox_file) -> Tuple[str, str]:
    """Process uploaded .vox file and convert to .glb"""
    if vox_file is None:
        return "", "Please upload a .vox file"
    
    try:
        converter = VoxToGlbConverter()
        
        if hasattr(vox_file, 'name'):
            real_file_path = vox_file.name
            
            if os.path.isfile(real_file_path):
                glb_path, message = converter.vox_to_glb(real_file_path)
                return glb_path, message
            else:
                return "", "Could not access uploaded file"
        else:
            return "", "Invalid file format"
            
    except Exception as e:
        return "", f"Error: {str(e)}"

def create_gradio_interface():
    with gr.Blocks(title="VOX to GLB Converter with Preview") as app:
        gr.Markdown("""
        # 🧊 MagicaVoxel VOX to GLB Converter with 3D Preview
        Convert your MagicaVoxel `.vox` files to `.glb` format and preview them in 3D
        """)
        
        with gr.Row():
            with gr.Column():
                vox_input = gr.File(label="Upload VOX File", file_types=[".vox"], file_count="single")
                convert_btn = gr.Button("🔄 Convert to GLB", variant="primary")
                status_output = gr.Textbox(label="Status", interactive=False, placeholder="Ready...")
                
            with gr.Column():
                glb_output = gr.File(label="Download GLB File", file_types=[".glb"], interactive=False)
        
        # 3D Preview Section - Using Gradio's built-in 3D viewer
        gr.Markdown("### 🎮 3D Preview")
        
        # FIXED: Use correct variable name
        model_3d = gr.Model3D(
            label="GLB Preview",
            height=300
        )
        
        # Connect conversion to preview
        def convert_with_preview(vox_input):
            glb_path, message = process_vox_file(vox_input)
            if glb_path and os.path.exists(glb_path):
                return glb_path, message, glb_path  # Return path for both file and 3D viewer
            else:
                return None, message, None
        
        convert_btn.click(
            fn=convert_with_preview,
            inputs=[vox_input],  # FIXED: Use correct variable name
            outputs=[glb_output, status_output, model_3d]
        )
        
        gr.Markdown("""
        ### 📋 How to Use
        1. Upload your `.vox` file
        2. Click "Convert to GLB"
        3. Download the GLB file
        4. Preview your voxel model in 3D above
        """)
    
    return app

if __name__ == "__main__":
    app = create_gradio_interface()
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
        show_error=True,
        theme=gr.themes.Soft()
    )