#ifndef PLS_INTERNAL_PROFILING_DAG_NODE_H_
#define PLS_INTERNAL_PROFILING_DAG_NODE_H_

#include "pls/internal/base/system_details.h"

#include <memory>
#include <list>
#include <algorithm>
#include <iostream>

namespace pls::internal::profiling {

struct dag_node {
  dag_node(unsigned spawning_thread_id) : spawning_thread_id_{spawning_thread_id} {};

  unsigned spawning_thread_id_;
  unsigned long max_memory_{0};
  unsigned long total_runtime_{0};

  std::unique_ptr<dag_node> next_node_;
  std::list<dag_node> child_nodes_;

  base::system_details::pointer_t node_print_id() {
    return reinterpret_cast<base::system_details::pointer_t >(this);
  }

  void dag_compact() {
    while (next_node_ && next_node_->child_nodes_.empty()) {
      max_memory_ = std::max(max_memory_, next_node_->max_memory_);
      total_runtime_ += next_node_->total_runtime_;
      next_node_ = std::move(next_node_->next_node_);
    }
    if (next_node_) {
      next_node_->dag_compact();
    }

    for (auto &child : child_nodes_) {
      child.dag_compact();
    }
  }

  void dag_print(std::ostream &stream, unsigned rank) {
    stream << node_print_id() << " [label=" << spawning_thread_id_ << ", rank=" << rank << "];" << std::endl;
    for (auto &child : child_nodes_) {
      child.dag_print(stream, rank + 1);
      stream << node_print_id() << " -> " << child.node_print_id() << ";" << std::endl;
    }
    if (next_node_) {
      next_node_->dag_print(stream, rank);
      stream << node_print_id() << " -> " << next_node_->node_print_id() << ";" << std::endl;
    }
  }

  unsigned dag_max_memory() {
    unsigned max = max_memory_;
    if (next_node_) {
      max = std::max(max, next_node_->dag_max_memory());
    }
    for (auto &child : child_nodes_) {
      max = std::max(max, child.dag_max_memory());
    }
    return max;
  }

  unsigned long dag_total_user_time() {
    unsigned long total_user_time = total_runtime_;
    if (next_node_) {
      total_user_time += next_node_->dag_total_user_time();
    }
    for (auto &child : child_nodes_) {
      total_user_time += child.dag_total_user_time();
    }
    return total_user_time;
  }

  unsigned long dag_critical_path() {
    unsigned long critical_path = total_runtime_;
    if (next_node_) {
      critical_path += next_node_->dag_critical_path();
    }

    unsigned long slowest_child = 0;
    for (auto &child : child_nodes_) {
      slowest_child = std::max(slowest_child, child.dag_critical_path());
    }
    critical_path += slowest_child;

    return critical_path;
  }
};

}

#endif //PLS_INTERNAL_PROFILING_DAG_NODE_H_