#include <stdio.h>
#include <inttypes.h>
#include <assert.h>
#include <string.h>

#include "bc_interpreter.h"

void ostack_init ( OperandStack * stack ) {
  stack -> elements = (Value**) malloc ( OPERAND_STACK_INIT_CAPACITY * sizeof (Value*) );
  stack -> size = 0;
  stack -> capacity = OPERAND_STACK_INIT_CAPACITY;
}

Value * ostack_pop ( OperandStack * stack  ) {
  assert ( stack -> size );
  stack -> size--;
  return stack -> elements [ stack -> size ];
}

Value * ostack_top ( OperandStack * stack  ) {
  return stack -> elements [ stack -> size - 1 ];
}

void ostack_push ( OperandStack * stack , Value * value ) {
  if ( stack -> size == stack -> capacity ) {
    Value ** tmp = (Value**) malloc ( 2 * stack -> capacity * sizeof (Value*) );
    memcpy ( tmp, stack -> elements, stack -> size * sizeof (Value*) );
    free ( stack -> elements );
    stack -> elements = tmp;
  }
  stack -> elements [ stack -> size++ ] = value;
}

void ostack_destroy ( OperandStack * stack ) {
  free ( stack -> elements );
}

void vm_init ( VM * vm, size_t heap_size ) {
  heap_init ( & vm -> heap, heap_size );
  make_null ( & vm -> heap );
  ostack_init ( & vm -> stack );
  vm -> local_frame = NULL;
  vm -> instruction_pointer = 0;
}

void vm_destroy ( VM * vm ) {
  heap_destroy ( & vm -> heap );
  ostack_destroy ( & vm -> stack );
}

// void * heap_alloc_alligned ( Heap * heap, size_t len, size_t align ) {
//   size_t pos = (size_t) heap -> next;
//   size_t rem = pos % align;
//   if ( rem )
//     heap -> next = heap -> next + align - rem;
//   if ( heap -> next + len >= heap -> end )
//     return NULL;  
//   void * ret = heap -> next;
//   heap -> next += len;
//   return ret;
// }
// 
// void * heap_alloc ( Heap * heap, size_t len ) {
//   size_t pos = (size_t) heap -> next;
//   size_t rem = pos % 8;
//   if ( rem )
//     heap -> next = heap -> next + 8 - rem;
//   if ( heap -> next + len >= heap -> end )
//     return NULL;  
//   void * ret = heap -> next;
//   heap -> next += len;
//   return ret;
// }
// 
// void heap_init ( Heap * heap, size_t heap_size ) {
//   heap -> begin = (u8*) malloc ( heap_size );
//   heap -> next = heap -> begin;
//   heap -> end = heap -> begin + heap_size;
// }
// 
// void heap_destroy ( Heap * heap ) {
//   free ( heap -> begin );
// }

//Value * make_null ( ASTInterpreterState * state ) {
//  NullValue * value = heap_alloc ( state -> heap, sizeof (NullValue) );
//  *value = (NullValue) { .kind = (Value) { .kind = VALUE_NULL } };
//  return (Value*) value;
//}
//
//Value * make_int ( ASTInterpreterState * state, i32 val ) {
//  IntValue * value = heap_alloc ( state -> heap, sizeof (IntValue) );
//  *value = (IntValue) { .kind = (Value) { .kind = VALUE_INTEGER }, .integer = val };
//  return (Value*) value;
//}
//
//Value * make_bool ( ASTInterpreterState * state, bool val ) {
//  BoolValue * value = heap_alloc ( state -> heap, sizeof (BoolValue) );
//  *value = (BoolValue) { .kind = (Value) { .kind = VALUE_BOOLEAN }, .boolean = val };
//  return (Value*) value;
//}

