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
import shutil

# 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:
            # Handle both file paths and file-like objects
            if hasattr(self.file_path, 'read'):
                # File-like object from Gradio
                data = self.file_path.read()
            else:
                # File path
                with open(self.file_path, 'rb') as f:
                    data = f.read()
            
            self.buffer = data
            self.view = memoryview(data)
            self.offset = 0
            
            # Read header
            header = self.read_string(4)
            if header != 'VOX ':
                raise ValueError("Invalid VOX file header")
            
            version = self.read_int32()
            
            # Parse chunks
            while self.offset < len(data):
                chunk_id = self.read_string(4)
                chunk_size = self.read_int32()
                child_size = self.read_int32()
                
                if chunk_id == 'SIZE':
                    self.parse_size(chunk_size)
                elif chunk_id == 'XYZI':
                    self.parse_voxels(chunk_size)
                elif chunk_id == 'RGBA':
                    self.parse_palette(chunk_size)
                else:
                    self.offset += chunk_size
                
                # Safety check
                if self.offset >= len(data):
                    break
                
            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, size: int):
        """Parse size chunk"""
        if self.offset + 12 <= len(self.buffer):
            self.size = {
                'x': self.read_int32(),
                'y': self.read_int32(),
                'z': self.read_int32()
            }
    
    def parse_voxels(self, size: int):
        """Parse voxel data"""
        if self.offset + 4 <= len(self.buffer):
            num_voxels = self.read_int32()
            
            for i in range(num_voxels):
                if self.offset + 4 <= len(self.buffer):
                    x, y, z, color_index = struct.unpack('BBBB', self.buffer[self.offset:self.offset+4])
                    self.voxels.append({
                        'x': x,
                        'y': y, 
                        'z': z,
                        'color_index': color_index
                    })
                    self.offset += 4
    
    def parse_palette(self, size: int):
        """Parse palette data"""
        for i in range(256):
            if self.offset + 4 <= len(self.buffer):
                r, g, b, a = struct.unpack('BBBB', self.buffer[self.offset:self.offset+4])
                self.palette.append({
                    'r': r,
                    'g': g,
                    'b': b,
                    'a': a
                })
                self.offset += 4
    
    def _default_palette(self) -> list:
        """Default MagicaVoxel palette if none found"""
        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

    def read_string(self, length: int) -> str:
        """Read string from buffer"""
        if self.offset + length <= len(self.buffer):
            result = self.buffer[self.offset:self.offset+length].tobytes().decode('ascii', errors='ignore')
            self.offset += length
            return result
        return ""
    
    def read_int32(self) -> int:
        """Read 32-bit integer"""
        if self.offset + 4 <= len(self.buffer):
            result = struct.unpack('<I', self.buffer[self.offset:self.offset+4])[0]
            self.offset += 4
            return result
        return 0

class VoxToGlbConverter:
    """Convert MagicaVoxel files to GLB format"""
    
    def __init__(self):
        self.voxel_size = 1.0
        
    def vox_to_glb(self, vox_file) -> Tuple[str, str]:
        """
        Convert .vox file to .glb file
        
        Args:
            vox_file: Gradio file object or file path
            
        Returns:
            Tuple of (glb_file_path, status_message)
        """
        try:
            # Create a temporary file to handle the upload
            with tempfile.NamedTemporaryFile(delete=False, suffix='.vox') as tmp_file:
                if hasattr(vox_file, 'name'):
                    # Gradio file object - copy to temp file
                    shutil.copy(vox_file.name, tmp_file.name)
                    temp_vox_path = tmp_file.name
                else:
                    # Direct file path
                    temp_vox_path = vox_file
                
                # Parse the .vox file
                parser = VoxParser(temp_vox_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 to a permanent location
                output_path = str(Path(tempfile.gettempdir()) / f"converted_{os.path.basename(temp_vox_path)}.glb")
                mesh.export(output_path)
                
                # Clean up temp file
                if temp_vox_path != vox_file:
                    os.unlink(temp_vox_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)
        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]
        )
        
        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
        """)
    
    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
    )