#!/usr/bin/env python3
"""
Parse Valheim .db world file and extract explored/unexplored map data.
Based on reverse-engineered format from ValheimMapCombiner and valheim-save-tools.
"""
import struct, sys, os
from PIL import Image

DB_PATH = "/tmp/Nerdhalla.db"
OUTPUT_DIR = "/tmp/valheim_map_output"

def read_save_file(path):
    """Parse a Valheim .db file and return structured data."""
    with open(path, 'rb') as f:
        raw = f.read()
    
    offset = 0
    
    # 4 bytes: file length (without SHA512 hash)
    file_len = struct.unpack_from('<I', raw, offset)[0]
    offset += 4
    
    # 20 bytes: data1 (player ID?)
    data1 = raw[offset:offset+20]
    offset += 20
    
    # 4 bytes: world count
    world_count = struct.unpack_from('<I', raw, offset)[0]
    offset += 4
    
    worlds = []
    for _ in range(world_count):
        # 8 bytes: world key (world ID)
        world_key = raw[offset:offset+8]
        offset += 8
        
        # 51 bytes: world data 1
        world_data1 = raw[offset:offset+51]
        offset += 51
        
        # 1 byte: map data bool
        map_data_bool = raw[offset]
        offset += 1
        
        # 4 bytes: map data length
        map_data_len = struct.unpack_from('<I', raw, offset)[0]
        offset += 4
        
        # map data
        map_data_raw = raw[offset:offset+map_data_len]
        offset += map_data_len
        
        map_data = parse_map_data(map_data_raw)
        
        worlds.append({
            'world_key': world_key,
            'world_data1': world_data1,
            'map_data_bool': map_data_bool,
            'map_data_len': map_data_len,
            'map_data': map_data,
        })
    
    # Remaining data (minus 64 bytes SHA512 hash)
    data2 = raw[offset:-64]
    sha512_hash = raw[-64:]
    
    return {
        'file_len': file_len,
        'data1': data1,
        'world_count': world_count,
        'worlds': worlds,
        'data2': data2,
        'sha512_hash': sha512_hash,
    }

def parse_map_data(map_data):
    """Parse the map data section containing explored bitmap and pins."""
    offset = 0
    
    # num1 (int32) - version?
    num1 = struct.unpack_from('<i', map_data, offset)[0]
    offset += 4
    
    # num2 (int32) - ?
    num2 = struct.unpack_from('<i', map_data, offset)[0]
    offset += 4
    
    # Explored data - remaining until pin data
    # The explored bitmap is 4096*4096/8 = 2,097,152 bytes
    # But it's stored as int32 array, so we need to find the pin count
    
    # Read int32s until we find what looks like pin count
    # Pin count is typically a small number (< 1000)
    # The explored bitmap is 4096*4096 bits = 2,097,152 bytes = 524,288 int32s
    
    # Let's read the whole thing as int32 array and find the pin count
    remaining = len(map_data) - offset
    int32_count = remaining // 4
    
    # The explored bitmap is 4096*4096 bits = 2,097,152 bytes
    # But it might be stored differently. Let's read it all and figure it out.
    
    # Actually, the explored data is stored as a bitmap where each bit = one map pixel
    # Valheim map is 4096x4096 = 16,777,216 pixels
    # At 1 bit per pixel = 2,097,152 bytes = 524,288 int32s
    
    # But the actual data might be smaller if the map hasn't been fully explored
    # Let's just read the raw bytes and try to interpret
    
    explored_bytes = map_data[offset:]
    
    return {
        'num1': num1,
        'num2': num2,
        'explored_raw': explored_bytes,
    }

def extract_explored_bitmap(save_data):
    """Extract the explored/unexplored bitmap from parsed save data."""
    if not save_data['worlds']:
        print("No worlds found in save file")
        return None
    
    world = save_data['worlds'][0]
    map_data = world['map_data']
    
    raw = map_data['explored_raw']
    print(f"Explored data size: {len(raw)} bytes")
    
    # Valheim map is 4096x4096 pixels
    # Explored data is a bitmap: 1 bit per pixel
    # Expected size: 4096 * 4096 / 8 = 2,097,152 bytes
    
    # But the data might be stored as int32s (4 bytes per value)
    # Let's check if it looks like a bitmap or int32 array
    
    # Try interpreting as bitmap directly
    expected_bitmap = 4096 * 4096 // 8
    
    if len(raw) >= expected_bitmap:
        print(f"Data is at least {expected_bitmap} bytes - treating as bitmap")
        bitmap = raw[:expected_bitmap]
    else:
        print(f"Data is {len(raw)} bytes, less than expected {expected_bitmap}")
        print("Trying to interpret as int32 array...")
        
        # Read as int32s and see what we get
        vals = struct.unpack_from(f'<{len(raw)//4}I', raw)
        print(f"First 20 int32 values: {vals[:20]}")
        print(f"Max value: {max(vals)}, Min value: {min(vals)}")
        
        # If values are 0 or 1, it might be a per-pixel explored flag
        if max(vals) <= 1:
            print("Values are 0/1 - treating as per-pixel explored flags")
            bitmap_bytes = bytearray()
            for i in range(0, len(vals), 8):
                byte = 0
                for j in range(8):
                    if i + j < len(vals) and vals[i + j]:
                        byte |= (1 << j)
                bitmap_bytes.append(byte)
            bitmap = bytes(bitmap_bytes)
        else:
            # Try treating raw bytes as bitmap directly
            bitmap = raw[:expected_bitmap]
    
    return bitmap

