#include "common.hh"

#include <mutex>
#include <unordered_map>

#include "x.hh"

namespace x {

namespace {

class AtomsImpl : public Atoms {
public:
  explicit AtomsImpl(shared_connection const& conn)
    : conn_(conn) {
  }

  void preload(std::string const& name, bool create) override {
    std::lock_guard<std::mutex> lock(mutex_);
    auto& entry = data_[name];
    if (entry.state_ == UNINITIALIZED) {
      entry.cookie_ = xcb_intern_atom(
          conn_.get(), create ? 0 : 1, name.size(), name.c_str());
      entry.state_ = REQUESTED;
    }
  }

  xcb_atom_t get(std::string const& name, bool create) override {
    std::lock_guard<std::mutex> lock(mutex_);
    return get_locked(name, create);
  }

private:
  enum State {
    UNINITIALIZED = 0,
    REQUESTED,
    RESOLVED,
  };

  struct Entry {
    State state_;
    xcb_atom_t atom_;
    xcb_intern_atom_cookie_t cookie_;

    Entry()
      : state_(UNINITIALIZED) {
    }
  };

  xcb_atom_t get_locked(std::string const& name, bool create) {
    auto& entry = data_[name];
    switch (entry.state_) {
    case UNINITIALIZED:
      entry.cookie_ = xcb_intern_atom(
          conn_.get(), create ? 0 : 1, name.size(), name.c_str());
      entry.state_ = REQUESTED;
      // Fallthrough
    case REQUESTED: {
      auto reply = xcb_intern_atom_reply(conn_.get(), entry.cookie_, nullptr);
      if (reply) {
        entry.atom_ = reply->atom;
        free(reply);
      }
      entry.state_ = RESOLVED;
      // Fallthrough
    }
    case RESOLVED:
      if (create && entry.atom_ == XCB_ATOM_NONE) {
        entry.state_ = UNINITIALIZED;
        return get_locked(name, create);
      }
      break;
    }
    return entry.atom_;
  }

  shared_connection conn_;
  std::mutex mutex_;
  std::unordered_map<std::string, Entry> data_;
};

class EwmhImpl : public Ewmh {
public:
  EwmhImpl(shared_connection const& conn, int screen)
    : conn_(conn), screen_(screen) {
    if (conn_) {
      cookies_ = xcb_ewmh_init_atoms(conn_.get(), &ewmh_);
    } else {
      cookies_ = nullptr;
    }
  }

  ~EwmhImpl() override {
    ensure_cookies();
    if (supported_) xcb_ewmh_get_atoms_reply_wipe(supported_.get());
    if (conn_) xcb_ewmh_connection_wipe(&ewmh_);
  }

  xcb_ewmh_connection_t* conn() override {
    ensure_supported();
    return conn_ ? &ewmh_ : nullptr;
  }

  bool supported(xcb_atom_t atom) override {
    ensure_supported();
    if (!conn_) return false;
    for (uint32_t i = 0; i < supported_->atoms_len; i++) {
      if (supported_->atoms[i] == atom) return true;
    }
    return false;
  }

private:
  void ensure_cookies() {
    if (!cookies_) return;
    if (!xcb_ewmh_init_atoms_replies(&ewmh_, cookies_, nullptr)) {
      conn_.reset();
    }
    cookies_ = nullptr;
  }

  void ensure_supported() {
    if (supported_ || !conn_) return;
    ensure_cookies();
    auto cookie = xcb_ewmh_get_supported(&ewmh_, screen_);
    supported_.reset(new xcb_ewmh_get_atoms_reply_t());
    if (!xcb_ewmh_get_supported_reply(&ewmh_, cookie, supported_.get(),
                                      nullptr)) {
      conn_.reset();
    }
  }

