#include <raylib.h>
#include <raymath.h>

#define RAYGUI_IMPLEMENTATION
#include "includes/raygui.h"

#include <iostream>
#include <vector>
#include <fstream>
#include <unordered_map>
#include <cstdint>
#include <cstring>
#include <thread>
#include <chrono>
#include "PlayerController.hpp"
#include "net/NetworkManager.hpp"

constexpr int WORLD_SIZE = 100;
constexpr float WORLD_SCALE = 10.0f;
constexpr float MOVE_SPEED = 15.0f;

struct Heightmap {
	std::vector<float> data;
	int size{0};

	float getHeight(float x, float z) const {
		auto hx = static_cast<int>((x / WORLD_SIZE + 0.5f) * (size - 1));
		auto hz = static_cast<int>((z / WORLD_SIZE + 0.5f) * (size - 1));

		if (hx < 0 || hx >= size || hz < 0 || hz >= size) return 0.0f;
		return data[hz * size + hx];
	}

	bool load(const std::string& filename) {
		std::ifstream file(filename, std::ios::binary);
		if (!file) return false;

		int32_t fileSize;
		file.read(reinterpret_cast<char*>(&fileSize), sizeof(fileSize));
		size = fileSize;

		data.resize(size * size);
		file.read(reinterpret_cast<char*>(data.data()), data.size() * sizeof(float));
		return true;
	}

	Mesh generateMesh() const {
		auto mesh = Mesh{};
		int vertexCount = size * size;
		int triangleCount = (size - 1) * (size - 1) * 2;

		mesh.vertexCount = vertexCount;
		mesh.triangleCount = triangleCount;
		mesh.vertices = (float*)MemAlloc(vertexCount * 3 * sizeof(float));
		mesh.texcoords = (float*)MemAlloc(vertexCount * 2 * sizeof(float));
		mesh.normals = (float*)MemAlloc(vertexCount * 3 * sizeof(float));
		mesh.indices = (unsigned short*)MemAlloc(triangleCount * 3 * sizeof(unsigned short));

		// Generate vertices
		for (int z = 0; z < size; z++) {
			for (int x = 0; x < size; x++) {
				int idx = z * size + x;
				mesh.vertices[idx * 3] = (float(x) / (size - 1) - 0.5f) * WORLD_SIZE;
				mesh.vertices[idx * 3 + 1] = data[idx];
				mesh.vertices[idx * 3 + 2] = (float(z) / (size - 1) - 0.5f) * WORLD_SIZE;

				mesh.texcoords[idx * 2] = static_cast<float>(x) / size;
				mesh.texcoords[idx * 2 + 1] = static_cast<float>(z) / size;
			}
		}

		// Generate indices
		int triIdx = 0;
		for (int z = 0; z < size - 1; z++) {
			for (int x = 0; x < size - 1; x++) {
				int topLeft = z * size + x;
				int topRight = topLeft + 1;
				int bottomLeft = (z + 1) * size + x;
				int bottomRight = bottomLeft + 1;

				mesh.indices[triIdx++] = topLeft;
				mesh.indices[triIdx++] = bottomLeft;
				mesh.indices[triIdx++] = topRight;

				mesh.indices[triIdx++] = topRight;
				mesh.indices[triIdx++] = bottomLeft;
				mesh.indices[triIdx++] = bottomRight;
			}
		}

		// Calculate normals
		for (int i = 0; i < vertexCount; i++) {
			mesh.normals[i * 3] = 0;
			mesh.normals[i * 3 + 1] = 1;
			mesh.normals[i * 3 + 2] = 0;
		}

		UploadMesh(&mesh, false);
		return mesh;
	}
};

enum GameState {
	STATE_LOGIN,
	STATE_CONNECTING,
	STATE_PLAYING
};

class Game {
	PlayerController playerController;
	Model terrainModel;
	Model playerModel;
	Heightmap heightmap;
	NetworkManager network;
	Vector3 playerPos{0, 0, 0};
	Texture2D terrainTexture;
	std::unordered_map<std::string, Texture2D> playerTextures;
	std::unordered_map<uint32_t, Model> remotePlayerModels;

