#include <ui.h>

#include <cassert.h>
#include <cstring.h>
#include <font.h>
#include <list.h>

#include <stdlib.h>

static void* uiAlloc(size_t count, size_t size) {
  void* mem = calloc(count, size);
  ASSERT(mem);
  return mem;
}

#define UI_NEW(TYPE) (TYPE*)uiAlloc(1, sizeof(TYPE))
#define UI_DEL(ppWidget)       \
  {                            \
    assert(ppWidget);          \
    void* widget_ = *ppWidget; \
    if (widget_) {             \
      free(widget_);           \
      *ppWidget = 0;           \
    }                          \
  }

DEF_LIST(Widget, uiWidget*)

/// Base widget type.
typedef struct uiWidget {
  uiWidgetType type;
  uiRect       rect;
  Widget_list  children;
} uiWidget;

/// Button.
typedef struct uiButton {
  uiWidget widget;
  string   text;
} uiButton;

/// Frame.
typedef struct uiFrame {
  uiWidget widget;
} uiFrame;

/// Label.
typedef struct uiLabel {
  uiWidget widget;
  string   text;
} uiLabel;

/// Table cell.
typedef struct uiCell {
  uiWidget* child;
} uiCell;

/// Table.
typedef struct uiTable {
  uiWidget widget;
  int      rows;
  int      cols;
  int*     widths; /// Width, in pixels, for each each column.
  uiCell*  header; /// If non-null, row of 'cols' header cells.
  uiCell*  cells;  /// Array of 'rows * cols' cells.
} uiTable;

typedef struct uiLibrary {
  FontAtlas* font;
} uiLibrary;

// -----------------------------------------------------------------------------
// Library.

uiLibrary g_ui = {0};

bool uiInit(void) {
  // TODO: Embed the font into the library instead.
  const char* font_path = "../ui/fontbaker/NK57.bin";
  if (!(g_ui.font = LoadFontAtlas(font_path))) {
    return false;
  }

  // TODO: Remove.
  const FontHeader* header     = &g_ui.font->header;
  const int         glyph_size = header->glyph_width * header->glyph_height;
  const int         atlas_size = header->num_glyphs * glyph_size;
  printf("Loaded font: %s\n", font_path);
  printf(
      "Glyph: %dx%d (%d bytes)\n", header->glyph_width, header->glyph_height,
      glyph_size);
  printf(
      "Atlas: %dx%d (%d bytes)\n", header->num_glyphs * header->glyph_width,
      header->glyph_height, atlas_size);

  return true;
}

void uiShutdown(void) {}

// -----------------------------------------------------------------------------
// Widget.

static uiButton* uiGetButtonPtr(uiWidgetPtr ptr) {
  assert(ptr.type == uiTypeButton);
  assert(ptr.button);
  return ptr.button;
}

static uiFrame* uiGetFramePtr(uiWidgetPtr ptr) {
  assert(ptr.type == uiTypeFrame);
  assert(ptr.frame);
  return ptr.frame;
}

static uiLabel* uiGetLabelPtr(uiWidgetPtr ptr) {
  assert(ptr.type == uiTypeLabel);
  assert(ptr.label);
  return ptr.label;
}

static uiTable* uiGetTablePtr(uiWidgetPtr ptr) {
  assert(ptr.type == uiTypeTable);
  assert(ptr.table);
  return ptr.table;
}

static void DestroyWidget(uiWidget** ppWidget) {
  assert(ppWidget);

  uiWidget* widget = *ppWidget;
  if (widget) {
    list_foreach_mut(widget->children, child, { DestroyWidget(&child); });
  }
  UI_DEL(ppWidget);
}

uiWidgetPtr uiMakeButtonPtr(uiButton* button) {
  assert(button);
  return (uiWidgetPtr){.type = uiTypeButton, .button = button};
}

uiWidgetPtr uiMakeFramePtr(uiFrame* frame) {
  assert(frame);
  return (uiWidgetPtr){.type = uiTypeFrame, .frame = frame};
}

