#include "forth_internal.h"

static enum fh_error w_recurse(struct fh_thread_s *fh, const struct fh_word_s *w)
{
  (void) w;
  enum fh_error rv;
  ENSURE_STATE(FH_STATE_COMPILE);

  TRY(fh_put_instr(fh, FH_INSTR_WORD, fh->dict_last));
  return FH_OK;
}


static enum fh_error w_if(struct fh_thread_s *fh, const struct fh_word_s *w)
{
  (void) w;
  enum fh_error rv;
  ENSURE_STATE(FH_STATE_COMPILE);

  TRY(cs_push(fh, fh->here));
  TRY(fh_put_instr(fh, FH_INSTR_JUMPZERO, MAGICADDR_UNRESOLVED));
  return FH_OK;
}

static enum fh_error w_else(struct fh_thread_s *fh, const struct fh_word_s *w)
{
  (void) w;
  enum fh_error rv;
  ENSURE_STATE(FH_STATE_COMPILE);

  uint32_t ifaddr = 0;
  TRY(cs_pop(fh, &ifaddr));
  struct fh_instruction_s *if_instr = fh_instr_at(fh, ifaddr);
  if (!if_instr || if_instr->data != MAGICADDR_UNRESOLVED) {
    LOGE("IF-ELSE control stack corruption");
    return FH_ERR_INTERNAL;
  }

  if_instr->data = fh->here + INSTR_SIZE;

  TRY(cs_push(fh, fh->here));
  TRY(fh_put_instr(fh, FH_INSTR_JUMP, MAGICADDR_UNRESOLVED));
  return FH_OK;
}

static enum fh_error w_then(struct fh_thread_s *fh, const struct fh_word_s *w)
{
  (void) w;
  enum fh_error rv;
  ENSURE_STATE(FH_STATE_COMPILE);

  uint32_t ifaddr = 0;
  TRY(cs_pop(fh, &ifaddr));
  struct fh_instruction_s *if_instr = fh_instr_at(fh, ifaddr);
  if (!if_instr || if_instr->data != MAGICADDR_UNRESOLVED) {
    LOGE("IF-ELSE control stack corruption");
    return FH_ERR_INTERNAL;
  }

  if_instr->data = fh->here;
  return FH_OK;
}

static enum fh_error w_until(struct fh_thread_s *fh, const struct fh_word_s *w)
{
  (void) w;
  enum fh_error rv;
  ENSURE_STATE(FH_STATE_COMPILE);

  uint32_t destaddr = 0;
  TRY(cs_pop(fh, &destaddr));
  TRY(fh_put_instr(fh, FH_INSTR_JUMPZERO, destaddr));
  return FH_OK;
}

static enum fh_error w_begin(struct fh_thread_s *fh, const struct fh_word_s *w)
{
  (void) w;
  enum fh_error rv;
  ENSURE_STATE(FH_STATE_COMPILE);
  TRY(cs_push(fh, fh->here)); /* dest */
  return FH_OK;
}

static enum fh_error wp_do(struct fh_thread_s *fh, const struct fh_word_s *w)
{
  (void) w;
  enum fh_error rv;
  ENSURE_STATE(FH_STATE_COMPILE);

  TRY(fh_put_instr(fh, w->param ? FH_INSTR_DO_QUESTION : FH_INSTR_DO, MAGICADDR_UNRESOLVED));
  // R will now contain: limit,index

  TRY(cs_push(fh, fh->here)); // mark the start position (after init)
  return FH_OK;
}

static enum fh_error wp_loop(struct fh_thread_s *fh, const struct fh_word_s *w)
{
  (void) w;
  enum fh_error rv;
  struct fh_instruction_s *ii;
  ENSURE_STATE(FH_STATE_COMPILE);

  uint32_t loopendaddr = fh->here;

  uint32_t startaddr;
  TRY(cs_pop(fh, &startaddr));

  TRY(fh_put_instr(fh, w->param ? FH_INSTR_LOOP_PLUS : FH_INSTR_LOOP, startaddr));

