#!/usr/bin/env python3
"""
Decode Valheim map tiles (from valheim-map.world export) into PNG images.
Tile format per sample: uint16_t biome + float height + float forestFactor = 10 bytes.
1024 x 1024 samples per tile, 4x4 tiles = 24000x24000 world.
"""
import struct, gzip, os, json, math
from PIL import Image

TILES_DIR = "/tmp/valheim_tiles/tiles"
OUTPUT_DIR = "/tmp/valheim_map_output"
MAP_JSON_PATH = "/tmp/valheim_map_data/map.json"

# Biome colors (distinct, map-friendly)
BIOME_COLORS = {
    0:   (30, 30, 40),     # None (deep ocean/void)
    1:   (107, 142, 35),   # Meadows - olive green
    2:   (85, 107, 47),    # Swamp - dark olive
    4:   (139, 137, 137),  # Mountain - grey
    8:   (34, 100, 34),    # BlackForest - dark green
    16:  (210, 180, 60),   # Plains - gold
    32:  (139, 69, 19),    # AshLands - dark brown/rust
    64:  (200, 220, 255),  # DeepNorth - ice blue
    256: (25, 60, 120),    # Ocean - deep blue
    512: (120, 80, 150),   # Mistlands - purple
}

# Biome to height-based tint mapping for compositing
BIOME_HEIGHT_KEYS = {
    1:   'meadows',
    2:   'swamp',
    4:   'mountain',
    8:   'blackforest',
    16:  'plains',
    32:  'ashlands',
    64:  'deepnorth',
    256: 'ocean',
    512: 'mistlands',
}