uiWidgetPtr uiMakeLabelPtr(uiLabel* label) {
  assert(label);
  return (uiWidgetPtr){.type = uiTypeLabel, .label = label};
}

uiWidgetPtr uiMakeTablePtr(uiTable* table) {
  assert(table);
  return (uiWidgetPtr){.type = uiTypeTable, .table = table};
}

void uiWidgetSetParent(uiWidgetPtr child_, uiWidgetPtr parent_) {
  uiWidget* child  = child_.widget;
  uiWidget* parent = parent_.widget;

  assert(child);
  assert(parent);

  list_add(parent->children, child);
}

// -----------------------------------------------------------------------------
// Button.

uiButton* uiMakeButton(const char* text) {
  assert(text);

  uiButton* button = UI_NEW(uiButton);

  *button = (uiButton){
      .widget =
          (uiWidget){
                     .type = uiTypeButton,
                     .rect = {0},
                     },
      .text = string_new(text),
  };
  return button;
}

// -----------------------------------------------------------------------------
// Label.

uiLabel* uiMakeLabel(const char* text) {
  assert(text);

  uiLabel* label = UI_NEW(uiLabel);

  *label = (uiLabel){
      .widget =
          (uiWidget){
                     .type = uiTypeLabel,
                     },
      .text = string_new(text),
  };
  return label;
}

// -----------------------------------------------------------------------------
// Frame.

uiFrame* uiMakeFrame(void) {
  uiFrame* frame     = UI_NEW(uiFrame);
  frame->widget.type = uiTypeFrame;
  return frame;
}

void uiDestroyFrame(uiFrame** ppFrame) { DestroyWidget((uiWidget**)ppFrame); }

uiSize uiGetFrameSize(const uiFrame* frame) {
  assert(frame);
  return (uiSize){
      .width  = frame->widget.rect.width,
      .height = frame->widget.rect.height,
  };
}

// -----------------------------------------------------------------------------
// Table.

static const uiCell* GetCell(const uiTable* table, int row, int col) {
  assert(table);
  return table->cells + (row * table->cols) + col;
}

static uiCell* GetCellMut(uiTable* table, int row, int col) {
  assert(table);
  return (uiCell*)GetCell(table, row, col);
}

static uiCell* GetLastRow(uiTable* table) {
  assert(table);
  assert(table->rows > 0);
  return &table->cells[table->cols * (table->rows - 1)];
}

uiTable* uiMakeTable(int rows, int cols, const char** header) {
  uiTable* table = UI_NEW(uiTable);

  *table = (uiTable){
      .widget = (uiWidget){.type = uiTypeTable},
      .rows   = rows,
      .cols   = cols,
      .widths = (cols > 0) ? calloc(cols, sizeof(int)) : 0,
      .header = header ? calloc(cols, sizeof(uiCell)) : 0,
      .cells  = (rows * cols > 0) ? calloc(rows * cols, sizeof(uiCell)) : 0,
  };

  if (header) {
    for (int col = 0; col < cols; ++col) {
      table->header[col].child = (uiWidget*)uiMakeLabel(header[col]);
    }
  }

  return table;
}

void uiTableAddRow(uiTable* table, const char** row) {
  assert(table);

  table->rows++;

  uiCell* cells =
      realloc(table->cells, table->rows * table->cols * sizeof(uiCell));
  assert(cells);
  table->cells = cells;

  uiCell* cell = GetLastRow(table);
  for (int col = 0; col < table->cols; ++col, ++cell) {
    cell->child = (uiWidget*)uiMakeLabel(row[col]);
  }
}

void uiTableSet(uiTable* table, int row, int col, uiWidgetPtr child) {
  assert(table);
  assert(child.widget);

  GetCellMut(table, row, col)->child = child.widget;
}

const uiWidget* uiTableGet(const uiTable* table, int row, int col) {
  assert(table);
  return GetCell(table, row, col)->child;
}

