package main

import (
	"encoding/binary"
	"flag"
	"fmt"
	"image"
	"image/color"
	"image/png"
	"log"
	"math"
	"math/rand"
	"os"
	"server/net"
)

const WorldSize = 100

func generateHeightmap(size int) [][]float32 {
	heightmap := make([][]float32, size)
	for i := range heightmap {
		heightmap[i] = make([]float32, size)
	}

	// Simple perlin-like noise
	for y := range heightmap {
		for x := range heightmap[y] {
			nx := float64(x) / float64(size) * 4
			ny := float64(y) / float64(size) * 4
			heightmap[y][x] = float32(
				math.Sin(nx*2+rand.Float64())*0.5+
					math.Cos(ny*3+rand.Float64())*0.3+
					rand.Float64()*0.2) * 10.0
		}
	}

	// Smooth the heightmap
	for range 3 {
		newHeightmap := make([][]float32, size)
		for y := range newHeightmap {
			newHeightmap[y] = make([]float32, size)
			for x := range newHeightmap[y] {
				sum := heightmap[y][x]
				count := float32(1)
				for dy := -1; dy <= 1; dy++ {
					for dx := -1; dx <= 1; dx++ {
						nx, ny := x+dx, y+dy
						if nx >= 0 && nx < size && ny >= 0 && ny < size {
							sum += heightmap[ny][nx]
							count++
						}
					}
				}
				newHeightmap[y][x] = sum / count
			}
		}
		heightmap = newHeightmap
	}

	return heightmap
}

func saveHeightmapPNG(heightmap [][]float32, filename string) error {
	size := len(heightmap)
	img := image.NewGray(image.Rect(0, 0, size, size))

	// Find min/max for normalization
	minH, maxH := heightmap[0][0], heightmap[0][0]
	for y := range heightmap {
		for x := range heightmap[y] {
			if heightmap[y][x] < minH {
				minH = heightmap[y][x]
			}
			if heightmap[y][x] > maxH {
				maxH = heightmap[y][x]
			}
		}
	}

	// Apply normalized values to image
	for y := range heightmap {
		for x := range heightmap[y] {
			normalized := (heightmap[y][x] - minH) / (maxH - minH)
			img.SetGray(x, y, color.Gray{Y: uint8(normalized * 255)})
		}
	}

	file, err := os.Create(filename)
	if err != nil {
		return fmt.Errorf("failed to create file %s: %w", filename, err)
	}
	defer file.Close()

	if err := png.Encode(file, img); err != nil {
		return fmt.Errorf("failed to encode PNG: %w", err)
	}
	return nil
}

func saveHeightmapBinary(heightmap [][]float32, filename string) error {
	size := len(heightmap)
	file, err := os.Create(filename)
	if err != nil {
		return fmt.Errorf("failed to create file %s: %w", filename, err)
	}
	defer file.Close()

	if err := binary.Write(file, binary.LittleEndian, int32(size)); err != nil {
		return fmt.Errorf("failed to write size: %w", err)
	}

	for y := range heightmap {
		for x := range heightmap[y] {
			if err := binary.Write(file, binary.LittleEndian, heightmap[y][x]); err != nil {
				return fmt.Errorf("failed to write heightmap data: %w", err)
			}
		}
	}
	return nil
}

// loadHeightmapBinary loads a heightmap from a binary file
func loadHeightmapBinary(filename string) ([][]float32, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, fmt.Errorf("failed to open file %s: %w", filename, err)
	}
	defer file.Close()

	var size int32
	if err := binary.Read(file, binary.LittleEndian, &size); err != nil {
		return nil, fmt.Errorf("failed to read size: %w", err)
	}

	heightmap := make([][]float32, size)
	for y := range heightmap {
		heightmap[y] = make([]float32, size)
		for x := range heightmap[y] {
			if err := binary.Read(file, binary.LittleEndian, &heightmap[y][x]); err != nil {
				return nil, fmt.Errorf("failed to read heightmap data: %w", err)
			}
		}
	}
	return heightmap, nil
}

func main() {
	// Parse command-line flags
	var (
		port        = flag.String("port", "9999", "UDP port to listen on")
		worldSize   = flag.Int("size", WorldSize, "World size for heightmap generation")
		skipGen     = flag.Bool("skip-gen", false, "Skip heightmap generation (fail if none exists)")
		forceGen    = flag.Bool("force-gen", false, "Force heightmap regeneration even if one exists")
		assetsPath  = flag.String("assets", "../assets", "Path to assets directory")
	)
	flag.Parse()

	// Setup logging
	log.SetPrefix("[Game] ")
	log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds)

	var heightmap [][]float32
	binPath := fmt.Sprintf("%s/heightmap.bin", *assetsPath)

	// Check if heightmap exists and should be loaded
	if _, err := os.Stat(binPath); err == nil && !*forceGen && !*skipGen {
		// Heightmap exists and we're not forcing regeneration
		log.Printf("Found existing heightmap at %s, loading...", binPath)
		heightmap, err = loadHeightmapBinary(binPath)
		if err != nil {
			log.Printf("Failed to load existing heightmap: %v, generating new one...", err)
			heightmap = nil
		} else {
			log.Printf("Successfully loaded existing heightmap")
		}
	} else if *forceGen {
		log.Printf("Force regeneration requested, ignoring existing heightmap")
	}

	// Generate new heightmap if needed (not loaded, force generation, or doesn't exist)
	if heightmap == nil && !*skipGen {
		log.Printf("Generating new %dx%d heightmap...", *worldSize, *worldSize)
		heightmap = generateHeightmap(*worldSize)

		pngPath := fmt.Sprintf("%s/heightmap.png", *assetsPath)
		if err := saveHeightmapPNG(heightmap, pngPath); err != nil {
			log.Printf("Warning: Failed to save PNG heightmap: %v", err)
		} else {
			log.Printf("Saved heightmap PNG to %s", pngPath)
		}

		if err := saveHeightmapBinary(heightmap, binPath); err != nil {
			log.Fatalf("Failed to save binary heightmap: %v", err)
		}
		log.Printf("Saved heightmap binary to %s", binPath)
	} else if *skipGen && heightmap == nil {
		// skip-gen was specified but no heightmap exists
		log.Fatalf("No existing heightmap found and generation was skipped (--skip-gen flag)")
	}

	server, err := net.NewServer(*port, heightmap)
	if err != nil {
		log.Fatalf("Failed to create server: %v", err)
	}

	log.Printf("Starting game server on port %s", *port)
	if err := server.Run(); err != nil {
		log.Fatalf("Server failed: %v", err)
	}
}