#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 );
  for ( size_t i = 0; i < vm -> program . constant_count; ++i ) {
    Constant constant = vm -> program . constant_pool [ i ];
    if ( constant . kind == CONSTANT_CLASS )
      free ( constant . class . members );
  }
  free ( vm -> program . globals . members );
  free ( vm -> program . constant_pool );
}

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;
}

ObjectValue * create_object ( VM * vm, Class * class ) {
  ObjectValue * object = (ObjectValue*) make_object ( & vm -> heap, class -> member_count );
  Constant * pool =  vm -> program . constant_pool;
  for ( size_t i = 0; i < class -> member_count; ++i ) {
    Constant constant = pool [ class -> members [ i ] ];
    assert ( constant . kind == CONSTANT_STRING );
    object -> members [ i ] = (SimpleEntry) { .name = constant . string, .value = NULL };
  }
  return object;
}

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 );
  }
  free ( args );
  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 (255 & bytes [ 3 ]) << 24 | (255 & bytes [ 2 ]) << 16 | (255 & bytes [ 1 ]) << 8 | (255 & bytes [ 0 ]);
}

static inline Class read_class ( const u8 * data, size_t * pos ) {
  Class class;
  u16 elements = read_u16 ( data, pos );
  class . member_count = elements;
  class . members = (u16*) malloc ( elements * sizeof (u16) );
  for ( size_t i = 0; i < elements; ++i )
    class . members [ i ] = read_u16 ( data, pos );
  return class;
}

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 ) {
  program -> constant_count = read_u16 ( input . str, pos );
  u8 tag;
  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 ) };
        break;
      case 0x01:
        constant_pool [ i ] = (Constant) { .kind = CONSTANT_NULL };
        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 } };
        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 } };
        break;
      }
      case 0x04:
        constant_pool [ i ] = (Constant) { .kind = CONSTANT_BOOLEAN, .boolean = read_byte ( input . str, pos ) };
        break;
      case 0x05: {
        constant_pool [ i ] = (Constant) { .kind = CONSTANT_CLASS, .class = read_class ( input . str, pos ) };
      }
      default:
        break;
    }
  }
  program -> constant_pool = constant_pool;
}

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
  program -> globals = read_class ( input . str, &pos );
  // read EP
  program -> entry_point = read_u16  ( input . str, &pos );
}

