package main

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"image"
	"image/color"
	"image/png"
	"math"
	"math/rand"
	"net"
	"os"
	"sync"
	"time"
	"slices"
)

type Vec3 struct {
	X, Y, Z float32
}

type Player struct {
	ID       uint32
	Position Vec3
	Velocity Vec3
	Color    string
	Address  *net.UDPAddr
	LastSeen time.Time
}

type GameServer struct {
	conn      *net.UDPConn
	players   map[uint32]*Player
	heightmap [][]float32
	mutex     sync.RWMutex
	nextID    uint32
}

const (
	MSG_LOGIN        = 0x01
	MSG_POSITION     = 0x02
	MSG_SPAWN        = 0x03
	MSG_MOVE         = 0x04
	MSG_UPDATE       = 0x05
	MSG_PLAYER_JOINED = 0x06
	MSG_PLAYER_LEFT  = 0x07
	MSG_PLAYER_LIST  = 0x08
	MSG_CHANGE_COLOR = 0x09
	MSG_COLOR_CHANGED = 0x0A

	WORLD_SIZE   = 100
	WORLD_SCALE  = 10.0
	MOVE_SPEED   = 15.0
	GRAVITY      = -9.8
	PLAYER_HEIGHT = 1.0
)

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

	// Simple perlin-like noise
	for y := 0; y < size; y++ {
		for x := 0; x < size; x++ {
			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 i := 0; i < 3; i++ {
		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) {
	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]
			}
		}
	}

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

	file, _ := os.Create(filename)
	defer file.Close()
	png.Encode(file, img)
}

func saveHeightmapBinary(heightmap [][]float32, filename string) {
	size := len(heightmap)
	file, _ := os.Create(filename)
	defer file.Close()

	binary.Write(file, binary.LittleEndian, int32(size))
	for y := range heightmap {
		for x := range heightmap[y] {
			binary.Write(file, binary.LittleEndian, heightmap[y][x])
		}
	}
}

func (s *GameServer) getHeightAt(x, z float32) float32 {
	// Convert world coords to heightmap coords with bilinear interpolation
	size := float32(len(s.heightmap))
	fx := (x/WORLD_SIZE + 0.5) * (size - 1)
	fz := (z/WORLD_SIZE + 0.5) * (size - 1)

	// Get integer coordinates
	x0 := int(math.Floor(float64(fx)))
	z0 := int(math.Floor(float64(fz)))
	x1 := x0 + 1
	z1 := z0 + 1

	// Clamp to bounds
	if x0 < 0 || x1 >= len(s.heightmap) || z0 < 0 || z1 >= len(s.heightmap) {
		return 0
	}

	// Get fractional parts
	tx := fx - float32(x0)
	tz := fz - float32(z0)

	// Bilinear interpolation
	h00 := s.heightmap[z0][x0]
	h10 := s.heightmap[z0][x1]
	h01 := s.heightmap[z1][x0]
	h11 := s.heightmap[z1][x1]

	h0 := h00*(1-tx) + h10*tx
	h1 := h01*(1-tx) + h11*tx

	return h0*(1-tz) + h1*tz
}

func (s *GameServer) loadPlayerPositions() {
	data, err := os.ReadFile("players.json")
	if err != nil {
		return
	}

	var savedPlayers map[uint32]Vec3
	json.Unmarshal(data, &savedPlayers)

	for id, pos := range savedPlayers {
		if id > s.nextID {
			s.nextID = id
		}
		s.players[id] = &Player{
			ID:       id,
			Position: pos,
			LastSeen: time.Now(),
		}
	}
}

func (s *GameServer) sendPlayerList(addr *net.UDPAddr, players []*Player) {
	if len(players) == 0 {
		return
	}

	msg := make([]byte, 1024)
	msg[0] = MSG_PLAYER_LIST
	msg[1] = uint8(len(players))

	offset := 2
	for _, p := range players {
		binary.LittleEndian.PutUint32(msg[offset:], p.ID)
		binary.LittleEndian.PutUint32(msg[offset+4:], math.Float32bits(p.Position.X))
		binary.LittleEndian.PutUint32(msg[offset+8:], math.Float32bits(p.Position.Y))
		binary.LittleEndian.PutUint32(msg[offset+12:], math.Float32bits(p.Position.Z))

		colorBytes := []byte(p.Color)
		msg[offset+16] = uint8(len(colorBytes))
		copy(msg[offset+17:], colorBytes)

		offset += 17 + len(colorBytes)
		if offset > 1000 {
			break // Prevent overflow
		}
	}

	s.conn.WriteToUDP(msg[:offset], addr)
}