  shared_connection conn_;
  int const screen_;
  xcb_ewmh_connection_t ewmh_;
  xcb_intern_atom_cookie_t* cookies_;
  std::unique_ptr<xcb_ewmh_get_atoms_reply_t> supported_;
};

}  // namespace

// static
Atoms* Atoms::create(shared_connection const& conn) {
  return new AtomsImpl(conn);
}

// static
Ewmh* Ewmh::create(shared_connection const& conn, int screen) {
  return new EwmhImpl(conn, screen);
}

xcb_screen_t const* get_screen(xcb_connection_t* conn, int screen) {
  auto setup = xcb_get_setup(conn);
  xcb_screen_iterator_t iter = xcb_setup_roots_iterator(setup);
  while (screen--) {
    xcb_screen_next(&iter);
  }
  return iter.data;
}

void prefetch_extensions(xcb_connection_t* conn) {
  xcb_prefetch_extension_data(conn, &xcb_render_id);
}

enum PixelFormat {
  ABGR_8888,
  ARGB_8888,
  BGRA_8888,
  RGBA_8888,
};

Format get_best_format(xcb_connection_t* conn, xcb_screen_t const* screen,
                       bool allow_render) {
  assert(conn);
  uint8_t depth = 0;
  xcb_visualtype_t* visual = nullptr;
  xcb_render_pictforminfo_t* render = nullptr;
  xcb_render_pictforminfo_t* render_alpha = nullptr;
  xcb_render_query_pict_formats_reply_t* render_formats = nullptr;
  PixelFormat format;

  auto render_reply = xcb_get_extension_data(conn, &xcb_render_id);
  if (render_reply && allow_render && render_reply->present) {
    auto version = xcb_render_query_version_unchecked(
        conn, XCB_RENDER_MAJOR_VERSION, XCB_RENDER_MINOR_VERSION);
    auto formats = xcb_render_query_pict_formats_unchecked(conn);
    auto version_reply = xcb_render_query_version_reply(conn, version, nullptr);
    if (!version_reply || (version_reply->major_version == 0
                           && version_reply->minor_version < 6)) {
      xcb_discard_reply(conn, formats.sequence);
    } else {
      render_formats = xcb_render_query_pict_formats_reply(
          conn, formats, nullptr);
    }
    free(version_reply);
  }
  auto count = xcb_screen_allowed_depths_length(screen);
  if (count == 0) {
    free(render_formats);
    return Format();
  }
  auto iter = xcb_screen_allowed_depths_iterator(screen);
  while (count--) {
    auto d = iter.data;
    if (d->depth <= depth) continue;
    auto visuals_count = xcb_depth_visuals_length(d);
    if (visuals_count == 0) continue;
    auto v = xcb_depth_visuals(d);
    for (; visuals_count--; ++v) {
      switch (v->_class) {
      case XCB_VISUAL_CLASS_TRUE_COLOR:
      case XCB_VISUAL_CLASS_DIRECT_COLOR:
        break;
      default:
        continue;
      }
      PixelFormat f;
      if (v->bits_per_rgb_value != 8) continue;
      if (v->red_mask == 0xff0000 && v->green_mask == 0xff00
          && v->blue_mask == 0xff) {
        f = ARGB_8888;
      } else if (v->red_mask == 0xff && v->green_mask == 0xff00
                 && v->blue_mask == 0xff0000) {
        f = ABGR_8888;
      } else if (v->red_mask == 0xff000000 && v->green_mask == 0xff0000
                 && v->blue_mask == 0xff00) {
        f = RGBA_8888;
      } else if (v->red_mask == 0xff00 && v->green_mask == 0xff0000
                 && v->blue_mask == 0xff000000) {
        f = BGRA_8888;
      } else {
        continue;
      }
      visual = v;
      format = f;
      depth = d->depth;
      if (!render_formats) break;
      auto info = xcb_render_query_pict_formats_formats(render_formats);
      auto count = xcb_render_query_pict_formats_formats_length(render_formats);
      for (; count--; ++info) {
        if (info->type != XCB_RENDER_PICT_TYPE_DIRECT) continue;
        if (info->depth != d->depth) continue;
        if (info->direct.alpha_mask != 0
            || (static_cast<uint32_t>(info->direct.red_mask)
                << info->direct.red_shift) != visual->red_mask
            || (static_cast<uint32_t>(info->direct.green_mask)
                << info->direct.green_shift) != visual->green_mask
            || (static_cast<uint32_t>(info->direct.blue_mask)
                << info->direct.blue_shift) != visual->blue_mask) continue;
        render = info;
        break;
      }
      if (render) break;
    }
    xcb_depth_next(&iter);
  }
  if (!depth) {
    free(render_formats);
    return Format();
  }
  if (render && render_formats) {
    auto info = xcb_render_query_pict_formats_formats(render_formats);
    auto count = xcb_render_query_pict_formats_formats_length(render_formats);
    for (; count--; ++info) {
      if (info->type != XCB_RENDER_PICT_TYPE_DIRECT) continue;
      if (info->depth != 32) continue;
      if (info->direct.alpha_mask != 0xff
          || info->direct.red_mask != 0xff
          || info->direct.green_mask != 0xff
          || info->direct.blue_mask != 0xff) continue;

      PixelFormat f;
      if (info->direct.red_shift == 0
          && info->direct.green_shift == 8
          && info->direct.blue_shift == 16
          && info->direct.alpha_shift == 24) {
        f = ABGR_8888;
      } else if (info->direct.red_shift == 16
                 && info->direct.green_shift == 8
                 && info->direct.blue_shift == 0
                 && info->direct.alpha_shift == 24) {
        f = ARGB_8888;
      } else if (info->direct.red_shift == 8
                 && info->direct.green_shift == 16
                 && info->direct.blue_shift == 24
                 && info->direct.alpha_shift == 0) {
        f = BGRA_8888;
      } else if (info->direct.red_shift == 24
                 && info->direct.green_shift == 16
                 && info->direct.blue_shift == 8
                 && info->direct.alpha_shift == 0) {
        f = RGBA_8888;
      } else {
        continue;
      }

      if (!render_alpha || f == format) {
        render = info;
      }
    }
    if (!render) {
      render = nullptr;
    }
  }
  // Leak render_formats here to be able to return render and render_alpha
  // as xcb_render_pictforminfo_t pointers
  // free(render_formats);

  return Format(depth, visual, render, render_alpha);
}

}  // namespace x