#include "monitor.hh"

#include "dbus_common.hh"
#include "xcb_atoms.hh"
#include "xcb_colors.hh"
#include "xcb_connection.hh"
#include "xcb_event.hh"
#include "xcb_resource.hh"
#include "xcb_resources.hh"
#include "xcb_xkb.hh"

#include <algorithm>
#include <cassert>
#include <condition_variable>
#include <cstddef>
#include <cstdint>
#include <deque>
#include <iostream>
#include <map>
#include <memory>
#include <mutex>
#include <optional>
#include <sdbus-c++/sdbus-c++.h>
#include <string>
#include <thread>
#include <utility>
#include <vector>
#include <xcb/xcb_event.h>
#include <xcb/xcb_icccm.h>
#include <xcb/xproto.h>

namespace {

constexpr char const kTitle[] = "Claude Code Monitor";
constexpr char const kClass[] = "org.the_jk.claudemon";

class Server {
 public:
  class Delegate {
   public:
    virtual ~Delegate() = default;

    virtual void notify(std::string const& session_id,
                        std::string const& event_name,
                        std::optional<std::string> const& cwd,
                        std::optional<std::string> const& notification_type,
                        std::optional<std::string> const& desktop) = 0;

   protected:
    Delegate() = default;
  };

  explicit Server(Delegate* delegate) : delegate_(delegate) {}

  virtual ~Server() {
    if (conn_)
      conn_->leaveEventLoop();
    thread_.join();
  }

  bool init() {
    try {
      conn_ = sdbus::createBusConnection(dbus::kServiceName);
      object_ = sdbus::createObject(*conn_, dbus::kObjectPath);

      object_
          ->addVTable(sdbus::registerMethod("notify").implementedAs(
              [this](std::string const& session_id,
                     std::string const& event_name, std::string const& cwd,
                     std::string const& notification_type,
                     std::string const& desktop) {
                delegate_->notify(
                    session_id, event_name,
                    cwd.empty() ? std::optional<std::string>() : cwd,
                    notification_type.empty() ? std::optional<std::string>()
                                              : notification_type,
                    desktop.empty() ? std::optional<std::string>() : desktop);
              }))
          .forInterface(dbus::kInterfaceName);

      thread_ = std::thread(&Server::run, this);

      return true;
    } catch (sdbus::Error const& err) {
      std::cerr << "Failed to init server: " << err.what() << '\n';
      return false;
    }
  }

 private:
  void run() { conn_->enterEventLoop(); }

  Delegate* const delegate_;
  std::shared_ptr<sdbus::IConnection> conn_;
  std::unique_ptr<sdbus::IObject> object_;
  std::thread thread_;
};

class Ui : public Server::Delegate {
  enum class State : uint8_t {
    IDLE,
    PROMPT,
    BUSY,
  };

  struct Session {
    State state{State::IDLE};
    std::string cwd{"/"};
    std::string desktop{"0"};
    std::string name;
  };

 public:
  Ui() : server_(std::make_unique<Server>(this)) {}

  bool init() { return server_->init(); }

