#include <imgui.h>

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <glm/glm.hpp>
#include <glm/gtc/constants.hpp>
#include <glm/gtc/matrix_transform.hpp>
#include <iostream>
#include <random>
#include <sophia/application.hpp>
#include <sophia/components.hpp>
#include <sophia/entity.hpp>
#include <sophia/event.hpp>
#include <sophia/key_event.hpp>
#include <sophia/scene_view.hpp>
#include <sophia/ui/panel.hpp>
#include <stdexcept>
#include <string>
#include <vector>

struct Boid {
  glm::vec2 position;
  glm::vec2 velocity;
  sophia::Entity entity;
};

// Shows a render target filling the whole panel.
class ViewportPanel : public sophia::Panel {
 public:
  ViewportPanel(sophia::UI& ui, sophia::RenderTargetHandle target)
      : Panel{{.name = "Boids",
               .windowFlags = ImGuiWindowFlags_NoScrollbar |
                              ImGuiWindowFlags_NoScrollWithMouse}},
        ui{ui},
        target{target} {}

  void onUpdate(const sophia::FrameInfo&) override { ui.image(target); }

 private:
  sophia::UI& ui;
  sophia::RenderTargetHandle target;
};

constexpr uint32_t DEFAULT_NUMBER_OF_BOIDS = 25;

class Boids2D : public sophia::Application {
 public:
  static constexpr float SCROLL_SPEED = 200.0f;

  // World is a fixed box independent of the camera; boids wrap at its edges.
  static constexpr float WORLD_HALF_WIDTH = 1000.0f;
  static constexpr float WORLD_HALF_HEIGHT = 1000.0f;

  static constexpr float MAX_SPEED = 120.0f;
  static constexpr float MAX_FORCE = 200.0f;
  static constexpr float PERCEPTION_RADIUS = 150.0f;
  static constexpr float SEPARATION_RADIUS = 40.0f;
  static constexpr float ALIGNMENT_WEIGHT = 1.0f;
  static constexpr float COHESION_WEIGHT = 1.0f;
  static constexpr float SEPARATION_WEIGHT = 1.5f;

  static constexpr float BOID_LENGTH = 20.0f;
  static constexpr float BOID_WIDTH = 12.0f;

  // Neighbor lookups use a uniform grid instead of an O(n^2) scan. Cell size
  // equals PERCEPTION_RADIUS so that a boid's neighbors can only ever fall in
  // its own cell or the 8 cells surrounding it.
  static constexpr float CELL_SIZE = PERCEPTION_RADIUS;
  static constexpr int GRID_WIDTH =
      static_cast<int>(2.0f * WORLD_HALF_WIDTH / CELL_SIZE) + 1;
  static constexpr int GRID_HEIGHT =
      static_cast<int>(2.0f * WORLD_HALF_HEIGHT / CELL_SIZE) + 1;

  Boids2D() : Application{{.width = 750, .height = 1000, .name = "Boids2D"}} {
    uint32_t numberOfBoids = DEFAULT_NUMBER_OF_BOIDS;
    if (args.size() > 1) {
      numberOfBoids = static_cast<uint32_t>(std::stoul(args[1]));
    }
    boids.resize(numberOfBoids);

    sophia::EventSystem::get().addListener<sophia::KeyPressedEvent>(
        [this](sophia::KeyPressedEvent& e) {
          setScrollKeyHeld(e.getKeyCode(), true);
        });
    sophia::EventSystem::get().addListener<sophia::KeyReleasedEvent>(
        [this](sophia::KeyReleasedEvent& e) {
          setScrollKeyHeld(e.getKeyCode(), false);
        });

    debugUI().show();
  }