uiWidget* uiTableGetMut(uiTable* table, int row, int col) {
  assert(table);
  return GetCellMut(table, row, col)->child;
}

// -----------------------------------------------------------------------------
// Layout and resizing.

static void ResizeTable(uiTable* table, int width, int height) {
  assert(table);

  if (table->cols == 0) {
    return;
  }

  // Surface width: W.
  // Columns: N
  //
  // First, find the minimum width of each column based on their contents.
  //
  // If the sum of column widths < N, then distribute the extra space first
  // among the smallest columns and building up towards the larger.
  //
  // If the sum of column widths > N, subtract from the largest column first and
  // move towards the smaller ones to distribute the space as evenly as
  // possible.

  // Find the minimum width for each column.
  int* widths = table->widths;
  // Header.
  for (int col = 0; col < table->cols; ++col) {
    const uiCell*  cell   = &table->header[col];
    const uiLabel* label  = (uiLabel*)cell->child;
    const int      length = (int)string_length(label->text);

    widths[col] = length;
  }
  // Table contents.
  for (int row = 0; row < table->rows; ++row) {
    for (int col = 0; col < table->cols; ++col) {
      const uiCell* cell = GetCell(table, row, col);
      if (cell->child) {
        const uiLabel* label  = (uiLabel*)cell->child;
        const int      length = (int)string_length(label->text);

        widths[col] = length > widths[col] ? length : widths[col];
      }
    }
  }
  // Multiply string lengths times glyph width to compute pixel size.
  for (int col = 0; col < table->cols; ++col) {
    widths[col] *= g_ui.font->header.glyph_width;
  }

  // Find the sum of widths.
  int used_width = 0;
  for (int col = 0; col < table->cols; ++col) {
    used_width += widths[col];
  }

  // Pad if available width is larger than sum of widths.
  if (used_width < width) {
    // Divide evenly among columns.
    // const int extra = width - used_width;
    //    const int pad   = extra / table->cols;
    //    const int mod   = extra % table->cols;
    //    for (int col = 0; col < table->cols; ++col) {
    //      table->widths[col] += pad + (col < mod ? 1 : 0);
    //    }

    int extra = width - used_width;
    while (extra > 0) {
      // Find smallest column.
      int smallest = 0;
      for (int col = 1; col < table->cols; ++col) {
        if (widths[col] < widths[smallest]) {
          smallest = col;
        }
      }
      // Pad it and subtract from the budget.
      widths[smallest] += 1;
      extra--;
    }
  }
  // Shrink if available width is smaller than the sum of widths.
  else if (used_width > width) {
    int deficit = used_width - width;
    while (deficit > 0) {
      // Find largest column.
      int largest = 0;
      for (int col = 1; col < table->cols; ++col) {
        if (widths[col] > widths[largest]) {
          largest = col;
        }
      }
      // Shrink it and subtract from the deficit.
      widths[largest] -= 1;
      deficit--;
    }
  }
}

static void ResizeWidget(uiWidget* widget, int width, int height) {
  assert(widget);

  widget->rect.width  = width;
  widget->rect.height = height;

  switch (widget->type) {
  case uiTypeButton:
    break;
  case uiTypeFrame:
    list_foreach_mut(
        widget->children, child, { ResizeWidget(child, width, height); });
    break;
  case uiTypeLabel:
    break;
  case uiTypeTable:
    ResizeTable((uiTable*)widget, width, height);
    break;
  case uiTypeMax:
    TRAP();
    break;
  }
}

void uiResizeFrame(uiFrame* frame, int width, int height) {
  assert(frame);
  ResizeWidget(&frame->widget, width, height);
}

// -----------------------------------------------------------------------------
// Rendering.

static const uiPixel uiBlack = {40, 40, 40, 255};
static const uiPixel uiWhite = {255, 255, 255, 255};
static const uiPixel uiPink  = {128, 0, 128, 255};