  int run(std::optional<std::string> display,
          std::optional<std::string> font_name) {
    int screen_index = 0;
    conn_ = xcb::make_shared_conn(
        xcb_connect(display.has_value() ? display.value().c_str() : nullptr,
                    &screen_index));

    {
      auto err = xcb_connection_has_error(conn_.get());
      if (err) {
        std::cerr << "Unable to connect to X display: " << err << '\n';
        return EXIT_FAILURE;
      }
    }

    auto atoms = xcb::Atoms::create(conn_);
    string_atom_ = atoms->get("STRING");
    wm_protocols_ = atoms->get("WM_PROTOCOLS");
    wm_delete_window_ = atoms->get("WM_DELETE_WINDOW");

    auto* screen = xcb::get_screen(conn_.get(), screen_index);
    assert(screen);

    auto colors = xcb::Colors::create(conn_, screen->default_colormap);
    background_ =
        colors->get_with_fallback(0x00, 0x00, 0x00, screen->black_pixel);
    foreground_busy_ =
        colors->get_with_fallback(0x96, 0x96, 0x96, screen->white_pixel);
    foreground_prompt_ =
        colors->get_with_fallback(0x96, 0x00, 0x00, screen->white_pixel);
    foreground_idle_ =
        colors->get_with_fallback(0x00, 0x96, 0x00, screen->white_pixel);

    keyboard_ = xcb::Keyboard::create(conn_.get());
    if (!keyboard_) {
      std::cerr << "Failed to initialize XKB.\n";
      return EXIT_FAILURE;
    }

    wnd_ = xcb::make_unique_wnd(conn_);
    gc_ = xcb::make_unique_gc(conn_);
    font_ = xcb::make_unique_font(conn_);

    if (!font_name.has_value()) {
      auto resources = xcb::Resources::create(conn_);
      font_name = resources->get_string("claudemon.font", "");
      if (!font_name.has_value()) {
        font_name = "fixed";
      }
    }

    auto font_cookie = xcb_open_font_checked(
        conn_.get(), font_->id(), font_name->size(), font_name->data());

    if (!atoms->sync()) {
      std::cerr << "Failed to get X atoms.\n";
      return EXIT_FAILURE;
    }
    if (!colors->sync()) {
      std::cerr << "Failed to get X colors.\n";
      return EXIT_FAILURE;
    }

    uint32_t value_list[3];
    uint32_t value_mask = 0;
    value_mask |= XCB_CW_BACK_PIXEL;
    value_list[0] = screen->black_pixel;
    value_mask |= XCB_CW_EVENT_MASK;
    value_list[1] = XCB_EVENT_MASK_EXPOSURE | XCB_EVENT_MASK_KEY_PRESS |
                    XCB_EVENT_MASK_BUTTON_PRESS |
                    XCB_EVENT_MASK_STRUCTURE_NOTIFY;
    xcb_create_window(conn_.get(), XCB_COPY_FROM_PARENT, wnd_->id(),
                      screen->root, 0, 0, wnd_width_, wnd_height_, 0,
                      XCB_WINDOW_CLASS_INPUT_OUTPUT, screen->root_visual,
                      value_mask, value_list);

    xcb_icccm_set_wm_name(conn_.get(), wnd_->id(), string_atom_->get(), 8,
                          sizeof(kTitle) - 1, kTitle);
    xcb_icccm_set_wm_class(conn_.get(), wnd_->id(), sizeof(kClass) - 1, kClass);
    xcb_atom_t atom_list[1];
    atom_list[0] = wm_delete_window_->get();
    xcb_icccm_set_wm_protocols(conn_.get(), wnd_->id(), wm_protocols_->get(), 1,
                               atom_list);

    value_mask = XCB_GC_FONT;
    value_list[0] = font_->id();
    xcb_create_gc(conn_.get(), gc_->id(), wnd_->id(), value_mask, value_list);

    {
      xcb::generic_error error(xcb_request_check(conn_.get(), font_cookie));
      if (error) {
        std::cerr << "Failed to load font " << *font_name << ": "
                  << static_cast<int>(error->error_code) << '\n';
        return EXIT_FAILURE;
      }
    }

    update_lines();

    xcb_map_window(conn_.get(), wnd_->id());
    xcb_flush(conn_.get());

    std::thread xcb_thread(&Ui::run_xcb, this);

    bool quit = false;
    while (!quit) {
      ActionData action;
      {
        std::unique_lock<std::mutex> lock(action_lock_);
        while (action_.empty()) {
          action_cond_.wait(lock);
        }
        action = action_.front();
        action_.pop_front();
      }

      switch (action.action) {
        case Action::QUIT:
          quit = true;
          break;
        case Action::DRAW:
          if (action.last) {
            draw();
          }
          break;
        case Action::RESIZE:
          wnd_width_ = action.rect.width;
          wnd_height_ = action.rect.height;
          break;
        case Action::NOTIFY: {
          if (action.event_name == "SessionEnd") {
            sessions_.erase(action.session_id);
          } else {
            auto& session = sessions_[action.session_id];
            auto cwd = action.cwd.value_or("/");
            if (session.cwd != cwd) {
              session.cwd = std::move(cwd);
              session.name.clear();
            }
            auto desktop = action.desktop.value_or("0");
            if (session.desktop != desktop) {
              session.desktop = std::move(desktop);
              session.name.clear();
            }
            if (action.event_name == "Notification" &&
                action.notification_type.has_value()) {
              if (action.notification_type.value() == "permission_prompt" ||
                  action.notification_type.value() == "elicitation_dialog") {
                session.state = State::PROMPT;
              } else if (action.notification_type.value() == "idle_prompt") {
                session.state = State::IDLE;
              }
            } else if (action.event_name == "UserPromptSubmit") {
              session.state = State::BUSY;
            }
          }
          update_lines();
          draw();
          break;
        }
        case Action::RESET:
          sessions_.clear();
          update_lines();
          draw();
          break;
      }
    }

    xcb_thread.join();
    return true;
  }