Value * alloc_constant ( VM * vm, u16 index ) {
  assert ( vm -> program . constant_count > index );
  Constant constant = vm -> program . constant_pool [ index ];
  switch ( constant . kind ) {
    case CONSTANT_INTEGER:
      return make_int ( & vm -> heap, constant . integer );
    case CONSTANT_NULL:
      return make_null ( & vm -> heap );
    case CONSTANT_BOOLEAN:
      return make_bool ( & vm -> heap, constant . boolean );
    case CONSTANT_FUNCTION:
      return make_cfunction ( & vm -> heap, constant . function );
    default:
      assert ( false );
  }
  return NULL;
}

void print_operand_value ( Value * value ) {
  switch ( value -> kind ) {
    case VALUE_INTEGER:
      printf ( "%" PRIi32, ((IntValue*) value ) -> integer );
      break;
    case VALUE_BOOLEAN:
      if ( ((BoolValue*) value ) -> boolean )
        printf ( "true" );
      else
        printf ( "false" );
      break;
    case VALUE_NULL:
      printf ( "null" );
      break;
    case VALUE_FUNCTION:
      printf ( "function" );
      break;
    /*case VALUE_ARRAY:
      putchar ( '[' );
      ArrayValue * arr = (ArrayValue*) value;
      for ( i32 i = 0; i < arr -> length; ++i ) {
        if ( i != 0 )
          printf ( ", " );
        print_operand_value ( arr -> elements [ i ] );
      }
      putchar ( ']' );
      break;
    case VALUE_OBJECT:
      printf ( "object(");
      ObjectValue * obj = (ObjectValue*) value;
      bool parent = false;
      if ( obj -> extends -> kind != VALUE_NULL ) {
        parent = true;
        printf ( "..=" );
        print_operand_value ( obj -> extends );
      }
      if ( obj -> member_cnt )
        qsort ( obj -> members, obj -> member_cnt, sizeof (SimpleEntry), compare_entry );
      for ( size_t i = 0; i < obj -> member_cnt; ++i ) {
        if ( i != 0 || parent )
          printf ( ", " );
        printf ( "%.*s=", (int) obj -> members [ i ] . name . len, obj -> members [ i ] . name . str );
        print_operand_value ( obj -> members [ i ] . value );
      }
      putchar ( ')' );
      break;*/
    default:
      assert ( false );
  }
}

void print_call ( VM * vm, u16 index, u8 arguments  ) {
  assert ( vm -> program . constant_count > index );
  Constant constant = vm -> program . constant_pool [ index ];
  assert ( constant . kind == CONSTANT_STRING );
  assert ( arguments <= vm -> stack . size );
  Value ** args = (Value**) malloc ( arguments * sizeof (Value*) );
  for ( size_t i = arguments; i != 0; --i )
    args [ i - 1 ] = ostack_pop ( & vm -> stack );
  size_t arg = 0;
  Str format = constant . string;
  u8 c;
  for ( size_t i = 0; i < format . len; ++i ) {
    c = format . str [ i ];
    if ( c == '\\' ) {
      ++i;
      c = format . str [ i ];
      switch ( c ) {
        case '~':
          putchar ( '~' );
          break;
        case 'n':
          putchar ( '\n' );
          break;
        case '"':
          putchar ( '"' );
          break;
        case 'r':
          putchar ( '\r' );
          break;
        case 't':
          putchar ( '\t' );
          break;
        case '\\':
          putchar ( '\\' );
          break;
        default:
          assert ( false );
      }
    } else if ( c == '~' ) {
      assert ( arg != arguments ); 
      print_operand_value ( args [ arg++ ] );
    } else
      putchar ( c );
  }
  assert ( arg == arguments );
}


static inline u8 read_byte ( const u8 * data, size_t * pos ) {
  return data [ (*pos)++ ];
}

static inline u16 read_u16 ( const u8 * data, size_t * pos ) {
  u8 bytes [ 2 ] = { read_byte ( data, pos ), read_byte ( data, pos ) };
  return (uint16_t) (255 & bytes [ 1 ]) << 8 | (255 & bytes [ 0 ]);
}

