#include "torlinkc/engine/torrent_engine.hpp"

#include <limits>
#include <unordered_set>

#include <libtorrent/add_torrent_params.hpp>
#include <libtorrent/alert_types.hpp>
#include <libtorrent/error_code.hpp>
#include <libtorrent/load_torrent.hpp>
#include <libtorrent/magnet_uri.hpp>
#include <libtorrent/settings_pack.hpp>
#include <libtorrent/torrent_info.hpp>
#include <libtorrent/torrent_status.hpp>
#include <libtorrent/write_resume_data.hpp>

namespace torlinkc {

namespace {

bool startsWith(const std::string& s, const std::string& prefix) {
  return s.size() >= prefix.size() && s.compare(0, prefix.size(), prefix) == 0;
}

std::string errorMessage(const lt::error_code& ec) { return ec.message(); }

}  // namespace

TorrentEngine::TorrentEngine() {
  lt::settings_pack pack;
  pack.set_int(lt::settings_pack::alert_mask,
               lt::alert_category::status | lt::alert_category::error);
  session_ = std::make_unique<lt::session>(pack);
}

TorrentEngine::~TorrentEngine() { destroy(); }

void TorrentEngine::removeEntry(const std::string& id) {
  auto it = torrents_.find(id);
  if (it == torrents_.end()) return;
  infoHashToId_.erase(it->second.handle.info_hash());
  if (session_) session_->remove_torrent(it->second.handle);
  torrents_.erase(it);
}

void TorrentEngine::add(const std::string& id, const std::string& source, const std::string& dir,
                         AddHandlers handlers, const std::vector<std::string>& announce) {
  removeEntry(id);

  lt::error_code ec;
  lt::add_torrent_params params;
  if (startsWith(source, "magnet:")) {
    params = lt::parse_magnet_uri(source, ec);
  } else {
    // Everything that isn't a magnet URI is a path to a cached .torrent file
    // (the only other thing DownloadQueue ever hands us -- see startSeeding).
    try {
      params = lt::load_torrent_file(source);
    } catch (const lt::system_error& e) {
      ec = e.code();
    } catch (const std::exception& e) {
      handlers.onError ? handlers.onError(e.what()) : void();
      return;
    }
  }
  if (ec) {
    if (handlers.onError) handlers.onError(errorMessage(ec));
    return;
  }

  params.save_path = dir;
  if (!announce.empty()) {
    std::unordered_set<std::string> seen(params.trackers.begin(), params.trackers.end());
    for (const auto& t : announce) {
      if (seen.insert(t).second) params.trackers.push_back(t);
    }
  }

  lt::torrent_handle handle = session_->add_torrent(std::move(params), ec);
  if (ec) {
    if (handlers.onError) handlers.onError(errorMessage(ec));
    return;
  }

  infoHashToId_[handle.info_hash()] = id;
  torrents_[id] = Entry{handle, std::move(handlers)};
}

void TorrentEngine::pollAlerts() {
  if (!session_) return;
  std::vector<lt::alert*> alerts;
  session_->pop_alerts(&alerts);

  for (lt::alert* a : alerts) {
    lt::torrent_handle handle;
    if (auto* mt = lt::alert_cast<lt::metadata_received_alert>(a)) {
      handle = mt->handle;
    } else if (auto* dt = lt::alert_cast<lt::torrent_finished_alert>(a)) {
      handle = dt->handle;
    } else if (auto* et = lt::alert_cast<lt::torrent_error_alert>(a)) {
      handle = et->handle;
    } else {
      continue;
    }

    auto idIt = infoHashToId_.find(handle.info_hash());
    if (idIt == infoHashToId_.end()) continue;
    auto entryIt = torrents_.find(idIt->second);
    if (entryIt == torrents_.end()) continue;
    const AddHandlers& h = entryIt->second.handlers;

    if (lt::alert_cast<lt::metadata_received_alert>(a)) {
      if (!h.onMetadata) continue;
      TorrentMeta meta;
      auto ti = handle.torrent_file();
      if (ti) {
        meta.name = ti->name();
        meta.total = ti->total_size();
        meta.files = ti->num_files();
        lt::add_torrent_params atp;
        atp.ti = std::const_pointer_cast<lt::torrent_info>(ti);
        try {
          std::vector<char> buf = lt::write_torrent_file_buf(atp, {});
          meta.torrentFile = std::string(buf.begin(), buf.end());
        } catch (const std::exception&) {
          // No piece layers / incomplete info: skip caching the raw bytes,
          // the rest of the metadata is still valid.
        }
      }
      h.onMetadata(meta);
    } else if (lt::alert_cast<lt::torrent_finished_alert>(a)) {
      if (h.onDone) h.onDone();
    } else if (auto* et = lt::alert_cast<lt::torrent_error_alert>(a)) {
      if (h.onError) h.onError(et->message());
    }
  }
}

std::optional<TorrentProgress> TorrentEngine::stats(const std::string& id) {
  auto it = torrents_.find(id);
  if (it == torrents_.end()) return std::nullopt;
  if (!it->second.handle.is_valid()) return std::nullopt;

  lt::torrent_status st = it->second.handle.status();

  TorrentProgress p;
  p.progress = st.progress;
  p.downloaded = st.total_wanted_done;
  p.total = st.total_wanted;
  p.speed = st.download_rate;
  p.uploadSpeed = st.upload_rate;
  p.uploaded = st.total_upload;
  p.peers = st.num_peers;
  p.name = st.name;
  p.timeRemaining = std::numeric_limits<double>::infinity();
  if (st.download_rate > 0 && st.total_wanted > st.total_wanted_done) {
    p.timeRemaining =
        static_cast<double>(st.total_wanted - st.total_wanted_done) / st.download_rate * 1000.0;
  }
  return p;
}

void TorrentEngine::remove(const std::string& id) { removeEntry(id); }

void TorrentEngine::destroy() {
  torrents_.clear();
  infoHashToId_.clear();
  session_.reset();
}

int TorrentEngine::listenPort() const { return session_ ? session_->listen_port() : 0; }

}  // namespace torlinkc
