#ifndef INCLUDE_ast
#define INCLUDE_ast

#include "token.h"
#include <stdlib.h>
#include <string.h>
#include <stdio.h>


typedef enum { EX_NUMBER, EX_VAR, EX_BINOP, EX_CALL, EX_ADDR, EX_DEREF, EX_STRING, EX_INDEX } _EK;

typedef struct _EX {
  _EK kind;
  union {
    int value;  // for number
    char *name; // for var
    struct {
      struct _EX *l;
      _TK op;
      struct _EX *r;
    } binop;
    struct {
      char *func_name;
      struct _EX **args;
      int argc;
    } call;
    struct {
      struct _EX *expr; // &expr
    } addr;
    struct {
      struct _EX *expr; // *expr
    } deref;
    char *string; // for string literal
    struct {
      struct _EX *array; // array expression
      struct _EX *index; // index expression
    } index;
  };
} _EX;


typedef enum { STK_RETURN, STK_VAR_DECL, STK_ASSIGN, STK_EXPR, STK_BLOCK, STK_IF, STK_WHILE, STK_FOR } _STK;

typedef struct _STN {
  _STK kind;
  union {
    struct {
      char *name;
      _EX *init;
      _TY type;
    } var_decl;
    struct {
      _EX *lhs; // var or *expr
      _EX *expr;
    } assign;
    _EX *expr;         // STK_EXPR
    _EX *return_expr;  // STK_RETURN
    struct _STN *body; // STK_BLOCK
    struct {           // STK_IF
      _EX *cond;
      struct _STN *thenb;
      struct _STN *elseb; // may be NULL
    } ifs;
    struct {           // STK_WHILE
      _EX *cond;
      struct _STN *body;
    } whl;
    struct {           // STK_FOR
      struct _STN *init; // may be NULL (a stmt like decl/assign)
      _EX *cond;         // may be NULL
      struct _STN *step; // may be NULL (an expr stmt)
      struct _STN *body;
    } fr;
  };
  struct _STN *n; // linked list
} _STN;



typedef struct _FN {
  char *name;
  char **params;
  _TY *param_types;
  int pac;
  _STN *body;
  struct _FN *n;
} _FN;

typedef struct NameEntry {
    char *name;
    struct NameEntry *next;
} NameEntry;

#define HASH_SIZE 128

static unsigned long hash_str(const char *s) {
    unsigned long h = 5381;
    int c;
    while ((c = *s++))
        h = ((h << 5) + h) + (unsigned long)c;
    return h;
}

static int name_seen(NameEntry **table, const char *name) {
    unsigned long h = hash_str(name) % HASH_SIZE;
    for (NameEntry *e = table[h]; e; e = e->next) {
        if (strcmp(e->name, name) == 0)
            return 1;
    }
    NameEntry *ne = malloc(sizeof(*ne));
    if (!ne) {
        fprintf(stderr, "Error: malloc failed\n");
        exit(1);
    }
    ne->name = strdup(name);
    ne->next = table[h];
    table[h] = ne;
    return 0;
}

static void free_hash_table(NameEntry **table) {
    for (int i = 0; i < HASH_SIZE; i++) {
        NameEntry *e = table[i];
        while (e) {
            NameEntry *next = e->next;
            free(e->name);
            free(e);
            e = next;
        }
        table[i] = NULL;
    }
}

_FN *fnlist_prepare(_FN *head) {
    if (!head) return NULL;

    NameEntry *table[HASH_SIZE] = {0};

    _FN *main_node = NULL;
    _FN *main_prev = NULL;
    _FN *prev = NULL;

    for (_FN *cur = head; cur; cur = cur->n) {
        if (!cur->name) {
            fprintf(stderr, "[fnlist prepare] error: function with NULL name\n");
            free_hash_table(table);
            exit(1);
        }

        if (name_seen(table, cur->name)) {
            fprintf(stderr, "[fnlist prepare] error: duplicate function name '%s'\n", cur->name);
            free_hash_table(table);
            exit(1);
        }

        if (strcmp(cur->name, "main") == 0) {
            if (main_node) {
                fprintf(stderr, "[fnlist prepare] error: multiple 'main' functions found\n");
                free_hash_table(table);
                exit(1);
            }
            main_node = cur;
            main_prev = prev;
        }

        prev = cur;
    }

    free_hash_table(table);

    if (!main_node) return head;
    if (!main_prev) return head;

    main_prev->n = main_node->n;
    main_node->n = head;
    return main_node;
}





/* Generic alloc macros */
#define NEW_EX(k)                                                              \
  _EX *e = (_EX *)calloc(1, sizeof(_EX));                                      \
  e->kind = k
#define NEW_ST(k)                                                              \
  _STN *s = (_STN *)calloc(1, sizeof(_STN));                                   \
  s->kind = k
#define NEW_FN() _FN *f = (_FN *)calloc(1, sizeof(_FN))

/* Constructor declaration macros */
#define DEFINE_EX_CONSTRUCTOR(name, kind, ...)                                 \
  static inline _EX *ex_##name(__VA_ARGS__)

#define DEFINE_ST_CONSTRUCTOR(name, kind, ...)                                 \
  static inline _STN *st_##name(__VA_ARGS__)

#define DEFINE_FN_CONSTRUCTOR(name, ...)                                       \
  static inline _FN *fn_##name(__VA_ARGS__)



DEFINE_EX_CONSTRUCTOR(number, EX_NUMBER, int value) {
  NEW_EX(EX_NUMBER);
  e->value = value;
  return e;
}