	// Login UI state
	GameState gameState = STATE_LOGIN;
	char usernameBuffer[32] = "";
	bool editMode = false;
	std::string loginError = "";
	std::string currentUsername = "";

public:
	Game() {
		InitWindow(1280, 720, "Multiplayer Terrain Game");
		SetTargetFPS(60);

		// Load heightmap
		if (!heightmap.load("../assets/heightmap.bin")) {
			std::cerr << "Failed to load heightmap\n";
		}

		// Load textures
		terrainTexture = LoadTexture("../assets/textures/black.png");

		// Load all player color textures
		playerTextures["red"] = LoadTexture("../assets/textures/red.png");
		playerTextures["green"] = LoadTexture("../assets/textures/green.png");
		playerTextures["orange"] = LoadTexture("../assets/textures/orange.png");
		playerTextures["purple"] = LoadTexture("../assets/textures/purple.png");
		playerTextures["white"] = LoadTexture("../assets/textures/white.png");

		// Create terrain model
		auto terrainMesh = heightmap.generateMesh();
		terrainModel = LoadModelFromMesh(terrainMesh);
		terrainModel.materials[0].maps[MATERIAL_MAP_DIFFUSE].texture = terrainTexture;

		// Create player cube (texture will be set when we know our color)
		auto cubeMesh = GenMeshCube(1.0f, 2.0f, 1.0f);
		playerModel = LoadModelFromMesh(cubeMesh);
	}

	~Game() {
		// Send logout if we're connected
		if (gameState == STATE_PLAYING && network.isConnected()) {
			network.sendLogout();
		}
		
		UnloadTexture(terrainTexture);
		for (auto& [color, texture] : playerTextures) {
			UnloadTexture(texture);
		}
		UnloadModel(terrainModel);
		UnloadModel(playerModel);
		for (auto& [id, model] : remotePlayerModels) {
			UnloadModel(model);
		}
		CloseWindow();
	}

	void run() {
		while (!WindowShouldClose()) {
			update();
			render();
		}
		
		// Clean logout when window is closing
		if (gameState == STATE_PLAYING && network.isConnected()) {
			network.sendLogout();
			// Give the network time to send the packet
			std::this_thread::sleep_for(std::chrono::milliseconds(100));
		}
	}

private:
	void update() {
		switch (gameState) {
			case STATE_LOGIN:
				// Login screen - no game updates
				break;

			case STATE_CONNECTING:
				// Check if we've received a response from the server
				if (network.isConnected()) {
					gameState = STATE_PLAYING;
					currentUsername = std::string(usernameBuffer);
				} else if (network.hasLoginError()) {
					gameState = STATE_LOGIN;
					loginError = network.getLoginError();
				}
				break;

			case STATE_PLAYING:
				updateGame();
				break;
		}
	}

	void updateGame() {
		if (!network.isConnected()) {
			// Player disconnected, go back to login
			gameState = STATE_LOGIN;
			loginError = "Disconnected from server";
			return;
		}

		// Get server position and update player controller
		playerPos = network.getPosition();
		playerController.setPlayerPosition(playerPos);

		// Set player texture based on assigned color
		if (network.isConnected() && playerTextures.count(network.getPlayerColor()) > 0) {
			playerModel.materials[0].maps[MATERIAL_MAP_DIFFUSE].texture = playerTextures[network.getPlayerColor()];
		}

		// Update remote player models
		auto remotePlayers = network.getRemotePlayers();
		for (const auto& [id, player] : remotePlayers) {
			if (remotePlayerModels.find(id) == remotePlayerModels.end()) {
				// Create new model for this player
				auto cubeMesh = GenMeshCube(1.0f, 2.0f, 1.0f);
				remotePlayerModels[id] = LoadModelFromMesh(cubeMesh);
			}
			// Always update texture in case color changed
			if (playerTextures.count(player.color) > 0) {
				remotePlayerModels[id].materials[0].maps[MATERIAL_MAP_DIFFUSE].texture = playerTextures[player.color];
			}
		}

		// Remove models for players who left
		for (auto it = remotePlayerModels.begin(); it != remotePlayerModels.end();) {
			if (remotePlayers.find(it->first) == remotePlayers.end()) {
				UnloadModel(it->second);
				it = remotePlayerModels.erase(it);
			} else {
				++it;
			}
		}

		// Update player controller (handles camera)
		float deltaTime = GetFrameTime();
		playerController.update(deltaTime);

		// Get movement input from player controller
		Vector3 moveInput = playerController.getMoveInput();

		// Send normalized movement direction to server (server handles speed)
		if (moveInput.x != 0 || moveInput.z != 0) {
			network.sendMove(moveInput.x, 0, moveInput.z);
		}

		// Handle color change with arrow keys
		static int currentColorIndex = -1;
		if (IsKeyPressed(KEY_LEFT) || IsKeyPressed(KEY_RIGHT)) {
			// Get current color index if not set
			if (currentColorIndex == -1) {
				auto currentColor = network.getPlayerColor();
				for (size_t i = 0; i < NetworkManager::AVAILABLE_COLORS.size(); i++) {
					if (NetworkManager::AVAILABLE_COLORS[i] == currentColor) {
						currentColorIndex = i;
						break;
					}
				}
				if (currentColorIndex == -1) currentColorIndex = 0;
			}

			// Change color index
			if (IsKeyPressed(KEY_LEFT)) {
				currentColorIndex--;
				if (currentColorIndex < 0) {
					currentColorIndex = NetworkManager::AVAILABLE_COLORS.size() - 1;
				}
			} else if (IsKeyPressed(KEY_RIGHT)) {
				currentColorIndex++;
				if (currentColorIndex >= (int)NetworkManager::AVAILABLE_COLORS.size()) {
					currentColorIndex = 0;
				}
			}

			// Send color change to server
			network.sendColorChange(NetworkManager::AVAILABLE_COLORS[currentColorIndex]);
		}

		// Handle logout
		if (IsKeyPressed(KEY_ESCAPE)) {
			network.sendLogout();
			gameState = STATE_LOGIN;
			loginError = "";
		}
	}

