#!/usr/bin/env python3
"""
Estimate VRAM usage by analyzing Godot's imported texture files.
This works without running the game.
"""

import os
import re
import json
from pathlib import Path
from collections import defaultdict

def parse_import_file(import_path):
    """Parse a .import file to get texture dimensions and format."""
    dimensions = None
    format_type = "unknown"
    compressed = False
    
    try:
        with open(import_path, 'r') as f:
            content = f.read()
            
        # Look for dimension info in the import file
        # Godot stores this in the remap section
        dim_match = re.search(r'width=(\d+)', content)
        height_match = re.search(r'height=(\d+)', content)
        
        if dim_match and height_match:
            dimensions = (int(dim_match.group(1)), int(height_match.group(2)))
        
        # Check for compression
        if 'ctex' in import_path or 'compress' in content.lower():
            compressed = True
            format_type = "compressed"
        elif 'texture' in content.lower():
            format_type = "rgba"
            
    except Exception as e:
        pass
    
    return dimensions, format_type, compressed

def estimate_vram_size(width, height, format_type="rgba", compressed=False):
    """Estimate VRAM size in bytes."""
    if compressed:
        # BC7/DXT5 compression is roughly 4:1
        return (width * height * 4) // 4
    else:
        # RGBA = 4 bytes per pixel
        return width * height * 4

def scan_ego_resources(project_root):
    """Scan ego directory for all texture resources."""
    ego_path = Path(project_root) / "ego"
    
    if not ego_path.exists():
        print("ERROR: ego/ directory not found")
        return {}
    
    textures = []
    
    # Find all .import files
    for import_file in ego_path.rglob("*.import"):
        # Get the original PNG path
        png_path = import_file.with_suffix('.png')
        
        if png_path.exists():
            # Get file size
            file_size = png_path.stat().st_size
            
            # Try to get dimensions from the import file
            dimensions, format_type, compressed = parse_import_file(import_file)
            
            # Estimate VRAM
            if dimensions:
                vram_bytes = estimate_vram_size(*dimensions, format_type, compressed)
            else:
                # Fallback: estimate from file size (compressed textures are smaller)
                vram_bytes = file_size * 2  # Rough estimate
            
            textures.append({
                'path': str(png_path.relative_to(project_root)),
                'file_size': file_size,
                'dimensions': dimensions,
                'vram_estimate': vram_bytes,
                'compressed': compressed
            })
    
    return textures

def analyze_scene_resources(project_root, scene_path):
    """Analyze a scene file to count external texture resources."""
    scene_file = Path(project_root) / scene_path
    
    if not scene_file.exists():
        return []
    
    textures = []
    
    with open(scene_file, 'r') as f:
        content = f.read()
    
    # Find all ext_resource entries for Texture2D
    pattern = r'\[ext_resource type="Texture2D".*?path="([^"]+)"'
    matches = re.findall(pattern, content)
    
    for texture_path in matches:
        full_path = Path(project_root) / texture_path
        if full_path.exists():
            textures.append(str(texture_path))
    
    return textures

def main():
    project_root = Path("/home/noti/dev/ai-game-2")
    
    print("=" * 70)
    print("VRAM USAGE ESTIMATOR (Static Analysis)")
    print("=" * 70)
    print()
    
    # Analyze Ego resources
    print("Scanning ego/ directory...")
    ego_textures = scan_ego_resources(project_root)
    
    if not ego_textures:
        print("No textures found in ego/")
        return
    
    # Group by animation set
    animation_sets = defaultdict(list)
    for tex in ego_textures:
        if 'minstrel-v3' in tex['path']:
            animation_sets['minstrel-v3'].append(tex)
        elif 'minstrel-v5' in tex['path']:
            animation_sets['minstrel-v5'].append(tex)
        elif 'rosella' in tex['path']:
            animation_sets['rosella'].append(tex)
        elif 'portrait' in tex['path']:
            animation_sets['portrait'].append(tex)
        elif 'mask' in tex['path']:
            animation_sets['masks'].append(tex)
        else:
            animation_sets['other'].append(tex)
    
    print(f"\nFound {len(ego_textures)} texture files in ego/")
    print()
    
    # Calculate totals
    total_disk = 0
    total_vram_compressed = 0
    total_vram_uncompressed = 0
    
    print("Breakdown by Animation Set:")
    print("-" * 70)
    
    for set_name, textures in sorted(animation_sets.items()):
        if not textures:
            continue
            
        disk_sum = sum(t['file_size'] for t in textures)
        vram_compressed_sum = sum(t['vram_estimate'] for t in textures)
        vram_uncompressed_sum = sum(t['vram_estimate'] * 4 for t in textures)
        
        total_disk += disk_sum
        total_vram_compressed += vram_compressed_sum
        total_vram_uncompressed += vram_uncompressed_sum
        
        print(f"\n{set_name.upper()}:")
        print(f"  Files: {len(textures)}")
        print(f"  Disk Size: {disk_sum / (1024*1024):.2f} MB")
        print(f"  VRAM (compressed): {vram_compressed_sum / (1024*1024):.2f} MB")
        print(f"  VRAM (uncompressed): {vram_uncompressed_sum / (1024*1024):.2f} MB")
    
    print()
    print("=" * 70)
    print("TOTALS:")
    print("=" * 70)
    print(f"Total Files: {len(ego_textures)}")
    print(f"Total Disk Size: {total_disk / (1024*1024):.2f} MB")
    print(f"Total VRAM (compressed): {total_vram_compressed / (1024*1024):.2f} MB")
    print(f"Total VRAM (uncompressed): {total_vram_uncompressed / (1024*1024):.2f} MB")
    print()
    
    # Analyze main scene
    print("Analyzing Game.tscn...")
    game_textures = analyze_scene_resources(project_root, "Game.tscn")
    print(f"  External textures referenced: {len(game_textures)}")
    
    # Analyze Ego scene
    print("Analyzing Ego.tscn...")
    ego_scene_textures = analyze_scene_resources(project_root, "Ego.tscn")
    print(f"  External textures referenced: {len(ego_scene_textures)}")
    
    # Count unique minstrel-v3 frames
    minstrel_v3_count = len([t for t in ego_textures if 'minstrel-v3' in t['path']])
    print(f"\nMinstrel-v3 frames loaded at startup: {minstrel_v3_count}")
    
    # Estimate startup VRAM
    minstrel_v3_vram = sum(t['vram_estimate'] for t in ego_textures if 'minstrel-v3' in t['path'])
    other_vram = total_vram_compressed - minstrel_v3_vram
    
    print()
    print("=" * 70)
    print("STARTUP VRAM ESTIMATE (Compressed):")
    print("=" * 70)
    print(f"Ego (minstrel-v3 only): {minstrel_v3_vram / (1024*1024):.2f} MB")
    print(f"Other ego resources: {other_vram / (1024*1024):.2f} MB")
    print(f"Estimated room textures: ~5.00 MB (estimated)")
    print(f"Estimated UI/cursors: ~1.00 MB (estimated)")
    print(f"----------------------------------------")
    print(f"TOTAL STARTUP VRAM: ~{(minstrel_v3_vram + other_vram) / (1024*1024) + 6:.2f} MB")
    print()
    print("UNCOMPRESSED estimate: ~{(minstrel_v3_vram + other_vram) * 4 / (1024*1024) + 24:.2f} MB")
    print()
    print("Note: These are estimates. Actual VRAM may vary based on:")
    print("  - Godot's texture compression settings")
    print("  - GPU driver optimizations")
    print("  - Texture mipmap generation")
    print("=" * 70)

if __name__ == "__main__":
    main()