  // after the end LOOP
  uint32_t endaddr = fh->here;

  // resolve ?DO dest
  ii = fh_instr_at(fh, startaddr - INSTR_SIZE);
  if (!ii || ii->kind == FH_INSTR_DO_QUESTION && ii->data == MAGICADDR_UNRESOLVED) {
    ii->data = endaddr;
  }

  // Resolve LEAVEs
  while (startaddr < loopendaddr) {
    ii = fh_instr_at(fh, startaddr);
    if (!ii) {
      LOGE("WHAT?");
      return FH_ERR_INTERNAL;
    }
    if (ii->kind == FH_INSTR_LEAVE && ii->data == MAGICADDR_UNRESOLVED) {
      LOG("Resolve leave addr");
      ii->data = endaddr;
    }

    // forward, skipping strings safely
    if (ii->kind == FH_INSTR_TYPESTR || ii->kind == FH_INSTR_ALLOCSTR) {
      startaddr += INSTR_SIZE + ii->data;
      startaddr = WORDALIGNED(startaddr);
    } else {
      startaddr += INSTR_SIZE;
    }
  }

  return FH_OK;
}

static enum fh_error w_unloop(struct fh_thread_s *fh, const struct fh_word_s *w)
{
  (void) w;
  enum fh_error rv;
  uint32_t limit;
  TRY(rs_pop(fh, &limit));
  TRY(fh_loop_unnest(fh));
  return FH_OK;
}

static enum fh_error w_leave(struct fh_thread_s *fh, const struct fh_word_s *w)
{
  (void) w;
  enum fh_error rv;
  ENSURE_STATE(FH_STATE_COMPILE);

  TRY(fh_put_instr(fh, FH_INSTR_LEAVE, MAGICADDR_UNRESOLVED));
  return FH_OK;
}

static enum fh_error w_while(struct fh_thread_s *fh, const struct fh_word_s *w)
{
  (void) w;
  enum fh_error rv;

  ENSURE_STATE(FH_STATE_COMPILE);

  uint32_t destaddr = 0;
  TRY(cs_pop(fh, &destaddr));

  TRY(cs_push(fh, fh->here)); // orig
  TRY(cs_push(fh, destaddr)); // dest

  TRY(fh_put_instr(fh, FH_INSTR_JUMPZERO, MAGICADDR_UNRESOLVED));
  return FH_OK;
}

static enum fh_error w_repeat(struct fh_thread_s *fh, const struct fh_word_s *w)
{
  (void) w;
  enum fh_error rv;

  ENSURE_STATE(FH_STATE_COMPILE);

  uint32_t origaddr = 0;
  uint32_t destaddr = 0;
  TRY(cs_pop(fh, &destaddr));
  TRY(cs_pop(fh, &origaddr));

  struct fh_instruction_s *branch_instr = fh_instr_at(fh, origaddr);
  if (!branch_instr || branch_instr->data != MAGICADDR_UNRESOLVED) {
    LOGE("REPEAT control stack corruption");
    return FH_ERR_INTERNAL;
  }
  branch_instr->data = fh->here + INSTR_SIZE;

  TRY(fh_put_instr(fh, FH_INSTR_JUMP, destaddr));

  return FH_OK;
}

static enum fh_error w_again(struct fh_thread_s *fh, const struct fh_word_s *w)
{
  (void) w;
  enum fh_error rv;

  ENSURE_STATE(FH_STATE_COMPILE);

  uint32_t destaddr = 0;
  TRY(cs_pop(fh, &destaddr));
  TRY(fh_put_instr(fh, FH_INSTR_JUMP, destaddr));

  return FH_OK;
}

static enum fh_error wp_ij(struct fh_thread_s *fh, const struct fh_word_s *w)
{
  enum fh_error rv;
  TRY(ds_push(fh, w->param ? fh->loop_j : fh->loop_i));
  return FH_OK;
}