/// Render state.
///
/// Render functions are allowed to manipulate the state internally (e.g., the
/// subsurface), but must leave the state intact before returning, except, of
/// course, for the rendered pixels.
///
/// We store a subsurface separate from the surface so that we can always check
/// whether a given coordinate is within the bounds of the physical surface.
typedef struct RenderState {
  uiSurface surface;    /// Surface of pixels on which the UI is rendered.
  uiRect    subsurface; /// Subregion where the current UI widget is rendered.
  uiPoint   pen;        /// Current pen position relative to subsurface.
} RenderState;

static void RenderWidget(RenderState* state, const uiWidget* widget);

void PushSubsurface(
    RenderState* state, int width, int height, uiRect* original_subsurface,
    uiPoint* original_pen) {
  assert(state);
  assert(original_subsurface);
  assert(original_pen);

  *original_subsurface = state->subsurface;
  *original_pen        = state->pen;

  state->subsurface.x      = state->subsurface.x + state->pen.x;
  state->subsurface.width  = width;
  state->subsurface.height = height;
  state->pen.x             = 0;
}

void PopSubsurface(
    RenderState* state, const uiRect* original_subsurface,
    const uiPoint* original_pen) {
  assert(state);
  assert(original_subsurface);
  assert(original_pen);

  state->subsurface = *original_subsurface;
  state->pen        = *original_pen;
}

/// Checks whether pen + (w,h) is within the surface and subsurface.
static bool PenInSurface(const RenderState* state, int w, int h) {
  assert(state);

  // Surface.
  const bool in_surface =
      ((state->subsurface.x + state->pen.x + w) < state->surface.width) &&
      ((state->subsurface.y + state->pen.y + h) < state->surface.height);

  // Subsurface.
  const bool in_subsurface = ((state->pen.x + w) < state->subsurface.width) &&
                             ((state->pen.y + h) < state->subsurface.height);

  return in_surface && in_subsurface;
}

/// Get the pixel at (x,y).
static uiPixel* SurfaceXy(uiSurface* surface, int x, int y) {
  assert(surface);
  assert(x >= 0);
  assert(y >= 0);
  assert(x < surface->width);
  assert(y < surface->height);
  return surface->pixels + (surface->width * y) + x;
}

/// Get the pixel at pen + (x,y).
static uiPixel* PixelXy(RenderState* state, int x, int y) {
  assert(state);
  return SurfaceXy(
      &state->surface, state->subsurface.x + state->pen.x + x,
      state->subsurface.y + state->pen.y + y);
}

static void FillRect(const uiRect* rect, uiPixel colour, RenderState* state) {
  assert(rect);
  assert(state);
  assert(rect->width <= state->subsurface.width);
  assert(rect->height <= state->subsurface.height);

  for (int y = rect->y; y < rect->y + rect->height; ++y) {
    uiPixel* pixel = PixelXy(state, rect->x, y);
    for (int x = rect->x; x < rect->x + rect->width; ++x) {
      *pixel++ = colour;
    }
  }
}

/// Render a glyph.
/// The glyph is clamped to the surface's bounds.
static void RenderGlyph(
    const FontAtlas* atlas, unsigned char c, RenderState* state) {
  assert(atlas);
  assert(state);
  assert(atlas->header.glyph_width <= state->subsurface.width);
  assert(atlas->header.glyph_height <= state->subsurface.height);

  const int glyph_width  = atlas->header.glyph_width;
  const int glyph_height = atlas->header.glyph_height;

  const unsigned char* glyph = FontGetGlyph(atlas, c);

  for (int y = 0; (y < atlas->header.glyph_height) &&
                  PenInSurface(state, glyph_width - 1, glyph_height - 1);
       ++y) {
    for (int x = 0; (x < atlas->header.glyph_width) &&
                    PenInSurface(state, glyph_width - 1, glyph_height - 1);
         ++x, ++glyph) {
      uiPixel* pixel = PixelXy(state, x, y);
      if (*glyph > 0) {
        pixel->r = *glyph;
        pixel->g = *glyph;
        pixel->b = *glyph;
        pixel->a = 255;
      }
    }
  }
}