def read_tile(tile_path):
    """Read a compressed tile and return (biomes, heights, forest_factors) as 2D arrays.
    
    Format: flat array of struct { uint16_t biome; float height; float forestFactor; }
    = 10 bytes per sample. No header. 1024*1024 = 1,048,576 samples = 10,485,760 bytes.
    """
    with gzip.open(tile_path, 'rb') as f:
        raw = f.read()
    
    total_samples = 1024 * 1024
    expected = total_samples * 10
    actual = len(raw)
    
    if actual != expected:
        print(f"  NOTE: Got {actual} bytes, expected {expected} (off by {actual - expected})")
    
    biomes = [[0]*1024 for _ in range(1024)]
    heights = [[0.0]*1024 for _ in range(1024)]
    forests = [[0.0]*1024 for _ in range(1024)]
    
    n = min(total_samples, actual // 10)
    for i in range(n):
        offset = i * 10
        biome, height, forest = struct.unpack('<Hff', raw[offset:offset+10])
        row = i // 1024
        col = i % 1024
        biomes[row][col] = biome
        heights[row][col] = height
        forests[row][col] = forest
    
    return biomes, heights, forests

def make_biome_image(tile_biomes):
    """Create a biome classification image (1024x1024 RGB)."""
    img = Image.new('RGB', (1024, 1024))
    pixels = img.load()
    for y in range(1024):
        for x in range(1024):
            biome = tile_biomes[y][x]
            # Find the exact biome (bitfield - multiple can be set)
            # Common Valheim biomes are powers of 2; pick highest
            found = 0
            for b in [512, 256, 64, 32, 16, 8, 4, 2, 1]:
                if biome & b:
                    found = b
                    break
            pixels[x, y] = BIOME_COLORS.get(found, (50, 50, 50))
    return img

def make_height_image(tile_heights):
    """Create a heightmap image (grayscale, normalized to 0-255)."""
    all_heights = [h for row in tile_heights for h in row if not math.isnan(h) and h > -1000]
    if not all_heights:
        return Image.new('L', (1024, 1024), 128)
    h_min, h_max = min(all_heights), max(all_heights)
    h_range = h_max - h_min if h_max > h_min else 1
    
    img = Image.new('L', (1024, 1024))
    pixels = img.load()
    for y in range(1024):
        for x in range(1024):
            h = tile_heights[y][x]
            if math.isnan(h) or h < -1000:
                val = 0
            else:
                val = int((h - h_min) / h_range * 255)
            pixels[x, y] = max(0, min(255, val))
    return img

def make_composite_image(tile_biomes, tile_heights):
    """Create a combined biome+height map using height-shaded biome colors."""
    all_heights = [h for row in tile_heights for h in row if not math.isnan(h) and h > -1000]
    if not all_heights:
        return make_biome_image(tile_biomes)
    h_min, h_max = min(all_heights), max(all_heights)
    h_range = h_max - h_min if h_max > h_min else 1
    
    img = Image.new('RGB', (1024, 1024))
    pixels = img.load()
    for y in range(1024):
        for x in range(1024):
            biome = tile_biomes[y][x]
            h = tile_heights[y][x]
            
            # Find dominant biome
            found = 0
            for b in [512, 256, 64, 32, 16, 8, 4, 2, 1]:
                if biome & b:
                    found = b
                    break
            
            base = BIOME_COLORS.get(found, (50, 50, 50))
            
            # Height shading (darker = lower, lighter = higher)
            if math.isnan(h) or h < -1000:
                shade = 0.5
            else:
                shade = 0.4 + 0.6 * ((h - h_min) / h_range)
            
            r = int(base[0] * shade)
            g = int(base[1] * shade)
            b = int(base[2] * shade)
            pixels[x, y] = (max(0, min(255, r)),
                           max(0, min(255, g)),
                           max(0, min(255, b)))
    return img

def make_biome_and_height_tiles():
    """Process all 16 tiles into PNG images."""
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    
    tile_names = sorted([f for f in os.listdir(TILES_DIR) if f.endswith('.bin.gz')])
    print(f"Found {len(tile_names)} tiles")
    
    for tile_name in tile_names:
        tile_path = os.path.join(TILES_DIR, tile_name)
        base = tile_name.replace('.bin.gz', '')
        print(f"\nDecoding {tile_name}...")
        
        biomes, heights, forests = read_tile(tile_path)
        print(f"  Biomes: {len(set(b for row in biomes for b in row))} unique values")
        
        # Biome image
        biome_img = make_biome_image(biomes)
        biome_path = os.path.join(OUTPUT_DIR, f"{base}_biome.png")
        biome_img.save(biome_path)
        print(f"  Biome: {biome_path}")
        
        # Height image
        height_img = make_height_image(heights)
        height_path = os.path.join(OUTPUT_DIR, f"{base}_height.png")
        height_img.save(height_path)
        print(f"  Height: {height_path}")
        
        # Composite image
        composite_img = make_composite_image(biomes, heights)
        composite_path = os.path.join(OUTPUT_DIR, f"{base}_composite.png")
        composite_img.save(composite_path)
        print(f"  Composite: {composite_path}")
    
    print(f"\nAll tiles processed. Output in {OUTPUT_DIR}/")

def stitch_full_map():
    """Stitch the 4x4 tiles into a full world map."""
    tile_names = sorted([f for f in os.listdir(TILES_DIR) if f.endswith('.bin.gz')])
    # Parse tile coords - they're xx-yy.bin.gz
    tiles = {}
    for tn in tile_names:
        parts = tn.replace('.bin.gz', '').split('-')
        x, y = int(parts[0]), int(parts[1])
        tiles[(x, y)] = tn
    
    full_biome = Image.new('RGB', (4096, 4096))
    full_height = Image.new('L', (4096, 4096))
    full_composite = Image.new('RGB', (4096, 4096))
    
    for (tx, ty), tn in sorted(tiles.items()):
        print(f"Stitching tile {tn} at position ({tx}, {ty})...")
        biomes, heights, _ = read_tile(os.path.join(TILES_DIR, tn))
        
        biome_img = make_biome_image(biomes)
        full_biome.paste(biome_img, (tx * 1024, ty * 1024))
        
        height_img = make_height_image(heights)
        full_height.paste(height_img, (tx * 1024, ty * 1024))
        
        composite_img = make_composite_image(biomes, heights)
        full_composite.paste(composite_img, (tx * 1024, ty * 1024))
    
    world_biome_path = os.path.join(OUTPUT_DIR, "world_biome.png")
    world_height_path = os.path.join(OUTPUT_DIR, "world_height.png")
    world_composite_path = os.path.join(OUTPUT_DIR, "world_composite.png")
    
    full_biome.save(world_biome_path)
    full_height.save(world_height_path)
    full_composite.save(world_composite_path)
    
    print(f"\nFull world maps saved:")
    print(f"  Biome: {world_biome_path}")
    print(f"  Height: {world_height_path}")
    print(f"  Composite: {world_composite_path}")

if __name__ == '__main__':
    print("=" * 60)
    print("Valheim Tile Decoder")
    print("=" * 60)
    print(f"Reading tiles from: {TILES_DIR}")
    
    # First check if PIL is available
    try:
        from PIL import Image
    except ImportError:
        print("ERROR: Pillow not installed. Install with: pip install Pillow")
        sys.exit(1)
    
    # Check first tile structure
    sample_tile = sorted([f for f in os.listdir(TILES_DIR) if f.endswith('.bin.gz')])[0]
    with gzip.open(os.path.join(TILES_DIR, sample_tile), 'rb') as f:
        raw = f.read()
    print(f"Sample tile {sample_tile}: {len(raw)} bytes uncompressed")
    print(f"  Header (first 4 bytes): {raw[:4].hex()} = {struct.unpack('<I', raw[:4])[0]}")
    print(f"  Expected sample size: 10 bytes per pixel × {1024*1024} pixels + 4 header = {(10*1024*1024)+4}")
    print(f"  Actual match: {len(raw) == (10 * 1024 * 1024) + 4}")
    
    # Process individual tile images
    make_biome_and_height_tiles()
    
    # Stitch full world map
    print("\n" + "=" * 60)
    print("Stitching full world map (4096x4096)...")
    print("=" * 60)
    stitch_full_map()
    
    print("\nDone!")