  void onCreate() override {
    std::cout << "Boids2D created" << std::endl;

    std::mt19937 rng{std::random_device{}()};
    std::uniform_real_distribution<float> posX{-WORLD_HALF_WIDTH,
                                               WORLD_HALF_WIDTH};
    std::uniform_real_distribution<float> posY{-WORLD_HALF_HEIGHT,
                                               WORLD_HALF_HEIGHT};
    std::uniform_real_distribution<float> angle{0.0f, glm::two_pi<float>()};

    for (Boid& boid : boids) {
      boid.position = {posX(rng), posY(rng)};
      float theta = angle(rng);
      boid.velocity = glm::vec2{std::cos(theta), std::sin(theta)} * MAX_SPEED;
    }

    boidTexture = createBoidTexture();

    sceneHandle = scenes().create();
    sophia::Scene& scene = scenes().get(sceneHandle);

    cameraEntity = scene.createCamera();

    std::uniform_real_distribution<float> tint{0.0f, 1.0f};
    for (Boid& boid : boids) {
      boid.entity = scene.createEntity("Boid");
      boid.entity.addComponent<sophia::SpriteComponent>(sophia::SpriteComponent{
          .texture = boidTexture,
          .color = glm::vec4{glm::mix(glm::vec3{0.55f, 0.8f, 1.0f},
                                      glm::vec3{1.0f, 0.85f, 0.55f}, tint(rng)),
                             1.0f},
          .size = {BOID_LENGTH, BOID_WIDTH}});
    }

    target = resources().create<sophia::RenderTarget>();
    sceneViews().create({.scene = sceneHandle,
                         .target = target,
                         .clearColor = {0.01f, 0.01f, 0.02f, 1.0f}});

    panels().emplace<ViewportPanel>(ui(), target);
    panels().setDefaultLayout({.center = {"Boids"}});
  }

  void onUpdate(const sophia::FrameInfo& frameInfo) override {
    float scrollDistance =
        SCROLL_SPEED * static_cast<float>(frameInfo.deltaTime);
    glm::vec3& cameraPosition =
        cameraEntity.getComponent<sophia::TransformComponent>().position;
    if (scrollLeft) cameraPosition.x -= scrollDistance;
    if (scrollRight) cameraPosition.x += scrollDistance;
    if (scrollUp) cameraPosition.y += scrollDistance;
    if (scrollDown) cameraPosition.y -= scrollDistance;

    // One world unit per pixel, whatever size the viewport is.
    cameraEntity.getComponent<sophia::CameraComponent>().ortho.height =
        resources().getExtent(target).y;

    updateBoids(static_cast<float>(frameInfo.deltaTime));

    for (Boid& boid : boids) {
      syncTransform(boid);
    }
  }

 private:
  // Maps a world position to its grid cell, clamped to the grid bounds.
  glm::ivec2 cellOf(glm::vec2 position) const {
    int cx = static_cast<int>((position.x + WORLD_HALF_WIDTH) / CELL_SIZE);
    int cy = static_cast<int>((position.y + WORLD_HALF_HEIGHT) / CELL_SIZE);
    return {std::clamp(cx, 0, GRID_WIDTH - 1),
            std::clamp(cy, 0, GRID_HEIGHT - 1)};
  }

  void rebuildGrid() {
    for (std::vector<uint32_t>& cell : grid) cell.clear();

    for (uint32_t i = 0; i < boids.size(); ++i) {
      glm::ivec2 cell = cellOf(boids[i].position);
      grid[cell.y * GRID_WIDTH + cell.x].push_back(i);
    }
  }

  void updateBoids(float deltaTime) {
    rebuildGrid();

    std::vector<glm::vec2> accelerations(boids.size());

    for (size_t i = 0; i < boids.size(); ++i) {
      glm::vec2 separation{0.0f};
      glm::vec2 averageVelocity{0.0f};
      glm::vec2 averagePosition{0.0f};
      int neighborCount = 0;

      glm::ivec2 cell = cellOf(boids[i].position);

      for (int cy = std::max(cell.y - 1, 0);
           cy <= std::min(cell.y + 1, GRID_HEIGHT - 1); ++cy) {
        for (int cx = std::max(cell.x - 1, 0);
             cx <= std::min(cell.x + 1, GRID_WIDTH - 1); ++cx) {
          for (uint32_t j : grid[cy * GRID_WIDTH + cx]) {
            if (j == i) continue;

            glm::vec2 offset = boids[i].position - boids[j].position;
            float dist = glm::length(offset);
            if (dist >= PERCEPTION_RADIUS) continue;

            averageVelocity += boids[j].velocity;
            averagePosition += boids[j].position;
            ++neighborCount;

            if (dist > 0.0f && dist < SEPARATION_RADIUS) {
              separation += offset / (dist * dist);
            }
          }
        }
      }

      glm::vec2 acceleration =
          steerTowards(boids[i], separation) * SEPARATION_WEIGHT;

      if (neighborCount > 0) {
        averageVelocity /= static_cast<float>(neighborCount);
        acceleration +=
            steerTowards(boids[i], averageVelocity) * ALIGNMENT_WEIGHT;

        averagePosition /= static_cast<float>(neighborCount);
        glm::vec2 towardsGroup = averagePosition - boids[i].position;
        acceleration += steerTowards(boids[i], towardsGroup) * COHESION_WEIGHT;
      }

      accelerations[i] = acceleration;
    }

    for (size_t i = 0; i < boids.size(); ++i) {
      Boid& boid = boids[i];
      boid.velocity += accelerations[i] * deltaTime;

      float speed = glm::length(boid.velocity);
      if (speed > MAX_SPEED) {
        boid.velocity = boid.velocity / speed * MAX_SPEED;
      }

      boid.position += boid.velocity * deltaTime;
      wrapPosition(boid.position);
    }
  }

