#include "token.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#ifndef RELEASE
#define _CB_IMPLEMENTATION
#include <cbuild/cbuild.h>
static _CB_PROJECT *this = NULL;
#endif 
#include "ast.h"
#include "codegen_jit.h"
#include "lexer.h"
#include "parser.h"

static void indentf(int indent) {
  for (int i = 0; i < indent; i++)
    printf("  ");
}

void print_expr(_EX *e, int indent) {
  if (!e) {
    indentf(indent);
    printf("(null expr)\n");
    return;
  }

  indentf(indent);
  switch (e->kind) {
  case EX_NUMBER:
    printf("Number(%d)\n", e->value);
    break;
  case EX_VAR:
    printf("Var(%s)\n", e->name);
    break;
  case EX_BINOP:
    printf("BinOp(%s)\n", _TN[e->binop.op]);
    indentf(indent + 1);
    printf("Left:\n");
    print_expr(e->binop.l, indent + 2);
    indentf(indent + 1);
    printf("Right:\n");
    print_expr(e->binop.r, indent + 2);
    break;
  case EX_CALL:
    printf("Call(%s) argc=%d\n", e->call.func_name, e->call.argc);
    for (int i = 0; i < e->call.argc; i++) {
      indentf(indent + 1);
      printf("Arg[%d]:\n", i);
      print_expr(e->call.args[i], indent + 2);
    }
    break;
  case EX_ADDR:
    printf("AddressOf(&)\n");
    indentf(indent + 1);
    printf("Operand:\n");
    print_expr(e->addr.expr, indent + 2);
    break;
  case EX_DEREF:
    printf("Dereference(*)\n");
    indentf(indent + 1);
    printf("Operand:\n");
    print_expr(e->deref.expr, indent + 2);
    break;
  case EX_STRING:
    printf("String(\"%s\")\n", e->string);
    break;
  case EX_INDEX:
    printf("ArrayIndex([])\n");
    indentf(indent + 1);
    printf("Array:\n");
    print_expr(e->index.array, indent + 2);
    indentf(indent + 1);
    printf("Index:\n");
    print_expr(e->index.index, indent + 2);
    break;
  default:
    printf("UnknownExpr(kind=%d)\n", e->kind);
    break;
  }
}

void print_stmt(_STN *s, int indent) {
  if (!s)
    return;

  indentf(indent);
  switch (s->kind) {
  case STK_RETURN:
    printf("Return\n");
    if (s->return_expr) {
      indentf(indent + 1);
      printf("Value:\n");
      print_expr(s->return_expr, indent + 2);
    }
    break;
  case STK_VAR_DECL:
    printf("VarDecl(%s) type=[%s", s->var_decl.name,
           tybase_name(s->var_decl.type.base));
    for (int i = 0; i < s->var_decl.type.ptr_level; i++)
      printf("*");
    if (s->var_decl.type.array_size >= 0) {
      if (s->var_decl.type.array_size == 0) {
        printf("[]");
      } else {
        printf("[%d]", s->var_decl.type.array_size);
      }
    }
    printf("]\n");
    if (s->var_decl.init) {
      indentf(indent + 1);
      printf("Init:\n");
      print_expr(s->var_decl.init, indent + 2);
    }
    break;
  case STK_ASSIGN:
    printf("Assign\n");
    indentf(indent + 1);
    printf("LHS:\n");
    print_expr(s->assign.lhs, indent + 2);
    indentf(indent + 1);
    printf("RHS:\n");
    print_expr(s->assign.expr, indent + 2);
    break;
  case STK_EXPR:
    printf("ExprStmt\n");
    print_expr(s->expr, indent + 1);
    break;
  case STK_BLOCK:
    printf("Block\n");
    print_stmt(s->body, indent + 1);
    break;
  case STK_IF:
    printf("If\n");
    indentf(indent + 1);
    printf("Condition:\n");
    print_expr(s->ifs.cond, indent + 2);
    indentf(indent + 1);
    printf("Then:\n");
    print_stmt(s->ifs.thenb, indent + 2);
    if (s->ifs.elseb) {
      indentf(indent + 1);
      printf("Else:\n");
      print_stmt(s->ifs.elseb, indent + 2);
    }
    break;
  case STK_WHILE:
    printf("While\n");
    indentf(indent + 1);
    printf("Condition:\n");
    print_expr(s->whl.cond, indent + 2);
    indentf(indent + 1);
    printf("Body:\n");
    print_stmt(s->whl.body, indent + 2);
    break;
  case STK_FOR:
    printf("For\n");
    if (s->fr.init) {
      indentf(indent + 1);
      printf("Init:\n");
      print_stmt(s->fr.init, indent + 2);
    }
    if (s->fr.cond) {
      indentf(indent + 1);
      printf("Condition:\n");
      print_expr(s->fr.cond, indent + 2);
    }
    if (s->fr.step) {
      indentf(indent + 1);
      printf("Step:\n");
      print_stmt(s->fr.step, indent + 2);
    }
    indentf(indent + 1);
    printf("Body:\n");
    print_stmt(s->fr.body, indent + 2);
    break;
  default:
    printf("UnknownStmt(kind=%d)\n", s->kind);
    break;
  }

  if (s->n)
    print_stmt(s->n, indent);
}