func (s *GameServer) broadcastPlayerJoined(newPlayer *Player) {
	colorBytes := []byte(newPlayer.Color)
	msg := make([]byte, 18+len(colorBytes))
	msg[0] = MSG_PLAYER_JOINED
	binary.LittleEndian.PutUint32(msg[1:5], newPlayer.ID)
	binary.LittleEndian.PutUint32(msg[5:9], math.Float32bits(newPlayer.Position.X))
	binary.LittleEndian.PutUint32(msg[9:13], math.Float32bits(newPlayer.Position.Y))
	binary.LittleEndian.PutUint32(msg[13:17], math.Float32bits(newPlayer.Position.Z))
	msg[17] = uint8(len(colorBytes))
	copy(msg[18:], colorBytes)

	s.mutex.RLock()
	for _, p := range s.players {
		if p.ID != newPlayer.ID && p.Address != nil {
			s.conn.WriteToUDP(msg, p.Address)
		}
	}
	s.mutex.RUnlock()
}

func (s *GameServer) broadcastPlayerLeft(playerID uint32) {
	msg := make([]byte, 5)
	msg[0] = MSG_PLAYER_LEFT
	binary.LittleEndian.PutUint32(msg[1:5], playerID)

	s.mutex.RLock()
	for _, p := range s.players {
		if p.ID != playerID && p.Address != nil {
			s.conn.WriteToUDP(msg, p.Address)
		}
	}
	s.mutex.RUnlock()
}

func (s *GameServer) broadcastUpdate(player *Player) {
	msg := make([]byte, 17)
	msg[0] = MSG_UPDATE
	binary.LittleEndian.PutUint32(msg[1:5], player.ID)
	binary.LittleEndian.PutUint32(msg[5:9], math.Float32bits(player.Position.X))
	binary.LittleEndian.PutUint32(msg[9:13], math.Float32bits(player.Position.Y))
	binary.LittleEndian.PutUint32(msg[13:17], math.Float32bits(player.Position.Z))

	s.mutex.RLock()
	for _, p := range s.players {
		if p.Address != nil {
			s.conn.WriteToUDP(msg, p.Address)
		}
	}
	s.mutex.RUnlock()
}

func (s *GameServer) handleColorChange(data []byte, addr *net.UDPAddr) {
	if len(data) < 6 {
		return
	}

	playerID := binary.LittleEndian.Uint32(data[1:5])
	colorLen := data[5]

	if len(data) < 6+int(colorLen) {
		return
	}

	newColor := string(data[6 : 6+colorLen])

	// Validate color
	validColors := []string{"red", "green", "orange", "purple", "white"}
	isValid := slices.Contains(validColors, newColor)

	if !isValid {
		return
	}

	s.mutex.Lock()
	player, exists := s.players[playerID]
	if !exists {
		s.mutex.Unlock()
		return
	}

	player.Color = newColor
	s.mutex.Unlock()

	// Broadcast color change to all players
	s.broadcastColorChanged(playerID, newColor)

	fmt.Printf("Player %d changed color to %s\n", playerID, newColor)
}

func (s *GameServer) broadcastColorChanged(playerID uint32, color string) {
	colorBytes := []byte(color)
	msg := make([]byte, 6+len(colorBytes))
	msg[0] = MSG_COLOR_CHANGED
	binary.LittleEndian.PutUint32(msg[1:5], playerID)
	msg[5] = uint8(len(colorBytes))
	copy(msg[6:], colorBytes)

	s.mutex.RLock()
	for _, p := range s.players {
		if p.Address != nil {
			s.conn.WriteToUDP(msg, p.Address)
		}
	}
	s.mutex.RUnlock()
}

func (s *GameServer) savePlayerPositions() {
	s.mutex.RLock()
	savedPlayers := make(map[uint32]Vec3)
	for id, player := range s.players {
		savedPlayers[id] = player.Position
	}
	s.mutex.RUnlock()

	data, _ := json.Marshal(savedPlayers)
	os.WriteFile("players.json", data, 0644)
}

