#include "trace_lib.h"

#include <assert.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#include "buffer.h"

size_t trace_sizeof_signal_type(trace_signal_type_t type) {
  switch (type) {
  case TRACE_SIGNAL_TYPE_BOOL:
    return sizeof(int);
  case TRACE_SIGNAL_TYPE_INT:
    return sizeof(int);
  case TRACE_SIGNAL_TYPE_FLOAT:
    return sizeof(float);
  }
  /* UNREACHABLE */
  assert (false);
  return 0;
}

typedef struct trace_signal {
  char *name;
  trace_signal_type_t type;
  buffer_t *samples;
} trace_signal_t;

trace_signal_t *trace_signal_alloc(const char *name,
                               trace_signal_type_t type,
                               size_t initial_buffer_size) {
  trace_signal_t *res = malloc_checked(sizeof *res);
  res->name = strdup_checked(name);
  res->type = type;
  res->samples =
    buffer_alloc(initial_buffer_size * trace_sizeof_signal_type(type));
  return res;
}

void trace_signal_free(trace_signal_t *signal) {
  assert (signal);

  buffer_free(signal->samples);
  free(signal->name);
  free(signal);
}

void trace_add_samples(trace_signal_t *signal, void *samples, size_t count) {
  assert (signal);
  assert (samples);
  buffer_write(signal->samples,
               samples,
               count * trace_sizeof_signal_type(signal->type));
}

const char *trace_time_unit_repr(trace_time_unit_t u) {
  const char *table[] = { "s", "ms", "us", "ns", "ps", "fs" };
  assert (TRACE_TIME_UNIT_S <= u && u <= TRACE_TIME_UNIT_FS);
  return table[u];
}

typedef struct trace_file {
  buffer_t *signals;
  trace_time_unit_t time_unit;
  size_t time_unit_factor;
} trace_file_t;

trace_file_t *trace_file_alloc(trace_time_unit_t time_unit,
                               size_t time_unit_factor) {
  trace_file_t *trace = malloc_checked(sizeof *trace);
  trace->time_unit = time_unit;
  trace->time_unit_factor = time_unit_factor;
  trace->signals = buffer_alloc(10 * sizeof(trace_signal_t));

  return trace;
}

void trace_file_free(trace_file_t *trace) {
  assert (trace);

  /* Free buffers. */
  buffer_foreach (trace_signal_t *, psig, trace->signals)
    trace_signal_free(*psig);
  buffer_free(trace->signals);

  free(trace);
}

trace_signal_t *trace_file_lookup_signal(const trace_file_t *trace,
                                         const char *signal_name) {
  assert (trace);
  assert (signal_name);

  trace_signal_t *res = NULL;

  buffer_foreach (trace_signal_t *, psig, trace->signals) {
    if (strcmp((*psig)->name, signal_name) == 0) {
      res = *psig;
      break;
    }
  }

  return res;
}

bool trace_file_add_signal(const trace_file_t *trace, trace_signal_t *signal) {
  assert (trace);
  assert (signal);

  if (trace_file_lookup_signal(trace, signal->name))
    return false;

  buffer_write(trace->signals, &signal, sizeof signal);
  return true;
}

typedef void (trace_backend_write_header_f)(FILE *, trace_file_t *);
typedef void (trace_backend_write_cycle_beg_f)(FILE *, size_t);
typedef void (trace_backend_write_cycle_end_f)(FILE *, size_t);
typedef void (trace_backend_write_sample_f)(FILE *, trace_signal_t *, void *);
typedef void (trace_backend_write_sample_missing_f)(FILE *, trace_signal_t *);

typedef struct trace_backend {
  char *file_extension;
  trace_backend_write_header_f *write_header;
  trace_backend_write_cycle_beg_f *write_cycle_beg;
  trace_backend_write_cycle_end_f *write_cycle_end;
  trace_backend_write_sample_f *write_sample;
  trace_backend_write_sample_missing_f *write_sample_missing;
} trace_backend_t;

void trace_file_write_vcd_header(FILE *f, trace_file_t *trace) {
  time_t current_time;
  time(&current_time);

  fprintf(f, "$version Generated by trace.c $end\n");
  fprintf(f, "$date %s $end\n", ctime(&current_time));
  fprintf(f, "$timescale %zu %s $end\n",
          trace->time_unit_factor,
          trace_time_unit_repr(trace->time_unit));

  /* Dump signal declarations. */
  fprintf(f, "$scope module Top $end\n");
  buffer_foreach (trace_signal_t *, psig, trace->signals) {
    fprintf(f, "$var ");
    switch ((*psig)->type) {
    case TRACE_SIGNAL_TYPE_BOOL:
      fprintf(f, "wire 1");
      break;
    case TRACE_SIGNAL_TYPE_INT:
      fprintf(f, "integer %zu", 8 * sizeof(int));
      break;
    case TRACE_SIGNAL_TYPE_FLOAT:
      fprintf(f, "real 32");
      break;
    }
    fprintf(f, " %p %s $end\n", (*psig), (*psig)->name);
  }
  fprintf(f, "$upscope $end\n");
  fprintf(f, "$enddefinitions\n");

  /* Dump samples. */
  fprintf(f, "$dumpvars\n");
}