def bitmap_to_image(bitmap, width=4096, height=4096):
    """Convert explored bitmap to a grayscale image (white=explored, black=unexplored)."""
    if bitmap is None:
        return None
    
    img = Image.new('L', (width, height))
    pixels = img.load()
    
    bytes_per_row = (width + 7) // 8
    
    for y in range(height):
        row_start = y * bytes_per_row
        for x in range(width):
            byte_idx = row_start + (x // 8)
            bit_idx = x % 8
            if byte_idx < len(bitmap):
                explored = (bitmap[byte_idx] >> bit_idx) & 1
                pixels[x, y] = 255 if explored else 0
            else:
                pixels[x, y] = 0
    
    return img

def create_fog_overlay(explored_img):
    """Create a fog-of-war overlay: explored=transparent, unexplored=dark overlay."""
    if explored_img is None:
        return None
    
    # Create RGBA image: unexplored areas get a dark overlay
    overlay = Image.new('RGBA', (4096, 4096))
    overlay_pixels = overlay.load()
    explored_pixels = explored_img.load()
    
    for y in range(4096):
        for x in range(4096):
            if explored_pixels[x, y] > 128:
                # Explored - transparent
                overlay_pixels[x, y] = (0, 0, 0, 0)
            else:
                # Unexplored - dark semi-transparent
                overlay_pixels[x, y] = (0, 0, 0, 200)
    
    return overlay

def main():
    print("=" * 60)
    print("Valheim .db File Parser — Explored Map Data")
    print("=" * 60)
    
    print(f"\nReading: {DB_PATH}")
    print(f"File size: {os.path.getsize(DB_PATH)} bytes")
    
    save_data = read_save_file(DB_PATH)
    
    print(f"\nWorld count: {save_data['world_count']}")
    for i, world in enumerate(save_data['worlds']):
        print(f"\nWorld {i}:")
        print(f"  World key: {world['world_key'].hex()}")
        print(f"  Map data bool: {world['map_data_bool']}")
        print(f"  Map data length: {world['map_data_len']}")
        md = world['map_data']
        print(f"  num1: {md['num1']}")
        print(f"  num2: {md['num2']}")
        print(f"  Explored raw size: {len(md['explored_raw'])} bytes")
        
        # Show first few bytes of explored data
        print(f"  First 32 bytes: {md['explored_raw'][:32].hex()}")
    
    # Extract explored bitmap
    print("\n" + "=" * 60)
    print("Extracting explored bitmap...")
    bitmap = extract_explored_bitmap(save_data)
    
    if bitmap:
        print(f"\nBitmap size: {len(bitmap)} bytes")
        print(f"Expected for 4096x4096: {4096*4096//8} bytes")
        
        # Convert to image
        explored_img = bitmap_to_image(bitmap)
        if explored_img:
            explored_path = os.path.join(OUTPUT_DIR, "explored.png")
            explored_img.save(explored_path)
            print(f"\nExplored map saved: {explored_path}")
            
            # Count explored vs unexplored pixels
            explored_pixels = sum(1 for y in range(4096) for x in range(4096) if explored_img.getpixel((x, y)) > 128)
            total = 4096 * 4096
            pct = explored_pixels / total * 100
            print(f"Explored: {explored_pixels:,} / {total:,} pixels ({pct:.1f}%)")
            
            # Create fog overlay
            overlay = create_fog_overlay(explored_img)
            if overlay:
                overlay_path = os.path.join(OUTPUT_DIR, "fog_overlay.png")
                overlay.save(overlay_path)
                print(f"Fog overlay saved: {overlay_path}")
                
                # Also create a composite with the world map
                composite = Image.open(os.path.join(OUTPUT_DIR, "world_composite.png")).convert('RGBA')
                composite.paste(overlay, (0, 0), overlay)
                fogged_path = os.path.join(OUTPUT_DIR, "world_fogged.png")
                composite.save(fogged_path)
                print(f"Fogged world map saved: {fogged_path}")
    else:
        print("Could not extract explored bitmap")
        # Try alternative: read raw bytes and dump structure
        print("\nDumping raw structure for analysis...")
        with open(DB_PATH, 'rb') as f:
            raw = f.read()
        
        # Show first 200 bytes
        print(f"First 200 bytes hex dump:")
        for i in range(0, min(200, len(raw)), 16):
            hex_str = ' '.join(f'{b:02x}' for b in raw[i:i+16])
            ascii_str = ''.join(chr(b) if 32 <= b < 127 else '.' for b in raw[i:i+16])
            print(f"  {i:04x}: {hex_str:48s} {ascii_str}")

if __name__ == '__main__':
    main()