  void notify(std::string const& session_id, std::string const& event_name,
              std::optional<std::string> const& cwd,
              std::optional<std::string> const& notification_type,
              std::optional<std::string> const& desktop) override {
    post(Action::NOTIFY, session_id, event_name, cwd, notification_type,
         desktop);
  }

 private:
  enum class Action : uint8_t {
    QUIT,
    DRAW,
    RESIZE,
    NOTIFY,
    RESET,
  };

  struct Line {
    std::vector<xcb_char2b_t> text;
    xcb_query_text_extents_cookie_t text_extents_cookie;
    uint16_t text_extents_width{0};
    uint16_t text_extents_height{0};
    uint16_t text_extents_offset{0};
    uint32_t color{0};
  };

  struct ActionData {
    Action action;
    xcb_rectangle_t rect;
    bool last;
    std::string session_id;
    std::string event_name;
    std::optional<std::string> cwd;
    std::optional<std::string> notification_type;
    std::optional<std::string> desktop;

    ActionData() = default;

    explicit ActionData(Action action) : action(action) {
      assert(action == Action::QUIT || action == Action::RESET);
    }

    ActionData(Action action, xcb_rectangle_t rect, bool last)
        : action(action), rect(rect), last(last) {
      assert(action == Action::DRAW);
    }

    ActionData(Action action, uint16_t width, uint16_t height)
        : action(action),
          rect({.x = 0, .y = 0, .width = width, .height = height}) {
      assert(action == Action::RESIZE);
    }

    ActionData(Action action, std::string session_id, std::string event_name,
               std::optional<std::string> cwd,
               std::optional<std::string> notification_type,
               std::optional<std::string> desktop)
        : action(action),
          session_id(std::move(session_id)),
          event_name(std::move(event_name)),
          cwd(std::move(cwd)),
          notification_type(std::move(notification_type)),
          desktop(std::move(desktop)) {
      assert(action == Action::NOTIFY);
    }
  };