  // Steers boid towards desiredDirection at up to MAX_SPEED, clamped to
  // MAX_FORCE.
  glm::vec2 steerTowards(const Boid& boid, glm::vec2 desiredDirection) {
    float directionLength = glm::length(desiredDirection);
    if (directionLength < 1e-5f) return glm::vec2{0.0f};

    glm::vec2 desired = desiredDirection / directionLength * MAX_SPEED;
    glm::vec2 steer = desired - boid.velocity;

    float steerLength = glm::length(steer);
    if (steerLength > MAX_FORCE) {
      steer = steer / steerLength * MAX_FORCE;
    }
    return steer;
  }

  void wrapPosition(glm::vec2& position) {
    if (position.x > WORLD_HALF_WIDTH) position.x = -WORLD_HALF_WIDTH;
    if (position.x < -WORLD_HALF_WIDTH) position.x = WORLD_HALF_WIDTH;
    if (position.y > WORLD_HALF_HEIGHT) position.y = -WORLD_HALF_HEIGHT;
    if (position.y < -WORLD_HALF_HEIGHT) position.y = WORLD_HALF_HEIGHT;
  }

  void syncTransform(Boid& boid) {
    glm::vec2 direction = boid.velocity;
    float speed = glm::length(direction);
    direction = speed > 1e-5f ? direction / speed : glm::vec2{0.0f, -1.0f};
    float angle = std::atan2(direction.y, direction.x);

    auto& transform = boid.entity.getComponent<sophia::TransformComponent>();
    transform.position = glm::vec3(boid.position, 0.0f);
    transform.rotation = glm::angleAxis(angle, glm::vec3(0.0f, 0.0f, 1.0f));
  }

  // Boid shape, nose pointing along +u, as an anti-aliased white triangle on
  // a transparent background; each boid's sprite tints it.
  sophia::TextureHandle createBoidTexture() {
    constexpr uint32_t WIDTH = 64;
    constexpr uint32_t HEIGHT = 40;

    std::vector<uint8_t> pixels(WIDTH * HEIGHT * 4);
    for (uint32_t y = 0; y < HEIGHT; ++y) {
      for (uint32_t x = 0; x < WIDTH; ++x) {
        float u = (static_cast<float>(x) + 0.5f) / WIDTH;
        float v = (static_cast<float>(y) + 0.5f) / HEIGHT;
        // Distance inside the triangle's edge, in pixels.
        float inside = (0.5f * (1.0f - u) - std::abs(v - 0.5f)) * HEIGHT;
        float alpha = std::clamp(inside + 0.5f, 0.0f, 1.0f);

        uint8_t* pixel = &pixels[(y * WIDTH + x) * 4];
        pixel[0] = pixel[1] = pixel[2] = 255;
        pixel[3] = static_cast<uint8_t>(alpha * 255.0f);
      }
    }

    return resources().create<sophia::Texture>(WIDTH, HEIGHT,
                                               std::move(pixels));
  }

  void setScrollKeyHeld(sophia::KeyCode keycode, bool held) {
    if (keycode == sophia::Key::Left) scrollLeft = held;
    if (keycode == sophia::Key::Right) scrollRight = held;
    if (keycode == sophia::Key::Up) scrollUp = held;
    if (keycode == sophia::Key::Down) scrollDown = held;
  }

  std::vector<Boid> boids;
  std::vector<std::vector<uint32_t>> grid{
      static_cast<size_t>(GRID_WIDTH * GRID_HEIGHT)};
  sophia::TextureHandle boidTexture{};
  sophia::SceneHandle sceneHandle{};
  sophia::RenderTargetHandle target{};
  sophia::Entity cameraEntity;
  bool scrollLeft = false;
  bool scrollRight = false;
  bool scrollUp = false;
  bool scrollDown = false;
};

SOPHIA_MAIN(Boids2D)