DEFINE_EX_CONSTRUCTOR(var, EX_VAR, char *name) {
  NEW_EX(EX_VAR);
  e->name = name;
  return e;
}

DEFINE_EX_CONSTRUCTOR(binop, EX_BINOP, _EX *l, _TK op, _EX *r) {
  NEW_EX(EX_BINOP);
  e->binop.l = l;
  e->binop.r = r;
  e->binop.op = op;
  return e;
}

DEFINE_EX_CONSTRUCTOR(call, EX_CALL, char *func_name, _EX **args, int argc) {
  NEW_EX(EX_CALL);
  e->call.func_name = func_name;
  e->call.args = args;
  e->call.argc = argc;
  return e;
}

DEFINE_EX_CONSTRUCTOR(string, EX_STRING, char *str) {
  NEW_EX(EX_STRING);
  e->string = str;
  return e;
}

DEFINE_EX_CONSTRUCTOR(index, EX_INDEX, _EX *array, _EX *idx) {
  NEW_EX(EX_INDEX);
  e->index.array = array;
  e->index.index = idx;
  return e;
}


DEFINE_ST_CONSTRUCTOR(return, STK_RETURN, _EX *expr) {
  NEW_ST(STK_RETURN);
  s->return_expr = expr;
  return s;
}

DEFINE_ST_CONSTRUCTOR(var_decl, STK_VAR_DECL, char *name, _EX *init) {
  NEW_ST(STK_VAR_DECL);
  s->var_decl.name = name;
  s->var_decl.init = init;
  return s;
}

DEFINE_ST_CONSTRUCTOR(assign, STK_ASSIGN, _EX *lhs, _EX *expr) {
  NEW_ST(STK_ASSIGN);
  s->assign.lhs = lhs;
  s->assign.expr = expr;
  return s;
}

DEFINE_ST_CONSTRUCTOR(expr, STK_EXPR, _EX *expr) {
  NEW_ST(STK_EXPR);
  s->expr = expr;
  return s;
}

DEFINE_ST_CONSTRUCTOR(block, STK_BLOCK, _STN *body) {
  NEW_ST(STK_BLOCK);
  s->body = body;
  return s;
}

DEFINE_ST_CONSTRUCTOR(if, STK_IF, _EX *cond, _STN *thenb, _STN *elseb) {
  NEW_ST(STK_IF);
  s->ifs.cond = cond;
  s->ifs.thenb = thenb;
  s->ifs.elseb = elseb;
  return s;
}

DEFINE_ST_CONSTRUCTOR(while, STK_WHILE, _EX *cond, _STN *body) {
  NEW_ST(STK_WHILE);
  s->whl.cond = cond;
  s->whl.body = body;
  return s;
}

DEFINE_ST_CONSTRUCTOR(for, STK_FOR, _STN *init, _EX *cond, _STN *step, _STN *body) {
  NEW_ST(STK_FOR);
  s->fr.init = init;
  s->fr.cond = cond;
  s->fr.step = step;
  s->fr.body = body;
  return s;
}



DEFINE_FN_CONSTRUCTOR(new, char *name, char **params, _TY* params_types, int pac, _STN *body) {
  NEW_FN();
  f->name = name;
  f->params = params;
  f->param_types = params_types;
  f->pac = pac;
  f->body = body;
  return f;
}


static inline void ex_free(_EX *e) {
  if (!e)
    return;
  switch (e->kind) {
  case EX_NUMBER:
    break;
  case EX_VAR:
    free(e->name);
    break;
  case EX_BINOP:
    ex_free(e->binop.l);
    ex_free(e->binop.r);
    break;
  case EX_CALL:
    for (int i = 0; i < e->call.argc; i++)
      ex_free(e->call.args[i]);
    free(e->call.args);
    free(e->call.func_name);
    break;
  case EX_ADDR:
    ex_free(e->addr.expr);
    break;
  case EX_DEREF:
    ex_free(e->deref.expr);
    break;
  case EX_STRING:
    free(e->string);
    break;
  case EX_INDEX:
    ex_free(e->index.array);
    ex_free(e->index.index);
    break;
  }
  free(e);
}

static inline void st_free(_STN *s) {
  if (!s)
    return;
  switch (s->kind) {
  case STK_RETURN:
    ex_free(s->return_expr);
    break;
  case STK_VAR_DECL:
    ex_free(s->var_decl.init);
    free(s->var_decl.name);
    break;
  case STK_ASSIGN:
    ex_free(s->assign.expr);
    break;
  case STK_EXPR:
    ex_free(s->expr);
    break;
  case STK_BLOCK:
    st_free(s->body);
    break;
  case STK_IF:
    ex_free(s->ifs.cond);
    st_free(s->ifs.thenb);
    st_free(s->ifs.elseb);
    break;
  case STK_WHILE:
    ex_free(s->whl.cond);
    st_free(s->whl.body);
    break;
  case STK_FOR:
    st_free(s->fr.init);
    ex_free(s->fr.cond);
    st_free(s->fr.step);
    st_free(s->fr.body);
    break;
  }
  _STN *next = s->n;
  free(s);
  st_free(next);
}

static inline void fn_free(_FN *f) {
  while (f) {
    _FN *next = f->n;
    if (f->name)
      free(f->name);
    if (f->params) {
      for (int i = 0; i < f->pac; i++) {
        free(f->params[i]);
      }
      free(f->params);
    }
    free(f->param_types);
    st_free(f->body);
    free(f);
    f = next;
  }
}

#endif