"""
USD Batch Importer for Houdini Layout Asset Gallery
Automatically imports USD assets with thumbnail generation
Compatible with Houdini 20.5+
"""

import hou
import os
import json
from pathlib import Path

class USDBatchImporter:
    def __init__(self):
        self.thumb_resolution = (512, 512)
        
    def find_usd_folders(self, root_path):
        """Find all folders containing USD files matching the pattern"""
        usd_folders = []
        root = Path(root_path)
        
        print(f"Searching in: {root}")
        
        for folder in root.rglob('*'):
            if not folder.is_dir():
                continue
                
            # Check for required USD files
            files = list(folder.glob('*.usd*'))
            file_names = [f.stem for f in files]
            
            # Look for geo, mtl, payload, and KB3D_ pattern
            has_geo = any('geo' in name.lower() for name in file_names)
            has_mtl = any('mtl' in name.lower() for name in file_names)
            has_payload = any('payload' in name.lower() for name in file_names)
            has_kb3d = any(name.startswith('KB3D_') for name in file_names)
            
            if has_geo and has_mtl and has_payload and has_kb3d:
                # Find the main KB3D file
                kb3d_file = next((f for f in files if f.stem.startswith('KB3D_')), None)
                if kb3d_file:
                    usd_folders.append({
                        'folder': folder,
                        'main_usd': kb3d_file,
                        'name': kb3d_file.stem
                    })
                    print(f"  Found: {kb3d_file.stem}")
        
        return usd_folders
    
    def generate_thumbnail_simple(self, usd_path, output_path):
        """Generate thumbnail using simple viewport capture"""
        try:
            # Get or create stage context
            stage_context = hou.node('/stage')
            if stage_context is None:
                obj = hou.node('/obj')
                stage_manager = obj.createNode('stagelayout', 'temp_stage_manager')
                stage_context = hou.node('/stage')
            
            # Create or get temp LOP network
            stage = stage_context.node('temp_thumb_stage')
            if stage is None:
                stage = stage_context.createNode('lopnet', 'temp_thumb_stage')
            
            # Clear existing children
            for child in stage.children():
                child.destroy()
            
            # Create reference node
            ref_node = stage.createNode('reference', 'ref')
            ref_node.parm('filepath1').set(str(usd_path))
            
            # Set display flag
            ref_node.setDisplayFlag(True)
            ref_node.setRenderFlag(True)
            
            # Get scene viewer
            desktop = hou.ui.curDesktop()
            scene_viewer = desktop.paneTabOfType(hou.paneTabType.SceneViewer)
            
            if scene_viewer is None:
                # Try to create a scene viewer
                pane = desktop.panes()[0]
                scene_viewer = pane.createTab(hou.paneTabType.SceneViewer)
            
            # Set to view the stage
            scene_viewer.setPwd(stage)
            
            # Frame the geometry
            viewport = scene_viewer.curViewport()
            viewport.frameAll()
            
            # Small delay to let viewport update
            import time
            time.sleep(0.1)
            
            # Capture image
            image = viewport.snapshotToImage()
            image.save(str(output_path))
            
            return True
            
        except Exception as e:
            print(f"    Thumbnail error: {e}")
            import traceback
            traceback.print_exc()
            return False
    
    def create_gallery_json(self, usd_folders, gallery_dir, generate_thumbs=True):
        """Create gallery JSON file for Houdini"""
        
        gallery_data = {
            "name": "USD Layout Assets",
            "type": "solaris",
            "items": []
        }
        
        # Create thumbnail directory
        thumb_dir = Path(gallery_dir) / 'thumbnails'
        thumb_dir.mkdir(parents=True, exist_ok=True)
        
        print(f"\nProcessing {len(usd_folders)} assets...")
        
        successful = 0
        for i, usd_data in enumerate(usd_folders):
            try:
                name = usd_data['name']
                usd_path = str(usd_data['main_usd'])
                
                print(f"\n[{i+1}/{len(usd_folders)}] {name}")
                
                # Generate thumbnail
                thumb_rel_path = None
                if generate_thumbs:
                    thumb_path = thumb_dir / f"{name}.jpg"
                    if not thumb_path.exists():
                        print(f"    Generating thumbnail...")
                        if self.generate_thumbnail_simple(usd_data['main_usd'], thumb_path):
                            thumb_rel_path = f"thumbnails/{name}.jpg"
                            print(f"    ✓ Thumbnail created")
                        else:
                            print(f"    ✗ Thumbnail failed")
                    else:
                        thumb_rel_path = f"thumbnails/{name}.jpg"
                        print(f"    ✓ Using existing thumbnail")
                
                # Add to gallery data
                item = {
                    "name": name,
                    "file": usd_path,
                    "thumbnail": thumb_rel_path
                }
                
                gallery_data["items"].append(item)
                print(f"    ✓ Added to gallery")
                successful += 1
                
            except Exception as e:
                print(f"    ✗ ERROR: {e}")
        
        return gallery_data, successful
    
    def save_asset_list(self, usd_folders, output_dir):
        """Save a simple asset list file that can be referenced"""
        
        output_file = Path(output_dir) / 'usd_asset_list.json'
        
        asset_list = []
        for usd_data in usd_folders:
            asset_list.append({
                'name': usd_data['name'],
                'path': str(usd_data['main_usd']),
                'folder': str(usd_data['folder'])
            })
        
        with open(output_file, 'w') as f:
            json.dump(asset_list, f, indent=2)
        
        print(f"\nAsset list saved to: {output_file}")
        return output_file
    
    def create_reference_nodes(self, usd_folders, generate_thumbs=True):
        """Create reference nodes in a LOP network for easy access"""
        
        # Get or create stage context
        stage_context = hou.node('/stage')
        if stage_context is None:
            # Create stage context if it doesn't exist
            obj = hou.node('/obj')
            stage_manager = obj.createNode('stagelayout', 'stage_manager')
            stage_context = hou.node('/stage')
        
        # Check for existing network
        existing = stage_context.node('USD_Asset_Library')
        if existing:
            # Ask to replace
            replace = hou.ui.displayMessage(
                "USD_Asset_Library already exists. Replace it?",
                buttons=("Replace", "Cancel"),
                default_choice=0
            )
            if replace == 0:
                existing.destroy()
            else:
                return None
        
        # Create LOP network in stage context
        stage = stage_context.createNode('lopnet', 'USD_Asset_Library')
        
        # Create thumbnail directory
        thumb_dir = Path(hou.expandString('$HIP')) / 'usd_thumbnails'
        thumb_dir.mkdir(parents=True, exist_ok=True)
        
        print(f"\nCreating reference nodes...")
        
        # Create subnet for organization
        subnet = stage.createNode('subnet', 'all_assets')
        
        # Grid layout
        x_pos = 0
        y_pos = 0
        spacing = 2
        cols = 10
        
        successful = 0
        for i, usd_data in enumerate(usd_folders):
            try:
                name = usd_data['name']
                print(f"[{i+1}/{len(usd_folders)}] Creating node: {name}")
                
                # Create reference node
                ref = subnet.createNode('reference', name)
                ref.parm('filepath1').set(str(usd_data['main_usd']))
                
                # Position node
                ref.setPosition([x_pos, -y_pos])
                
                # Update grid position
                x_pos += spacing
                if (i + 1) % cols == 0:
                    x_pos = 0
                    y_pos += spacing
                
                # Generate thumbnail if requested
                if generate_thumbs:
                    thumb_path = thumb_dir / f"{name}.jpg"
                    if not thumb_path.exists():
                        ref.setDisplayFlag(True)
                        ref.setRenderFlag(True)
                        self.generate_thumbnail_simple(usd_data['main_usd'], thumb_path)
                
                successful += 1
                
            except Exception as e:
                print(f"  ERROR: {e}")
        
        # Layout nodes
        subnet.layoutChildren()
        
        print(f"\n✓ Created {successful} reference nodes in /stage/USD_Asset_Library")
        return stage
    
    def batch_import(self, root_path, generate_thumbs=True):
        """Main batch import function"""
        print(f"\n{'='*60}")
        print(f"USD BATCH IMPORTER")
        print(f"{'='*60}")
        print(f"Scanning: {root_path}\n")
        
        # Find all USD folders
        usd_folders = self.find_usd_folders(root_path)
        print(f"\nFound {len(usd_folders)} USD asset folders")
        
        if not usd_folders:
            hou.ui.displayMessage(
                "No USD folders found!\n\nLooking for folders containing:\n- geo.usd\n- mtl.usd\n- payload.usd\n- KB3D_*.usd",
                severity=hou.severityType.Warning
            )
            return
        
        # Confirm before processing
        proceed = hou.ui.displayMessage(
            f"Found {len(usd_folders)} assets.\n\nThis will create a LOP network with all assets as reference nodes.\n\nGenerate thumbnails: {'YES' if generate_thumbs else 'NO'}\n\nProceed?",
            buttons=("Create Network", "Cancel"),
            severity=hou.severityType.Message,
            default_choice=0,
            close_choice=1
        )
        
        if proceed == 1:
            print("Cancelled by user")
            return
        
        # Save asset list
        output_dir = Path(hou.expandString('$HIP'))
        if str(output_dir) == '$HIP':
            output_dir = Path(root_path)
        
        asset_list_file = self.save_asset_list(usd_folders, output_dir)
        
        # Create reference nodes
        stage = self.create_reference_nodes(usd_folders, generate_thumbs)
        
        if stage:
            # Final message
            msg = f"Import Complete!\n\n"
            msg += f"Created: /stage/USD_Asset_Library\n"
            msg += f"Assets: {len(usd_folders)}\n"
            msg += f"Asset list: {asset_list_file}\n\n"
            msg += f"You can now:\n"
            msg += f"1. Browse assets in the network editor\n"
            msg += f"2. Copy/paste nodes as needed\n"
            msg += f"3. Use Layout Asset Gallery node and import USD files"
            
            hou.ui.displayMessage(msg, severity=hou.severityType.Message)
            
            # Jump to the network
            desktop = hou.ui.curDesktop()
            network_editor = desktop.paneTabOfType(hou.paneTabType.NetworkEditor)
            if network_editor:
                network_editor.setPwd(stage)


def run_importer():
    """UI wrapper to run the importer"""
    print("\n" + "="*60)
    print("LAUNCHING USD BATCH IMPORTER")
    print("="*60 + "\n")
    
    # Prompt for root directory
    root_path = hou.ui.selectFile(
        title="Select Root USD Directory (containing all asset folders)",
        file_type=hou.fileType.Directory,
        chooser_mode=hou.fileChooserMode.Read
    )
    
    if not root_path:
        print("No directory selected. Cancelled.")
        return
    
    # Clean up path
    root_path = root_path.replace('file://', '').strip()
    
    # Confirm thumbnail generation
    generate_thumbs = hou.ui.displayMessage(
        "Generate thumbnails?\n\n(Recommended - creates preview images for each asset)",
        buttons=("Yes", "No", "Cancel"),
        severity=hou.severityType.Message,
        default_choice=0,
        close_choice=2
    )
    
    if generate_thumbs == 2:
        print("Cancelled by user")
        return
    
    generate_thumbs = (generate_thumbs == 0)
    
    # Run importer
    importer = USDBatchImporter()
    importer.batch_import(root_path, generate_thumbs)


# Execute immediately
run_importer()