  void run_xcb() {
    while (true) {
      xcb::generic_event event(xcb_wait_for_event(conn_.get()));
      if (!event) {
        auto err = xcb_connection_has_error(conn_.get());
        if (err) {
          std::cerr << "X connection had fatal error: " << err << '\n';
        } else {
          std::cerr << "X connection had fatal I/O error.\n";
        }
        break;
      }
      auto response_type = XCB_EVENT_RESPONSE_TYPE(event.get());
      if (response_type == XCB_EXPOSE) {
        auto* e = reinterpret_cast<xcb_expose_event_t*>(event.get());
        if (e->window == wnd_->id()) {
          xcb_rectangle_t rect;
          rect.x = static_cast<int16_t>(e->x);
          rect.y = static_cast<int16_t>(e->y);
          rect.width = e->width;
          rect.height = e->height;
          post(Action::DRAW, rect, e->count == 0);
        }
        continue;
      }
      if (response_type == XCB_KEY_PRESS) {
        auto* e = reinterpret_cast<xcb_key_press_event_t*>(event.get());
        if (e->event == wnd_->id()) {
          auto str = keyboard_->get_utf8(e);
          if (str == "q" || str == "\x1b" /* Escape */) {
            break;
          }
          if ((e->state & XCB_MOD_MASK_CONTROL) &&
              str == "\x12" /* Ctrl + R */) {
            post(Action::RESET);
          }
        }
        continue;
      }
      if (response_type == XCB_CONFIGURE_NOTIFY) {
        auto* e = reinterpret_cast<xcb_configure_notify_event_t*>(event.get());
        if (e->window == wnd_->id()) {
          post(Action::RESIZE, e->width, e->height);
        }
        continue;
      }
      if (response_type == XCB_REPARENT_NOTIFY) {
        // Ignored, part of XCB_EVENT_MASK_STRUCTURE_NOTIFY
        continue;
      }
      if (response_type == XCB_MAP_NOTIFY) {
        // Ignored, part of XCB_EVENT_MASK_STRUCTURE_NOTIFY
        continue;
      }
      if (keyboard_->handle_event(conn_.get(), event.get())) {
        continue;
      }
      if (response_type == XCB_CLIENT_MESSAGE) {
        auto* e = reinterpret_cast<xcb_client_message_event_t*>(event.get());
        if (e->window == wnd_->id() && e->type == wm_protocols_->get() &&
            e->format == 32) {
          if (e->data.data32[0] == wm_delete_window_->get()) {
            break;
          }
        }
        continue;
      }

#ifndef NDEBUG
      if (response_type == 0) {
        auto* e = reinterpret_cast<xcb_generic_error_t*>(event.get());
        std::cout << "Unhandled error: "
                  << xcb_event_get_error_label(e->error_code) << '\n';
      } else {
        std::cout << "Unhandled event: " << xcb_event_get_label(response_type)
                  << '\n';
      }
#endif
    }

    post(Action::QUIT);
  }

  template <class... Args>
  void post(Args&&... args) {
    bool notify;
    {
      std::scoped_lock lock(action_lock_);
      notify = action_.empty();
      action_.emplace_back(std::forward<Args>(args)...);
    }
    if (notify)
      action_cond_.notify_one();
  }

  void update_names() {
    if (sessions_.empty())
      return;

    if (std::ranges::all_of(sessions_,
                            [](auto& it) { return !it.second.name.empty(); })) {
      return;
    }

    auto it = sessions_.begin();
    std::string common_path{it->second.cwd};
    for (++it; it != sessions_.end(); ++it) {
      auto& cwd = it->second.cwd;
      size_t len = std::min(common_path.size(), cwd.size());
      size_t i = 0;
      for (; i < len; ++i) {
        if (common_path[i] != cwd[i])
          break;
      }
      common_path.resize(i);
    }

    if (!common_path.empty() && common_path.back() != '/') {
      auto last = common_path.rfind('/');
      if (last != std::string::npos) {
        common_path = common_path.substr(0, last + 1);
      } else {
        common_path = "/";
      }
    }

    for (auto& pair : sessions_) {
      pair.second.name.clear();
      pair.second.name.append(pair.second.desktop);
      pair.second.name.append(": ");
      auto suffix = pair.second.cwd.substr(common_path.size());
      auto end = suffix.find('/');
      if (end != std::string::npos)
        suffix = suffix.substr(0, end);
      pair.second.name.append(suffix);
    }
  }