static enum fh_error w_case(struct fh_thread_s *fh, const struct fh_word_s *w)
{
  (void) w;
  enum fh_error rv;
  TRY(cs_push(fh, fh->here)); // save marker for ENDCASE to resolve all the ENDOF's within
  ENSURE_STATE(FH_STATE_COMPILE);
  return FH_OK;
}

static enum fh_error w_of(struct fh_thread_s *fh, const struct fh_word_s *w)
{
  (void) w;
  enum fh_error rv;
  ENSURE_STATE(FH_STATE_COMPILE);
  TRY(cs_push(fh, fh->here)); // save the marker for ENDOF
  TRY(fh_put_instr(fh, FH_INSTR_OF, MAGICADDR_UNRESOLVED));
  return FH_OK;
}

static enum fh_error w_endof(struct fh_thread_s *fh, const struct fh_word_s *w)
{
  (void) w;
  enum fh_error rv;
  ENSURE_STATE(FH_STATE_COMPILE);
  uint32_t ofaddr;
  TRY(cs_pop(fh, &ofaddr));

  struct fh_instruction_s *of_instr = fh_instr_at(fh, ofaddr);
  if (!of_instr || of_instr->data != MAGICADDR_UNRESOLVED) {
    LOGE("CASE-OF control stack corruption");
    return FH_ERR_INTERNAL;
  }

  of_instr->data = fh->here + INSTR_SIZE; // next

  TRY(fh_put_instr(fh, FH_INSTR_JUMP, MAGICADDR_ENDCASE_UNRESOLVED)); // go to end of CASEs
  return FH_OK;
}

static enum fh_error w_endcase(struct fh_thread_s *fh, const struct fh_word_s *w)
{
  (void) w;
  enum fh_error rv;
  ENSURE_STATE(FH_STATE_COMPILE);
  uint32_t caseaddr;
  TRY(cs_pop(fh, &caseaddr));

  // Now walk the instructions and resolve every MAGICADDR_ENDCASE_UNRESOLVED

  uint32_t caseendaddr = fh->here;
  struct fh_instruction_s *ii;

  // Resolve ENDOF. TODO copied from LOOP impl, unify?
  while (caseaddr < caseendaddr) {
    ii = fh_instr_at(fh, caseaddr);
    if (!ii) {
      LOGE("WHAT?");
      return FH_ERR_INTERNAL;
    }
    if (ii->kind == FH_INSTR_JUMP && ii->data == MAGICADDR_ENDCASE_UNRESOLVED) {
      LOG("Resolve endof jump");
      ii->data = caseendaddr + INSTR_SIZE;
    }

    // forward, skipping strings safely
    if (ii->kind == FH_INSTR_TYPESTR || ii->kind == FH_INSTR_ALLOCSTR) {
      caseaddr += INSTR_SIZE + ii->data;
      caseaddr = WORDALIGNED(caseaddr);
    } else {
      caseaddr += INSTR_SIZE;
    }
  }

  TRY(fh_put_instr(fh, FH_INSTR_ENDCASE, 0));

  return FH_OK;
}

const struct name_and_handler fh_builtins_control[] = {
    {"i",       wp_ij,     0, 0},
    {"j",       wp_ij,     0, 1},
    {"if",      w_if,      1, 0},
    {"else",    w_else,    1, 0},
    {"then",    w_then,    1, 0},
    {"recurse", w_recurse, 1, 0},
    {"do",      wp_do,     1, 0},
    {"?do",     wp_do,     1, 1},
    {"loop",    wp_loop,   1, 0},
    {"leave",   w_leave,   1, 0},
    {"+loop",   wp_loop,   1, 1},
    {"begin",   w_begin,   1, 0},
    {"while",   w_while,   1, 0},
    {"repeat",  w_repeat,  1, 0},
    {"again",   w_again,   1, 0},
    {"until",   w_until,   1, 0},
    {"unloop",  w_unloop,  0, 0},
    {"case",    w_case,    1, 0},
    {"of",      w_of,      1, 0},
    {"endof",   w_endof,   1, 0},
    {"endcase", w_endcase, 1, 0},

    { /* end marker */ }
};