void print_func(_FN *f, int indent) {
  if (!f)
    return;

  indentf(indent);
  printf("Function(%s) params=[", f->name);
  for (int i = 0; i < f->pac; i++) {
    printf("%s", tybase_name(f->param_types[i].base));
    for (int j = 0; j < f->param_types[i].ptr_level; j++)
      printf("*");
    if (f->param_types[i].array_size >= 0) {
      if (f->param_types[i].array_size == 0) {
        printf("[]");
      } else {
        printf("[%d]", f->param_types[i].array_size);
      }
    }
    printf(" %s", f->params[i]);
    if (i < f->pac - 1)
      printf(", ");
  }
  printf("]\n");

  indentf(indent + 1);
  printf("Body:\n");
  print_stmt(f->body, indent + 2);

  if (f->n)
    print_func(f->n, indent);
}

int main(int argc, char *argv[]) {
#ifndef RELEASE
  _CB_CREATE_PROJECT(this, .name = "ccdjit", .build_type = BUILD_EXEC,
                     .files = CB_STRLIST("src/main.c"), .is_rebuild = 1,
                     .output = "bin/ccdjit");
  _CB_PROJECT_BUILD(.projects = CB_PROJECT_LIST(this), .run = 1);
#endif 
  if (argc == 1) {
    fprintf(stderr, "Usage: %s [file.c] [-- [args]] (optional)\n", argv[0]);
    fprintf(stderr, "%s add2.c -- 1 2\n", argv[0]);
    return 1;
  }
  const char *filename = argv[1];
  int _in_arg_c = 0;
  char **_in_arg_v = NULL;
  if (argc > 1) {
    for (int i = 2; i < argc; i++) {
      if (strcmp(argv[i], "--") == 0) {
        _in_arg_c = i + 1;
        _in_arg_v = &argv[i + 1];
        break;
      }
    }
  }

  printf("Input File Name: %s\n\tInput Args:\n", filename);
  for (int i = 0; i < _in_arg_c; i++)
    printf("\t%s\n", _in_arg_v[i]);

  FILE *file = fopen(filename, "r");
  if (file == NULL) {
    fprintf(stderr, "Failed to open file %s\n", filename);
    return 1;
  }
  fseek(file, 0, SEEK_END);
  size_t file_size = ftell(file);
  fseek(file, 0, SEEK_SET);
  char *src = malloc(file_size + 1);
  if (src == NULL) {
    fprintf(stderr, "Failed to allocate memory for file contents\n");
    fclose(file);
    return 1;
  }
  fread(src, file_size, 1, file);
  fclose(file);
  src[file_size] = '\0';

  _LX lx = {src, 0, 1, 0}; // line=1, col=0
  _FN *ast = parse_program(&lx);
  _FN *clean_ast = fnlist_prepare(ast);

#ifndef RELEASE
  print_func(clean_ast, 0);
#endif /* ifndef RELEASE */
  JIT jit;
  jit_init(&jit);

  jit_compile_all(&jit, clean_ast);

  int result = jit_run(&jit, argc, argv);

#ifndef RELEASE
  printf("JIT returned: %d\n", result);
#endif /* ifndef RELEASE */
  jit_free(&jit);
  fn_free(clean_ast);
  free(src);

#ifndef RELEASE
  return 0;
#else
  return result;
#endif
}