  void update_lines() {
    update_names();

    std::vector<std::map<std::string, Session>::const_iterator> sort;
    for (auto it = sessions_.begin(); it != sessions_.end(); ++it) {
      sort.emplace_back(it);
    }
    std::ranges::sort(
        sort, [](auto& a, auto& b) { return a->second.name < b->second.name; });

    lines_.resize(sessions_.size());
    for (size_t i = 0; i < sort.size(); ++i) {
      std::string line;
      line.append(sort[i]->second.name);
      line.append(": ");
      switch (sort[i]->second.state) {
        case State::BUSY:
          line.append("busy");
          lines_[i].color = foreground_busy_->get();
          break;
        case State::IDLE:
          line.append("idle");
          lines_[i].color = foreground_idle_->get();
          break;
        case State::PROMPT:
          line.append("prompt");
          lines_[i].color = foreground_prompt_->get();
          break;
      }

      lines_[i].text.resize(line.size());
      for (size_t j = 0; j < line.size(); ++j) {
        lines_[i].text[j].byte1 = 0;
        lines_[i].text[j].byte2 = line[j];
      }

      lines_[i].text_extents_cookie =
          xcb_query_text_extents(conn_.get(), font_->id(),
                                 lines_[i].text.size(), lines_[i].text.data());
      lines_[i].text_extents_width = 0;
    }
  }

  void draw() {
    xcb_rectangle_t r{0, 0, wnd_width_, wnd_height_};

    uint32_t values[2];
    values[0] = background_->get();
    xcb_change_gc(conn_.get(), gc_->id(), XCB_GC_FOREGROUND, values);
    xcb_poly_fill_rectangle(conn_.get(), wnd_->id(), gc_->id(), 1, &r);

    uint16_t tot_height = 0;
    uint16_t max_width = 0;
    uint16_t margin = 2;

    for (auto& line : lines_) {
      if (line.text_extents_width == 0) {
        xcb::reply<xcb_query_text_extents_reply_t> reply(
            xcb_query_text_extents_reply(conn_.get(), line.text_extents_cookie,
                                         nullptr));
        if (reply) {
          line.text_extents_width = reply->overall_width;
          line.text_extents_height = reply->font_ascent + reply->font_descent;
          line.text_extents_offset = reply->font_ascent;
        } else {
          line.text_extents_width = 1;
          line.text_extents_height = 0;
          line.text_extents_offset = 0;
        }
      }
      if (tot_height != 0)
        tot_height += margin;
      tot_height += line.text_extents_height;
      max_width = std::max(line.text_extents_width, max_width);
    }

    auto x = static_cast<int16_t>(
        max_width < wnd_width_ ? (wnd_width_ - max_width) / 2 : 0);
    auto y = static_cast<int16_t>(
        tot_height < wnd_height_ ? (wnd_height_ - tot_height) / 2 : 0);

    for (auto& line : lines_) {
      values[0] = line.color;
      values[1] = background_->get();
      xcb_change_gc(conn_.get(), gc_->id(),
                    XCB_GC_FOREGROUND | XCB_GC_BACKGROUND, values);

      xcb_image_text_16(conn_.get(), line.text.size(), wnd_->id(), gc_->id(), x,
                        static_cast<int16_t>(y + line.text_extents_offset),
                        line.text.data());

      y = static_cast<int16_t>(y + line.text_extents_height + margin);
    }

    xcb_flush(conn_.get());
  }

  xcb::shared_conn conn_;
  std::optional<xcb::Atoms::Reference> string_atom_;
  std::optional<xcb::Atoms::Reference> wm_protocols_;
  std::optional<xcb::Atoms::Reference> wm_delete_window_;
  std::optional<xcb::Colors::Color> background_;
  std::optional<xcb::Colors::Color> foreground_idle_;
  std::optional<xcb::Colors::Color> foreground_busy_;
  std::optional<xcb::Colors::Color> foreground_prompt_;
  std::unique_ptr<xcb::Keyboard> keyboard_;
  xcb::unique_wnd wnd_;
  xcb::unique_gc gc_;
  xcb::unique_font font_;
  uint16_t wnd_width_{100};
  uint16_t wnd_height_{150};

  std::map<std::string, Session> sessions_;

  std::vector<Line> lines_;

  std::mutex action_lock_;
  std::condition_variable action_cond_;
  std::deque<ActionData> action_;

  std::unique_ptr<Server> server_;
};

}  // namespace

int Monitor::run(std::optional<std::string> display,
                 std::optional<std::string> font_name) {
  Ui ui;
  if (!ui.init())
    return EXIT_FAILURE;
  return ui.run(std::move(display), std::move(font_name));
}