#ifndef TREEVIZ_H
#define TREEVIZ_H

// TreeViz context. Contains settings for the visualizer functions. This
// is wrapped in a struct so you don't have to always pass the same options
// to every funcion call. However, if you want to do that, you can use the
// `TV_vizl` function.
typedef struct TVCtx_T {
  // Output filename. Should end in `.pdf`
  char *filename;

  // Function that takes a reference to a node, and returns a string
  // representing the node's data. The representation should not include any
  // information about children, just the data stored in the node that the
  // function was called on.
  char *(*represent_data)(void *node);

  // A function that takes a node, and iteratively returns each child of the
  // node on each successive call, setting `done` to 1 when all children have
  // been processed. Be sure to use `cur_idx` to track the function's progress:
  // this function will be called in a recursive context, so using static memory
  // introduces some funny bugs.
  void *(*next_child)(void *root_node, int *done, int *cur_idx);
} TVCtx;

// Visualize the tree rooted at `root_node`, using the settings found in `ctx`.
void TV_vizv(void *root_node, TVCtx *ctx);

// Visualize the tree rooted at `root_node`, using the settings specifed. See
// the `TVCtx` struct.
void TV_vizl(void *root_node, char *filename,
             char *(*represent_data)(void *root_node),
             void *(*next_child)(void *root_node, int *done, int *cur_idx));

#endif

#ifdef TREEVIZ_IMPLEMENTATION
#ifndef TREEVIZ_IMPL_ALREADY
#define TREEVIZ_IMPL_ALREADY

#include <assert.h>
#include <errno.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/wait.h>
#include <unistd.h>

#define ERRCHECK(...)                                                          \
  if (errno != 0) {                                                            \
    fprintf(stderr, "error: %s (errno %d)\n", strerror(errno), errno);         \
    fprintf(stderr, __VA_ARGS__);                                              \
    fputc('\n', stderr);                                                       \
    assert(errno == 0);                                                        \
  }

#define NOT_NULL(x) assert(x != NULL)

void _TV_write_dot(void *root_node, int fd, TVCtx *ctx, size_t *id) {
  void *(*next_child)(void *root_node, int *done, int *cur_idx) =
      ctx->next_child;
  char *(*represent_data)(void *node) = ctx->represent_data;

  char *repr = represent_data(root_node);

  size_t cur_id = *id;
  ++(*id);

  errno = 0;
  if (root_node == NULL) {
    dprintf(fd, "%zu [style=filled,fillcolor=gray,label=\"%s\"];", cur_id,
            repr);
  } else {
    dprintf(fd, "%zu [label=\"%s\"];", cur_id, repr);
  }
  ERRCHECK("TreeViz: _TV_write_dot: could not write to pipe");

  free(repr);

  if (root_node == NULL) {
    return;
  }

  int done = 0;
  int iter = 0;
  while (!done) {
    void *child = next_child(root_node, &done, &iter);

    errno = 0;
    dprintf(fd, "%zu->%zu;", cur_id, *id);
    ERRCHECK("TreeViz: _TV_write_dot: could not write to pipe");

    _TV_write_dot(child, fd, ctx, id);
  }
}

void TV_vizv(void *root_node, TVCtx *ctx) {
  int pipe_stdin[2];

  errno = 0;
  pipe(pipe_stdin);
  ERRCHECK("TreeViz: TV_vizv: could not create pipe to dot")

  int child_stdin_write = pipe_stdin[1];
  int child_stdin_read = pipe_stdin[0];

  char o_argument[1024] = "-o";
  strncat(o_argument, ctx->filename, 1021);

  errno = 0;
  pid_t pid = fork();
  ERRCHECK("TreeViz: TV_vizv: could not fork process");

  if (pid == 0) {
    //// CHILD PROCESS
    close(child_stdin_write);
    ERRCHECK("TreeViz: TV_vizv: CHILD: could not close old pipe, write end");

    dup2(child_stdin_read, STDIN_FILENO);
    ERRCHECK("TreeViz: TV_vizv: CHILD: could not set stdin fileno");

    close(child_stdin_read);
    ERRCHECK("TreeViz: TV_vizv: CHILD: could not close old pipe, read end");

    char *run = "dot";
    execlp(run, run, "-Tpdf", o_argument);
    ERRCHECK("TreeViz: TV_vizv: CHILD: could not exec `%s`.\n\nDo you have "
             "graphviz installed and in your PATH?\n",
             run);
  } else {
    //// PARENT PROCESS
    close(child_stdin_read);
    ERRCHECK("TreeViz: TV_vizv: could not close old pipe, read end");

    dprintf(child_stdin_write, "digraph{node [shape=rect];");
    ERRCHECK("TreeViz: TV_vizv: could not write to pipe");

    size_t id = 0;
    _TV_write_dot(root_node, child_stdin_write, ctx, &id);

    errno = 0;
    dprintf(child_stdin_write, "}");
    ERRCHECK("TreeViz: TV_vizv: could not write to pipe");

    close(child_stdin_write);
    ERRCHECK("TreeViz: TV_vizv: could not close old pipe, write end");

    int status;
    wait(&status);
    ERRCHECK("TreeViz: TV_vizv: could not wait on child process");
    if (!WIFEXITED(status) || WEXITSTATUS(status) != 0) {
      fprintf(stderr, "TreeViz: TV_vizv: dot process failed. "
                      "terminating immediately.\n");
      assert(status == 0);
    }
  }
}

void TV_vizl(void *root_node, char *filename,
             char *(*represent_data)(void *root_node),
             void *(*next_child)(void *root_node, int *done, int *cur_idx)) {
  TVCtx ctx = {filename, represent_data, next_child};
  TV_vizv(root_node, &ctx);
}

#endif
#endif