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

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class VoxParser:
    """Parse MagicaVoxel .vox files"""
    def __init__(self, file_path):
        self.file_path = file_path
        self.voxels = []
        self.palette = []
        self.size = {}
        
    def parse(self) -> dict:
        """Parse the .vox file and extract voxel data"""
        try:
            with open(self.file_path, 'rb') as f:
                # Read header
                header = f.read(4).decode('ascii')
                if header != 'VOX ':
                    raise ValueError("Invalid VOX file header")
                
                version = struct.unpack('<I', f.read(4))[0]
                
                # Parse chunks
                while True:
                    chunk_data = f.read(8)
                    if len(chunk_data) < 8:
                        break
                        
                    chunk_id = chunk_data[:4].decode('ascii')
                    chunk_size, child_size = struct.unpack('<II', chunk_data[4:])
                    
                    chunk_content = f.read(chunk_size)
                    
                    if chunk_id == 'SIZE':
                        self.parse_size(chunk_content)
                    elif chunk_id == 'XYZI':
                        self.parse_voxels(chunk_content)
                    elif chunk_id == 'RGBA':
                        self.parse_palette(chunk_content)
                    else:
                        # Skip unknown chunks
                        if chunk_size > 0:
                            f.seek(chunk_size, 1)
                
                return {
                    'voxels': self.voxels,
                    'palette': self.palette or self._default_palette(),
                    'size': self.size
                }
        except Exception as e:
            logger.error(f"Error parsing VOX file: {e}")
            raise
    
    def parse_size(self, data: bytes):
        """Parse size chunk"""
        self.size = {
            'x': struct.unpack('<I', data[:4])[0],
            'y': struct.unpack('<I', data[4:8])[0],
            'z': struct.unpack('<I', data[8:12])[0]
        }
    
    def parse_voxels(self, data: bytes):
        """Parse voxel data"""
        num_voxels = struct.unpack('<I', data[: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])
                self.voxels.append({
                    'x': x,
                    'y': y, 
                    'z': z,
                    'color_index': color_index
                })
                offset += 4
    
    def parse_palette(self, data: bytes):
        """Parse palette data"""
        offset = 0
        for i in range(256):
            if offset + 4 <= len(data):
                r, g, b, a = struct.unpack('BBBB', data[offset:offset+4])
                self.palette.append({
                    'r': r,
                    'g': g,
                    'b': b,
                    'a': a
                })
                offset += 4
    
    def _default_palette(self) -> list:
        """Default MagicaVoxel palette if none found"""
        colors = []
        # Create a simple gradient palette
        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:
    """Convert MagicaVoxel files to GLB format"""
    
    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
        
        Args:
            vox_file_path: Path to the .vox file
            
        Returns:
            Tuple of (glb_file_path, status_message)
        """
        try:
            # Parse the .vox file
            parser = VoxParser(vox_file_path)
            voxel_data = parser.parse()
            
            if not voxel_data['voxels']:
                return "", "No voxels found in the file"
            
            # Create mesh from voxels
            mesh = self.create_mesh_from_voxels(voxel_data)
            
            # Save as GLB
            output_path = str(Path(vox_file_path).with_suffix('.glb'))
            mesh.export(output_path)
            
            return output_path, f"Successfully converted {len(voxel_data['voxels'])} voxels to GLB format"
            
        except Exception as e:
            logger.error(f"Conversion error: {e}")
            return "", f"Error converting file: {str(e)}"
    
    def create_mesh_from_voxels(self, voxel_data: dict) -> trimesh.Trimesh:
        """Create a trimesh from voxel data"""
        voxels = voxel_data['voxels']
        palette = voxel_data['palette']
        
        # Group voxels by color for better performance
        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)
        
        # Create meshes for each color group
        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}
            
            # Create instanced cubes for this color
            cube = trimesh.creation.box(extents=[self.voxel_size, self.voxel_size, self.voxel_size])
            
            for voxel in voxels:
                # Position the cube
                translation = trimesh.transformations.translation_matrix([
                    voxel['x'] * self.voxel_size,
                    voxel['z'] * self.voxel_size,  # Swap Y and Z for proper orientation
                    voxel['y'] * self.voxel_size
                ])
                
                transformed_cube = cube.copy()
                transformed_cube.apply_transform(translation)
                
                # Set vertex colors
                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)
        
        # Combine all meshes
        if meshes:
            combined = trimesh.util.concatenate(meshes)
            return combined
        else:
            # Fallback: create a simple cube
            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()
        glb_path, message = converter.vox_to_glb(vox_file.name)
        return glb_path, message
    except Exception as e:
        return "", f"Error: {str(e)}"

def create_gradio_interface():
    """Create the Gradio interface"""
    
    with gr.Blocks(title="VOX to GLB Converter", theme=gr.themes.Soft()) as app:
        gr.Markdown("""
        # 🧊 MagicaVoxel to GLB Converter
        
        Convert your MagicaVoxel `.vox` files to `.glb` format for preview and use in 3D applications.
        
        **Features:**
        - ✅ Preserves voxel colors and structure
        - ✅ Optimized for 3D preview
        - ✅ Handles legacy VOX formats (2019-2020)
        - ✅ Downloads as GLB file
        """)
        
        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 to convert..."
                )
                
            with gr.Column():
                glb_output = gr.File(
                    label="Download GLB File",
                    file_types=[".glb"],
                    interactive=False
                )
                
                preview_info = gr.Markdown("""
                **Preview Info:**
                - Download the GLB file
                - Drag into any 3D viewer
                - Compatible with Three.js, Babylon.js, Unity, Blender
                """)
        
        # Event handlers
        convert_btn.click(
            fn=process_vox_file,
            inputs=[vox_input],
            outputs=[glb_output, status_output]
        )
        
        # Examples
        gr.Markdown("### 📁 Example Usage")
        gr.Markdown("""
        **How to use:**
        1. Click "Upload VOX File" and select your `.vox` file
        2. Click "Convert to GLB" 
        3. Download the converted `.glb` file
        4. Preview in any 3D viewer or use in your projects
        
        **Supported formats:**
        - MagicaVoxel .vox files (all versions)
        - Preserves colors and voxel positions
        - Optimized mesh output
        """)
    
    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
    )