static inline i32 read_i32 ( const u8 * data, size_t * pos ) {
  u8 bytes [4];
  for ( size_t i = 0; i < 4; ++i )
    bytes [ i ] = read_byte ( data, pos );
  return bytes [ 3 ] << 24 | bytes [ 2 ] << 16 | bytes [ 1 ] << 8 | bytes [ 0 ];
}

void read_header ( Str input, size_t * pos ) {
  u8 bytes [4];
  for ( size_t i = 0; i < 4; ++i )
    bytes [ i ] = read_byte ( input . str, pos );
  Str header;
  header . str = bytes;
  header . len = 4;
  assert ( str_eq ( header, STR ( "FML\n" ) ) );
}

void read_constant_pool ( Program * program, Str input, size_t * pos ) {
  // printf ( "%u, %u\n", count_bytes [ 0 ], count_bytes [ 1 ] );
  program -> constant_count = read_u16 ( input . str, pos );
  // printf ( "%x, %x\n", count_bytes [ 0 ], count_bytes [ 1 ] );
  u8 tag;
  //printf ( "size is %u\n", program -> constant_count );
  Constant * constant_pool = (Constant*) malloc ( program -> constant_count * sizeof (Constant) );
  for ( size_t i = 0; i < program -> constant_count; ++i ) {
    tag = read_byte ( input . str, pos );
    switch ( tag ) {
      case 0x00:
        constant_pool [ i ] = (Constant) { .kind = CONSTANT_INTEGER, .integer = read_i32 ( input . str, pos ) };
        // printf ( "%d\n", constant_pool [ i ] . integer );
        break;
      case 0x01:
        constant_pool [ i ] = (Constant) { .kind = CONSTANT_NULL };
        // printf ( "null\n" );
        break;
      case 0x02: {
        u32 len = (u32) read_i32 ( input . str, pos );
        const u8 * start = input . str + *pos;
        *pos += len;
        constant_pool [ i ] = (Constant) { .kind = CONSTANT_STRING, .string = (Str) { .len = len, .str = start } };
        // printf ( "%.*s\n", (int) constant_pool [ i ] . string . len, constant_pool [ i ] . string . str );
        break;
      }
      case 0x03: {
        u8 parameters = read_byte ( input . str, pos );
        u16 locals = read_u16 ( input . str, pos );
        u32 len = (u32) read_i32 ( input . str, pos );
        const u8 * start = input . str + *pos;
        *pos += len;
        constant_pool [ i ] = (Constant) { .kind = CONSTANT_FUNCTION, .function = (ConstantFunction) { .parameters = parameters, .locals = locals, .len = len, .start = start } };
        // printf ( "func: p=%u, l=%u, len=%u, start=%lu\n", parameters, locals, len, *pos - len );
        break;
      }
      case 0x04:
        constant_pool [ i ] = (Constant) { .kind = CONSTANT_BOOLEAN, .boolean = read_byte ( input . str, pos ) };
        // printf ( "%s\n", constant_pool [ i ] . boolean == true ? "true" : "false" );
        break;
      default:
        // printf ( "unsupported type\n" );
        break;
    }
  }
  program -> constant_pool = constant_pool;
}

void read_globals ( Str input, size_t * pos ) {
  u16 elements = read_u16 ( input . str, pos );
  *pos += elements;
}

void parse_program ( Program * program, Str input ) {
  // read + check header
  size_t pos = 0;
  read_header ( input, &pos );
  // read constant pool 
  read_constant_pool ( program, input, &pos );
  read_globals ( input, &pos );
  program -> entry_point = read_u16  ( input . str, &pos );
  //printf ("EP: %u\n", program -> entry_point );
}