	void render() {
		BeginDrawing();
		ClearBackground(SKYBLUE);

		switch (gameState) {
			case STATE_LOGIN:
				renderLoginScreen();
				break;

			case STATE_CONNECTING:
				renderConnectingScreen();
				break;

			case STATE_PLAYING:
				renderGame();
				break;
		}

		EndDrawing();
	}

	void renderLoginScreen() {
		// Center the login box
		int boxWidth = 300;
		int boxHeight = 200;
		int boxX = (GetScreenWidth() - boxWidth) / 2;
		int boxY = (GetScreenHeight() - boxHeight) / 2;

		// Draw login panel
		GuiPanel((Rectangle){(float)boxX, (float)boxY, (float)boxWidth, (float)boxHeight}, "Login");

		// Username label and text box
		GuiLabel((Rectangle){(float)(boxX + 20), (float)(boxY + 50), 80, 30}, "Username:");
		if (GuiTextBox((Rectangle){(float)(boxX + 100), (float)(boxY + 50), 180, 30},
		               usernameBuffer, 32, editMode)) {
			editMode = !editMode;
		}

		// Login button
		if (GuiButton((Rectangle){(float)(boxX + 100), (float)(boxY + 100), 100, 30}, "Login")) {
			if (strlen(usernameBuffer) > 0) {
				// Send login with username
				network.sendLoginWithUsername(usernameBuffer);
				gameState = STATE_CONNECTING;
				loginError = "";
			} else {
				loginError = "Please enter a username";
			}
		}

		// Error message
		if (!loginError.empty()) {
			DrawText(loginError.c_str(), boxX + 20, boxY + 150, 20, RED);
		}

		// Instructions
		DrawText("Enter your username to join the game", 10, 10, 20, WHITE);
	}

	void renderConnectingScreen() {
		DrawText("Connecting to server...",
		         GetScreenWidth()/2 - 100, GetScreenHeight()/2 - 10, 20, WHITE);
	}

	void renderGame() {
		BeginMode3D(playerController.getCamera());

		// Draw terrain
		DrawModel(terrainModel, {0, 0, 0}, 1.0f, WHITE);

		// Draw player
		if (network.isConnected()) {
			DrawModel(playerModel, playerPos, 1.0f, WHITE);
		}

		// Draw remote players
		auto remotePlayers = network.getRemotePlayers();
		for (const auto& [id, player] : remotePlayers) {
			if (remotePlayerModels.find(id) != remotePlayerModels.end()) {
				DrawModel(remotePlayerModels[id], player.position, 1.0f, WHITE);
			}
		}

		EndMode3D();

		// UI
		DrawText(TextFormat("Logged in as: %s", currentUsername.c_str()), 10, 10, 20, WHITE);
		DrawText("WASD: Move | Q/E: Strafe | Right-Click: Rotate Camera", 10, 35, 20, WHITE);
		DrawText("Left/Right Arrow: Change Color | Mouse Wheel: Zoom | ESC: Logout", 10, 60, 20, WHITE);
		if (network.isConnected()) {
			std::string colorText = "Your color: " + network.getPlayerColor();
			DrawText(colorText.c_str(), 10, 85, 20, WHITE);
		}
		DrawFPS(10, 110);
	}
};

int main() {
	Game game;
	game.run();
	return 0;
}