#include "src/kernel/pipe.h"
#include "src/kernel/process.h"
#include "src/kernel/stream.h"
#include "src/lib/memory.h"

#define PIPE_BUFFER_SIZE 65536

typedef struct {
  char buffer[PIPE_BUFFER_SIZE];
  uint64_t begin;
  uint64_t end;
  uint8_t write_closed;
  uint8_t read_closed;
} pipe_t;

typedef struct {
  stream_t stream;
  pipe_t *pipe;
} pipe_read_stream_t;

typedef struct {
  stream_t stream;
  pipe_t *pipe;
} pipe_write_stream_t;

static uint64_t null_read(__attribute__((unused)) stream_t *self, __attribute__((unused)) uint64_t max,
                          __attribute__((unused)) char *to) {
  return 0;
}

static uint64_t null_write(__attribute__((unused)) stream_t *self, __attribute__((unused)) const char *from,
                           __attribute__((unused)) uint64_t bytes) {
  return 0;
}

static uint64_t pipe_read(stream_t *self, uint64_t max, char *to) {
  pipe_t *pipe = ((pipe_read_stream_t *)self)->pipe;

  while (pipe->begin == pipe->end) {
    if (pipe->write_closed) {
      return 0;
    }
    __asm__ volatile("sti");
    process_next();
    __asm__ volatile("cli");
  }

  uint64_t available = (pipe->end + PIPE_BUFFER_SIZE - pipe->begin) % PIPE_BUFFER_SIZE;
  uint64_t to_read = available < max ? available : max;

  if (pipe->begin + to_read <= PIPE_BUFFER_SIZE) {
    memory_copy(pipe->buffer + pipe->begin, to_read, to);
  } else {
    uint64_t chunk = PIPE_BUFFER_SIZE - pipe->begin;
    memory_copy(pipe->buffer + pipe->begin, chunk, to);
    memory_copy(pipe->buffer, to_read - chunk, to + chunk);
  }

  pipe->begin = (pipe->begin + to_read) % PIPE_BUFFER_SIZE;
  return to_read;
}

static uint64_t pipe_write(stream_t *self, const char *from, uint64_t bytes) {
  pipe_t *pipe = ((pipe_read_stream_t *)self)->pipe;

  uint64_t written = 0;
  while (written < bytes) {
    while (((pipe->end + 1) % PIPE_BUFFER_SIZE) == pipe->begin) {
      if (pipe->read_closed) {
        return written;
      }
      __asm__ volatile("sti");
      process_next();
      __asm__ volatile("cli");
    }

    uint64_t available = PIPE_BUFFER_SIZE - (pipe->end + PIPE_BUFFER_SIZE - pipe->begin) % PIPE_BUFFER_SIZE - 1;
    uint64_t to_write = available < bytes - written ? available : bytes - written;

    if (pipe->end + to_write <= PIPE_BUFFER_SIZE) {
      memory_copy(from, to_write, pipe->buffer + pipe->end);
    } else {
      uint64_t chunk = PIPE_BUFFER_SIZE - pipe->end;
      memory_copy(from, chunk, pipe->buffer + pipe->end);
      memory_copy(from + chunk, to_write - chunk, pipe->buffer);
    }

    pipe->end = (pipe->end + to_write) % PIPE_BUFFER_SIZE;
    written += to_write;
  }

  return written;
}

static void read_close(stream_t *self) {
  pipe_t *pipe = ((pipe_read_stream_t *)self)->pipe;

  pipe->read_closed = 1;

  if (pipe->read_closed && pipe->write_closed) {
    memory_free(pipe);
  }
  memory_free(self);
}

static void write_close(stream_t *self) {
  pipe_t *pipe = ((pipe_write_stream_t *)self)->pipe;

  pipe->write_closed = 1;

  if (pipe->read_closed && pipe->write_closed) {
    memory_free(pipe);
  }
  memory_free(self);
}

void pipe_init(stream_t **write, stream_t **read) {
  pipe_t *pipe = memory_allocate(sizeof(pipe_t));

  pipe_write_stream_t *ws = memory_allocate(sizeof(pipe_write_stream_t));
  ws->stream.read = null_read;
  ws->stream.write = pipe_write;
  ws->stream.close = write_close;
  ws->pipe = pipe;

  pipe_read_stream_t *rs = memory_allocate(sizeof(pipe_read_stream_t));
  rs->stream.read = pipe_read;
  rs->stream.write = null_write;
  rs->stream.close = read_close;
  rs->pipe = pipe;

  *write = (stream_t *)ws;
  *read = (stream_t *)rs;
}