#pragma once

#include "direct_ipu.hpp"

#include <algorithm>
#include <atomic>
#include <cstdint>
#include <cstring>
#include <fstream>
#include <functional>
#include <limits>
#include <sstream>
#include <stdexcept>
#include <string>
#include <vector>

namespace hostexchange {

struct Page {
  unsigned index;
  std::size_t size;
};

struct Slice {
  unsigned page;
  std::size_t pageOffset;
  std::size_t fileOffset;
  std::size_t size;
};

struct Call {
  std::string name;
  std::uint32_t id;
  unsigned phases;
  std::vector<Slice> inputs;
  std::vector<Slice> outputs;
};

struct Protocol {
  unsigned startupMark = 0;
  unsigned commandPage = 0;
  std::size_t commandOffset = 0;
  std::vector<Page> pages;
  std::vector<unsigned> attachOrder;
  std::vector<Call> calls;

  static Protocol read(const std::string &path) {
    std::ifstream input(path);
    if (!input)
      throw std::runtime_error("cannot open host-exchange manifest: " + path);

    Protocol result;
    std::string line;
    unsigned lineNumber = 0;
    bool header = false;
    while (std::getline(input, line)) {
      ++lineNumber;
      const auto comment = line.find('#');
      if (comment != std::string::npos)
        line.erase(comment);
      std::istringstream fields(line);
      std::string directive;
      if (!(fields >> directive))
        continue;
      const auto fail = [&](const std::string &message) {
        throw std::runtime_error(path + ":" + std::to_string(lineNumber) +
                                 ": " + message);
      };
      if (!header) {
        unsigned version;
        if (directive != "IPU-HOST-EXCHANGE" || !(fields >> version) ||
            version != 1)
          fail("expected IPU-HOST-EXCHANGE 1");
        header = true;
        continue;
      }

      if (directive == "startup-mark") {
        if (!(fields >> result.startupMark))
          fail("invalid startup mark");
      } else if (directive == "page") {
        Page page{};
        if (!(fields >> page.index >> page.size) || page.size == 0)
          fail("invalid page");
        result.pages.push_back(page);
      } else if (directive == "attach") {
        unsigned page;
        if (!(fields >> page))
          fail("invalid attach entry");
        result.attachOrder.push_back(page);
      } else if (directive == "command") {
        if (!(fields >> result.commandPage >> result.commandOffset))
          fail("invalid command location");
      } else if (directive == "call") {
        Call call{};
        if (!(fields >> call.name >> call.id >> call.phases) ||
            call.name.empty())
          fail("invalid call");
        result.calls.push_back(std::move(call));
      } else if (directive == "input" || directive == "output") {
        std::string name;
        Slice slice{};
        if (!(fields >> name >> slice.page >> slice.pageOffset >>
              slice.fileOffset >> slice.size) ||
            slice.size == 0)
          fail("invalid data slice");
        auto call = std::find_if(result.calls.begin(), result.calls.end(),
                                 [&](const Call &item) {
                                   return item.name == name;
                                 });
        if (call == result.calls.end())
          fail("slice refers to an unknown call");
        (directive == "input" ? call->inputs : call->outputs)
            .push_back(slice);
      } else {
        fail("unknown directive " + directive);
      }
      std::string trailing;
      if (fields >> trailing)
        fail("unexpected trailing field");
    }
    if (!header)
      throw std::runtime_error("empty host-exchange manifest: " + path);
    result.validate();
    return result;
  }

  const Call &call(const std::string &name) const {
    const auto found = std::find_if(calls.begin(), calls.end(),
                                    [&](const Call &item) {
                                      return item.name == name;
                                    });
    if (found == calls.end())
      throw std::runtime_error("unknown host-exchange call: " + name);
    return *found;
  }

  std::size_t inputSize(const Call &call) const {
    return dataSize(call.inputs);
  }
  std::size_t outputSize(const Call &call) const {
    return dataSize(call.outputs);
  }

private:
  const Page &page(unsigned index) const {
    const auto found = std::find_if(pages.begin(), pages.end(),
                                    [&](const Page &item) {
                                      return item.index == index;
                                    });
    if (found == pages.end())
      throw std::runtime_error("manifest refers to undefined page " +
                               std::to_string(index));
    return *found;
  }

  static std::size_t dataSize(const std::vector<Slice> &slices) {
    std::size_t result = 0;
    for (const auto &slice : slices) {
      if (slice.fileOffset > std::numeric_limits<std::size_t>::max() -
                                 slice.size)
        throw std::runtime_error("host-exchange file slice overflows");
      result = std::max(result, slice.fileOffset + slice.size);
    }
    return result;
  }