void trace_file_write_vcd_cycle_beg(FILE *f, size_t cycle) {
  fprintf(f, "#%zu\n", cycle);
}

void trace_file_write_vcd_cycle_end(FILE *f, size_t cycle) {
}

void trace_file_write_vcd_sample(FILE *f, trace_signal_t *sig, void *sample) {
  switch (sig->type) {
  case TRACE_SIGNAL_TYPE_BOOL:
    fprintf(f, "%d%p\n", (*(int *)sample ? 1 : 0), sig);
    break;
  case TRACE_SIGNAL_TYPE_INT:
    fprintf(f, "r%d %p\n", *(int *)sample, sig);
    break;
  case TRACE_SIGNAL_TYPE_FLOAT:
    fprintf(f, "r%.16g %p\n", *(float *)sample, sig);
    break;
  }
}

void trace_file_write_vcd_sample_missing(FILE *f, trace_signal_t *sig) {
}

void trace_file_write_csv_header(FILE *f, trace_file_t *trace) {
  buffer_foreach (trace_signal_t *, psig, trace->signals) {
    fprintf(f, "%s,", (*psig)->name);
  }
  fprintf(f, "\n");
}

void trace_file_write_csv_cycle_beg(FILE *f, size_t cycle) {
}

void trace_file_write_csv_cycle_end(FILE *f, size_t cycle) {
  fprintf(f, "\n");
}

void trace_file_write_csv_sample(FILE *f, trace_signal_t *sig, void *sample) {
  switch (sig->type) {
  case TRACE_SIGNAL_TYPE_BOOL:
    fprintf(f, "%d,", (*(int *)sample ? 1 : 0));
    break;
  case TRACE_SIGNAL_TYPE_INT:
    fprintf(f, "%d,", *(int *)sample);
    break;
  case TRACE_SIGNAL_TYPE_FLOAT:
    fprintf(f, "%f,", *(float *)sample);
    break;
  }
}

void trace_file_write_csv_sample_missing(FILE *f, trace_signal_t *sig) {
  fprintf(f, "XXX,");
}

trace_backend_t backends[] = {
  {
    ".vcd",
    trace_file_write_vcd_header,
    trace_file_write_vcd_cycle_beg,
    trace_file_write_vcd_cycle_end,
    trace_file_write_vcd_sample,
    trace_file_write_vcd_sample_missing,
  },
  {
    ".csv",
    trace_file_write_csv_header,
    trace_file_write_csv_cycle_beg,
    trace_file_write_csv_cycle_end,
    trace_file_write_csv_sample,
    trace_file_write_csv_sample_missing,
  },
};

bool trace_file_write(trace_file_t *trace, const char *file_name) {
  assert (trace);
  assert (file_name);

  FILE *f = fopen(file_name, "w");

  if (!f)
    return false;

  trace_backend_t *backend = NULL;
  char *file_ext;

  /* Search for the backend. */

  if (!(file_ext = strrchr(file_name, '.'))) {
    fprintf(stderr, "[trace] could not determine file extension of %s\n",
            file_name);
    return false;
  }

  for (size_t i = 0; i < sizeof(backends) / sizeof(trace_backend_t); i++) {
    if (!strcmp(backends[i].file_extension, file_ext)) {
      backend = &backends[i];
      break;
    }
  }

  if (!backend) {
    fprintf(stderr, "[trace] unknown file extension \"%s\"\n", file_ext);
    return false;
  }

  /* Write header. */
  backend->write_header(f, trace);

  /* Write samples. */

  /* We maintain an array of pointers to the current samples in each signal, and
     walk through them as long as none of them is finished. */
  size_t signal_count = trace->signals->occupancy / sizeof(trace_signal_t *);
  unsigned char **psamples = calloc(signal_count, sizeof(unsigned char *));
  assert (psamples);
  for (size_t i = 0; i < signal_count; i++)
    psamples[i] = ((trace_signal_t **)trace->signals->data)[i]->samples->data;

  /* Loop until all signals have been depleted, consuming one sample from each
     non-depleted signal at each cycle. */
  for (size_t cycle = 0;; cycle++) {
    /* Check if there is at least one active sample. */
    bool active = false;
    for (size_t i = 0; i < signal_count; i++) {
      trace_signal_t *sig = ((trace_signal_t **)trace->signals->data)[i];
      if (psamples[i] < sig->samples->data + sig->samples->occupancy) {
        active = true;
        break;
      }
    }

    if (!active)
      break;

    /* Write cycle-start marker. */
    backend->write_cycle_beg(f, cycle);

    /* Write each sample, including missing ones. */
    for (size_t i = 0; i < signal_count; i++) {
      trace_signal_t *sig = ((trace_signal_t **)trace->signals->data)[i];

      if (psamples[i] < sig->samples->data + sig->samples->occupancy) {
        active = true;
        backend->write_sample(f, sig, psamples[i]);
        psamples[i] += trace_sizeof_signal_type(sig->type);
      } else
        backend->write_sample_missing(f, sig);
    }

    /* Write cycle-stop marker. */
    backend->write_cycle_end(f, cycle);
  }

  free(psamples);

  fclose(f);
  return true;
}