func (s *GameServer) handleLogin(addr *net.UDPAddr) {
	s.mutex.Lock()
	s.nextID++
	playerID := s.nextID

	// Assign color based on player ID to ensure variety
	colors := []string{"red", "green", "orange", "purple", "white"}
	// Cycle through colors based on player ID
	colorIndex := (playerID - 1) % uint32(len(colors))
	color := colors[colorIndex]

	// Spawn at random position on heightmap
	x := rand.Float32() * WORLD_SIZE - WORLD_SIZE/2
	z := rand.Float32() * WORLD_SIZE - WORLD_SIZE/2
	y := s.getHeightAt(x, z) + PLAYER_HEIGHT

	player := &Player{
		ID:       playerID,
		Position: Vec3{x, y, z},
		Color:    color,
		Address:  addr,
		LastSeen: time.Now(),
	}

	// Send existing players to new player
	existingPlayers := make([]*Player, 0)
	for _, p := range s.players {
		if p.ID != playerID {
			existingPlayers = append(existingPlayers, p)
		}
	}

	s.players[playerID] = player
	s.mutex.Unlock()

	// Send spawn message with color
	colorBytes := []byte(color)
	msg := make([]byte, 18+len(colorBytes))
	msg[0] = MSG_SPAWN
	binary.LittleEndian.PutUint32(msg[1:5], playerID)
	binary.LittleEndian.PutUint32(msg[5:9], math.Float32bits(player.Position.X))
	binary.LittleEndian.PutUint32(msg[9:13], math.Float32bits(player.Position.Y))
	binary.LittleEndian.PutUint32(msg[13:17], math.Float32bits(player.Position.Z))
	msg[17] = uint8(len(colorBytes))
	copy(msg[18:], colorBytes)

	s.conn.WriteToUDP(msg, addr)

	// Send player list to new player
	s.sendPlayerList(addr, existingPlayers)

	// Notify other players about new player
	s.broadcastPlayerJoined(player)

	fmt.Printf("Player %d logged in at (%.2f, %.2f, %.2f) with color %s\n", playerID, x, y, z, color)

	s.savePlayerPositions()
}

func (s *GameServer) handleMove(data []byte, addr *net.UDPAddr) {
	if len(data) < 17 {
		return
	}

	playerID := binary.LittleEndian.Uint32(data[1:5])
	dx := math.Float32frombits(binary.LittleEndian.Uint32(data[5:9]))
	// dy := math.Float32frombits(binary.LittleEndian.Uint32(data[9:13])) // Not used - Y position is determined by terrain height
	dz := math.Float32frombits(binary.LittleEndian.Uint32(data[13:17]))

	s.mutex.Lock()
	player, exists := s.players[playerID]
	if !exists {
		s.mutex.Unlock()
		return
	}

	// Server-authoritative movement - server decides the actual speed
	// dx/dz from client are just normalized direction vectors
	deltaTime := float32(0.016) // Assume 60fps for now
	newX := player.Position.X + dx * MOVE_SPEED * deltaTime
	newZ := player.Position.Z + dz * MOVE_SPEED * deltaTime

	// Clamp to world bounds
	newX = float32(math.Max(float64(-WORLD_SIZE/2), math.Min(float64(WORLD_SIZE/2), float64(newX))))
	newZ = float32(math.Max(float64(-WORLD_SIZE/2), math.Min(float64(WORLD_SIZE/2), float64(newZ))))

	// Set Y to terrain height with some smoothing
	targetY := s.getHeightAt(newX, newZ) + PLAYER_HEIGHT
	// Smooth the Y transition
	smoothFactor := float32(0.15) // How quickly to adapt to new height
	newY := player.Position.Y + (targetY - player.Position.Y) * smoothFactor

	player.Position.X = newX
	player.Position.Y = newY
	player.Position.Z = newZ
	player.LastSeen = time.Now()

	s.mutex.Unlock()

	// Broadcast position update to all players
	s.broadcastUpdate(player)
}

func (s *GameServer) run() {
	buffer := make([]byte, 1024)

	// Periodic save
	go func() {
		ticker := time.NewTicker(10 * time.Second)
		for range ticker.C {
			s.savePlayerPositions()
		}
	}()

	for {
		n, addr, err := s.conn.ReadFromUDP(buffer)
		if err != nil {
			continue
		}

		if n < 1 {
			continue
		}

		msgType := buffer[0]

		switch msgType {
		case MSG_LOGIN:
			s.handleLogin(addr)
		case MSG_MOVE:
			s.handleMove(buffer[:n], addr)
		case MSG_CHANGE_COLOR:
			s.handleColorChange(buffer[:n], addr)
		}
	}
}

func main() {
	// Generate and save heightmap
	fmt.Println("Generating heightmap...")
	heightmap := generateHeightmap(WORLD_SIZE)
	saveHeightmapPNG(heightmap, "../assets/heightmap.png")
	saveHeightmapBinary(heightmap, "../assets/heightmap.bin")

	// Start UDP server
	addr, _ := net.ResolveUDPAddr("udp", ":9999")
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		panic(err)
	}
	defer conn.Close()

	server := &GameServer{
		conn:      conn,
		players:   make(map[uint32]*Player),
		heightmap: heightmap,
		nextID:    0,
	}

	server.loadPlayerPositions()

	fmt.Println("Server running on :9999")
	server.run()
}