  void validate() const {
    if (startupMark == 0 || startupMark > directipu::pci::hspMarkMask)
      throw std::runtime_error("invalid host-exchange startup mark");
    if (pages.empty() || attachOrder.size() != pages.size())
      throw std::runtime_error("host-exchange pages and attach order differ");
    for (std::size_t i = 0; i < pages.size(); ++i) {
      if (std::count_if(pages.begin(), pages.end(), [&](const Page &page) {
            return page.index == pages[i].index;
          }) != 1 ||
          std::count(attachOrder.begin(), attachOrder.end(), pages[i].index) !=
              1)
        throw std::runtime_error("duplicate or missing host-exchange page");
    }
    if (commandOffset > page(commandPage).size ||
        sizeof(std::uint32_t) > page(commandPage).size - commandOffset)
      throw std::runtime_error("host-exchange command location is out of range");
    for (std::size_t i = 0; i < calls.size(); ++i) {
      if (std::count_if(calls.begin(), calls.end(), [&](const Call &call) {
            return call.name == calls[i].name;
          }) != 1)
        throw std::runtime_error("duplicate host-exchange call name");
      for (const auto *slices : {&calls[i].inputs, &calls[i].outputs})
        for (const auto &slice : *slices)
          if (slice.pageOffset > page(slice.page).size ||
              slice.size > page(slice.page).size - slice.pageOffset)
            throw std::runtime_error("host-exchange page slice is out of range");
      (void)inputSize(calls[i]);
      (void)outputSize(calls[i]);
    }
  }
};

class Session {
public:
  Session(directipu::Device &device, Protocol protocol)
      : device_(device), protocol_(std::move(protocol)) {
    for (const auto &spec : protocol_.pages) {
      void *mapping = mmap(nullptr, spec.size, PROT_READ | PROT_WRITE,
                           MAP_PRIVATE | MAP_ANONYMOUS | MAP_POPULATE, -1, 0);
      if (mapping == MAP_FAILED) {
        release();
        throw directipu::systemError("allocate host-exchange page");
      }
      std::memset(mapping, 0, spec.size);
      pages_.push_back({spec, mapping, false});
    }
  }

  Session(const Session &) = delete;
  Session &operator=(const Session &) = delete;
  ~Session() { release(); }

  unsigned startupMark() const { return protocol_.startupMark; }

  void start(const std::function<void()> &releaseLoaderResources) {
    constexpr auto timeout = std::chrono::seconds(10);
    releaseLoaderResources();
    device_.setMark(1);
    device_.waitForMarkRegister(directipu::pci::hspGs2Control, 0, timeout);
    attach();
    device_.writeConfig(directipu::pci::hspGs2Control, 1);
    device_.waitForMarkRegister(directipu::pci::hspGs2Control, 0, timeout);
  }

  void attach() {
    if (attached_)
      throw std::logic_error("host-exchange pages already attached");
    try {
      for (const auto index : protocol_.attachOrder) {
        auto &entry = page(index);
        device_.writeConfig(directipu::pci::exchangeWindowBase,
                            directipu::pci::exchangeWindowHexopt);
        device_.attachBuffer(index, entry.mapping, entry.spec.size);
        entry.attached = true;
      }
      attached_ = true;
    } catch (...) {
      detach();
      throw;
    }
  }

  std::vector<std::uint8_t> invoke(const std::string &name,
                                   const std::vector<std::uint8_t> &input) {
    if (!attached_)
      throw std::logic_error("host-exchange session has not started");
    const auto &call = protocol_.call(name);
    const auto requiredInput = protocol_.inputSize(call);
    if (input.size() != requiredInput)
      throw std::runtime_error(name + " expects " +
                               std::to_string(requiredInput) +
                               " input bytes, got " +
                               std::to_string(input.size()));
    for (const auto &slice : call.inputs)
      std::memcpy(bytes(slice.page) + slice.pageOffset,
                  input.data() + slice.fileOffset, slice.size);
    for (const auto &slice : call.outputs)
      std::memset(bytes(slice.page) + slice.pageOffset, 0xa5, slice.size);
    setCommand(call.id);
    advance(call.phases, name);
    std::vector<std::uint8_t> output(protocol_.outputSize(call));
    std::atomic_thread_fence(std::memory_order_seq_cst);
    for (const auto &slice : call.outputs)
      std::memcpy(output.data() + slice.fileOffset,
                  bytes(slice.page) + slice.pageOffset, slice.size);
    return output;
  }

private:
  struct Mapping {
    Page spec;
    void *mapping;
    bool attached;
  };

  Mapping &page(unsigned index) {
    const auto found = std::find_if(pages_.begin(), pages_.end(),
                                    [&](const Mapping &item) {
                                      return item.spec.index == index;
                                    });
    if (found == pages_.end())
      throw std::logic_error("missing mapped host-exchange page");
    return *found;
  }

  std::uint8_t *bytes(unsigned index) {
    return static_cast<std::uint8_t *>(page(index).mapping);
  }

  void setCommand(std::uint32_t command) {
    std::memcpy(bytes(protocol_.commandPage) + protocol_.commandOffset,
                &command, sizeof(command));
    std::atomic_thread_fence(std::memory_order_seq_cst);
  }

  void advance(unsigned count, const std::string &name) {
    constexpr auto timeout = std::chrono::seconds(10);
    for (unsigned phase = 0; phase < count; ++phase) {
      try {
        device_.waitForMarkRegister(directipu::pci::hspGs2Control, 0,
                                    timeout);
        device_.writeConfig(directipu::pci::hspGs2Control, 1);
        device_.waitForMarkRegister(directipu::pci::hspGs2Control, 0,
                                    timeout);
      } catch (const std::exception &error) {
        throw std::runtime_error(name + " phase " + std::to_string(phase) +
                                 ": " + error.what());
      }
    }
  }

  void detach() {
    for (auto at = protocol_.attachOrder.rbegin();
         at != protocol_.attachOrder.rend(); ++at) {
      auto &entry = page(*at);
      if (!entry.attached)
        continue;
      device_.detachBuffer(entry.spec.index);
      entry.attached = false;
    }
    attached_ = false;
  }

  void release() {
    detach();
    for (auto &entry : pages_)
      if (entry.mapping)
        munmap(entry.mapping, entry.spec.size);
    pages_.clear();
  }

  directipu::Device &device_;
  Protocol protocol_;
  std::vector<Mapping> pages_;
  bool attached_ = false;
};

} // namespace hostexchange