static void RenderText(const char* text, size_t length, RenderState* state) {
  assert(text);
  assert(state);

  const FontAtlas* atlas = g_ui.font;

  const int glyph_width  = atlas->header.glyph_width;
  const int glyph_height = atlas->header.glyph_height;

  // Save the x-pen so that we can restore it after rendering the text.
  const int x0 = state->pen.x;

  // Truncate the text rendering if it exceeds the subsurface's width or height.
  const char* c = text;
  for (size_t i = 0;
       (i < length) && PenInSurface(state, glyph_width - 1, glyph_height - 1);
       ++i, ++c, state->pen.x += glyph_width) {
    RenderGlyph(atlas, *c, state);
  }

  state->pen.x = x0;
}

static void RenderFrame(const uiFrame* frame, RenderState* state) {
  assert(frame);

  FillRect(&frame->widget.rect, uiBlack, state);
}

static void RenderLabel(const uiLabel* label, RenderState* state) {
  assert(label);
  assert(state);

  RenderText(string_data(label->text), string_length(label->text), state);
}

static void RenderTable(const uiTable* table, RenderState* state) {
  assert(table);
  assert(state);

  const int x0 = state->pen.x;
  const int y0 = state->pen.y;

  uiRect  original_subsurface = {0};
  uiPoint original_pen        = {0};

  // Render header.
  if (table->header) {
    for (int col = 0; col < table->cols; ++col) {
      // Crop the column contents to the column width so that one column does
      // not spill into the next.
      PushSubsurface(
          state, table->widths[col], state->subsurface.height,
          &original_subsurface, &original_pen);

      const uiCell* cell = &table->header[col];
      RenderWidget(state, cell->child);

      // Reset the original subsurface and pen for subsequent columns.
      PopSubsurface(state, &original_subsurface, &original_pen);

      // Next column.
      state->pen.x += table->widths[col];
    }
  }
  state->pen.x = x0;
  state->pen.y += g_ui.font->header.glyph_height;

  // Render rows.
  for (int row = 0; (row < table->rows) && PenInSurface(state, 0, 0); ++row) {
    for (int col = 0; (col < table->cols) && PenInSurface(state, 0, 0); ++col) {
      // Crop the column contents to the column width so that one column does
      // not spill into the next.
      PushSubsurface(
          state, table->widths[col], state->subsurface.height,
          &original_subsurface, &original_pen);

      state->subsurface.x     = state->subsurface.x + state->pen.x;
      state->subsurface.width = table->widths[col];
      state->pen.x            = 0;

      const uiCell* cell = GetCell(table, row, col);
      RenderWidget(state, cell->child);

      // Reset the original subsurface and pen for subsequent columns.
      PopSubsurface(state, &original_subsurface, &original_pen);

      // Next column.
      state->pen.x += table->widths[col];
    }
    state->pen.x = x0;
    state->pen.y += g_ui.font->header.glyph_height;
  }
  state->pen.y = y0;
}

static void RenderWidget(RenderState* state, const uiWidget* widget) {
  assert(state);
  assert(widget);

  // Render this widget.
  switch (widget->type) {
  case uiTypeButton:
    break;
  case uiTypeFrame:
    RenderFrame((const uiFrame*)widget, state);
    break;
  case uiTypeLabel:
    RenderLabel((const uiLabel*)widget, state);
    break;
  case uiTypeTable:
    RenderTable((const uiTable*)widget, state);
    break;
  case uiTypeMax:
    TRAP();
    break;
  }

  // Render children.
  list_foreach(widget->children, child, { RenderWidget(state, child); });
}

void uiRender(const uiFrame* frame, uiSurface* surface) {
  assert(frame);
  assert(surface);

  RenderWidget(
      &(RenderState){
          .surface = *surface,
          .subsurface =
              (uiRect){
                       .x      = 0,
                       .y      = 0,
                       .width  = surface->width,
                       .height = surface->height},
          .pen = {.x = 0, .y = 0},
  },
      (const uiWidget*)frame);
}