void bc_function_call ( VM * vm, u8 arguments, bool is_function, Str * name ) {
  Value * tmp [256];
  arguments = is_function ? arguments : arguments - 1;
  for ( size_t i = arguments; i > 0; --i )
    tmp [ i - 1 ] = ostack_pop ( & vm -> stack );
  Value * callee = ostack_pop ( & vm -> stack );
  ConstantFunction function;
  if ( is_function ) { 
    assert ( callee -> kind == VALUE_FUNCTION );
    function = ((CFunctionValue*) callee ) -> function;
  } else {
    Value * field = NULL;
    while ( callee -> kind == VALUE_OBJECT ) {
      field = find_current_object_field ( callee, *name );
      if ( field )
        break;
      callee = ((ObjectValue*) callee ) -> extends;
    }
    if ( ! field ) {
      Value * result = try_operator ( & vm -> heap, callee, tmp, arguments, name );
      assert ( result );
      ostack_push ( & vm -> stack, result );
      return;
    }
    assert ( field -> kind == VALUE_FUNCTION );
    function = ((CFunctionValue*) field) -> 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 );
  if ( is_function )
    new_frame -> locals [ 0 ] = make_null ( & vm -> heap );
  else
    new_frame -> locals [ 0 ] = callee;
  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 0x03: {
        // printf ("array\n");
        Value * init = ostack_pop ( & vm -> stack );
        Value * size = ostack_pop ( & vm -> stack );
        assert ( size -> kind == VALUE_INTEGER );
        size_t len = ((IntValue*) size) -> integer;
        ArrayValue * array = (ArrayValue*) make_array ( & vm -> heap, len );
        for ( size_t i = 0; i < len; ++i )
          array -> elements [ i ] = init;
        ostack_push ( & vm -> stack, (Value*) array );
        break;
      }
      case 0x04: {
        // printf ("object\n");
        u16 index = read_u16 ( function . start, & vm -> instruction_pointer );
        Constant constant = vm -> program . constant_pool [ index ]; 
        assert ( constant . kind == CONSTANT_CLASS );
        assert ( constant . class . member_count < vm -> stack . size );
        ObjectValue * obj = create_object ( vm, & constant . class );
        for ( size_t i = obj -> member_cnt; i > 0; --i )
          obj -> members [ i - 1 ] . value = ostack_pop ( & vm -> stack );
        obj -> extends = ostack_pop ( & vm -> stack );
        ostack_push ( & vm -> stack, (Value*) obj );
        break;
      }
      case 0x05: {
        // printf ("get field\n");
        u16 index = read_u16 ( function . start, & vm -> instruction_pointer );
        assert ( index < vm -> program . constant_count );
        Constant name = vm -> program . constant_pool [ index ];
        assert ( name . kind == CONSTANT_STRING );
        Value * object = ostack_pop ( & vm -> stack );
        assert ( object -> kind == VALUE_OBJECT );
        Value ** field = get_object_field ( object, name . string );
        assert ( field );
        ostack_push ( & vm -> stack, *field );
        break;
      }
      case 0x06: {
        // printf ("set field\n");
        u16 index = read_u16 ( function . start, & vm -> instruction_pointer );
        assert ( index < vm -> program . constant_count );
        Constant name = vm -> program . constant_pool [ index ];
        assert ( name . kind == CONSTANT_STRING );
        Value * value = ostack_pop ( & vm -> stack );
        Value * object = ostack_pop ( & vm -> stack );
        assert ( object -> kind == VALUE_OBJECT );
        Value ** field = get_object_field ( object, name . string );
        assert ( field );
        *field = value;
        ostack_push ( & vm -> stack, value );
        break;
      }
      case 0x07: {
        // printf ("method call\n");
        u16 index = read_u16 ( function . start, & vm -> instruction_pointer );
        u8 arguments = read_byte ( function . start, & vm -> instruction_pointer );
        assert ( index < vm -> program . constant_count );
        Constant name = vm -> program . constant_pool [ index ];
        assert ( name . kind == CONSTANT_STRING );
        bc_function_call ( vm, arguments, false, & name . string );
        break;
      }
      case 0x08: {
        // printf ("function call\n");
        u8 arguments = read_byte ( function . start, & vm -> instruction_pointer );
        bc_function_call ( vm, arguments, true, NULL );
        break;
      }
      case 0x09: {
        // printf ("set local\n");
        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: {
        // printf ("get local\n");
        u16 index = read_u16 ( function . start, & vm -> instruction_pointer );
        assert ( index < local_count );
        ostack_push ( & vm -> stack, new_frame -> locals [ index ] );
        break;
      }
      case 0x0B: {
        // printf ("set global\n");
        u16 index = read_u16 ( function . start, & vm -> instruction_pointer );
        assert ( index < vm -> program . constant_count );
        Constant name = vm -> program . constant_pool [ index ];
        assert ( name . kind == CONSTANT_STRING );
        Value ** field = get_object_field ( vm -> globals, name . string );
        assert ( field );
        *field = ostack_top ( & vm -> stack );
        break;
      }
      case 0x0C: {
        // printf ("get global\n");
        u16 index = read_u16 ( function . start, & vm -> instruction_pointer );
        assert ( index < vm -> program . constant_count );
        Constant name = vm -> program . constant_pool [ index ];
        assert ( name . kind == CONSTANT_STRING );
        Value ** field = get_object_field ( vm -> globals, name . string );
        assert ( field );
        ostack_push ( & vm -> stack, *field );
        break;
      }
      case 0x0D: {
        // printf ("cond");
        i16 offset = (i16) read_u16 ( function . start, & vm -> instruction_pointer );
        if ( offset < 0 )
          assert ( (i32) vm -> instruction_pointer >= offset );
        else
          assert ( vm -> instruction_pointer + offset < function . len );
        Value * cond = ostack_pop ( & vm -> stack );
        if ( value_to_bool ( cond ) )
          vm -> instruction_pointer += offset;
        // printf (" %c\n", value_to_bool ( cond ) ? 'T' : 'F');
        break;
      }
      case 0x0E: {
        i16 offset = (i16) read_u16 ( function . start, & vm -> instruction_pointer );
        // printf ("jump %d from %lu\n", offset, vm -> instruction_pointer);
        assert ( (i32) vm -> instruction_pointer >= - offset && vm -> instruction_pointer + offset < function . len );
        vm -> instruction_pointer += offset;
        break;
      }
      case 0x0F: {
        // printf ("ret\n");
        Frame * tmp = vm -> local_frame;
        vm -> local_frame = tmp -> prev;
        vm -> instruction_pointer = tmp -> return_ip;
        free ( tmp );
        return;
      }
      default:
        assert ( false );
    }
  }
  free ( new_frame );
}

int vm_interpret ( VM * vm, Str input ) {
  parse_program ( & vm -> program, input );
  ostack_push ( & vm -> stack, alloc_constant ( vm, vm -> program . entry_point ) );
  ObjectValue * obj = create_object ( vm, & vm -> program . globals );
  for ( size_t i = 0; i < obj -> member_cnt; ++i )
    obj -> members [ i ] . value = make_null ( & vm -> heap );
  vm -> globals = (Value*) obj;
  bc_function_call ( vm, 0, true, NULL );
  
  return 0;
}