void bc_function_call ( VM * vm, u8 arguments ) {
  // assert ( index < vm -> program . constant_count );
  // assert ( vm -> program . constant_pool [ index ] . kind == CONSTANT_FUNCTION );
  Value * tmp [256];
  for ( size_t i = 0; i < arguments; ++i )
    tmp [ i ] = ostack_pop ( & vm -> stack );
  Value * fvalue = ostack_pop ( & vm -> stack );
  assert ( fvalue -> kind == VALUE_FUNCTION );
  ConstantFunction function = ((CFunctionValue*) fvalue ) -> function;
  size_t local_count = function . parameters + function . locals;
  Frame * new_frame = (Frame *) malloc ( sizeof (Frame) + local_count * sizeof (Value*) );
  new_frame -> prev = vm -> local_frame;
  new_frame -> return_ip = ++(vm -> instruction_pointer);
  vm -> local_frame = new_frame;
  assert ( arguments + 1 == function . parameters );
  new_frame -> locals [ 0 ] = make_null ( & vm -> heap );
  for ( size_t i = 1; i <= arguments; ++i )
    new_frame -> locals [ i ] = tmp [ i - 1 ];
  for ( size_t i = arguments + 1; i < local_count; ++i )
    new_frame -> locals [ i ] = make_null ( & vm -> heap );
  u8 opcode;
  for ( vm -> instruction_pointer = 0; vm -> instruction_pointer < function . len; ) {
    opcode = read_byte ( function . start, & vm -> instruction_pointer );
    switch ( opcode ) {
      case 0x00:
        //printf ("drop\n");
        ostack_pop ( & vm -> stack );
        break;
      case 0x01: {
        u16 index = read_u16 ( function . start, & vm -> instruction_pointer );
        //printf ("constant %u\n", index);
        ostack_push ( & vm -> stack, alloc_constant ( vm, index ) );
        break;
      }
      case 0x02: {
        //printf ("print\n");
        u16 index = read_u16 ( function . start, & vm -> instruction_pointer );
        u8 arguments = read_byte ( function . start, & vm -> instruction_pointer );
        print_call ( vm, index, arguments );
        ostack_push ( & vm -> stack, make_null ( & vm -> heap) );
        break;
      }
      case 0x08: {
        u8 arguments = read_byte ( function . start, & vm -> instruction_pointer );
        bc_function_call ( vm, arguments );
        break;
      }
      case 0x09: {
        u16 index = read_u16 ( function . start, & vm -> instruction_pointer );
        assert ( index < local_count );
        assert ( vm -> stack . size );
        new_frame -> locals [ index ] = ostack_top ( & vm -> stack );
        break;
      }
      case 0x0A: {
        u16 index = read_u16 ( function . start, & vm -> instruction_pointer );
        assert ( index < local_count );
        ostack_push ( & vm -> stack, new_frame -> locals [ index ] );
        break;
      }
      case 0x0D: {
        i16 offset = (i16) read_u16 ( function . start, & vm -> instruction_pointer );
        assert ( (i32) vm -> instruction_pointer >= offset && vm -> instruction_pointer + offset < function . len );
        Value * cond = ostack_top ( & vm -> stack );
        if ( value_to_bool ( cond ) )
          vm -> instruction_pointer += offset;
        break;
      }
      case 0x0E: {
        i16 offset = (i16) read_u16 ( function . start, & vm -> instruction_pointer );
        assert ( (i32) vm -> instruction_pointer >= offset && vm -> instruction_pointer + offset < function . len );
        vm -> instruction_pointer += offset;
        break;
      }
      case 0x0F: {
        Frame * tmp = vm -> local_frame;
        vm -> local_frame = tmp -> prev;
        vm -> instruction_pointer = tmp -> return_ip;
        free ( tmp );
        return;
      }
      default:
        //printf ( "%d\n", opcode );
        assert ( false );
    }
  }
}

int vm_interpret ( VM * vm, Str input ) {
  parse_program ( & vm -> program, input );
  ostack_push ( & vm -> stack, alloc_constant ( vm, vm -> program . entry_point ) );
  bc_function_call ( vm, 0 );
  //Constant start = vm -> program . constant_pool [ vm -> program . entry_point ];
  
  //u8 opcode;
  //size_t addr = f . start - input . str;
  //ostack_push ( & vm -> stack, (Value*) vm -> heap . begin );
  
  return 0;
}