#include "bc_compiler.h"
#include <assert.h>

void string_init    ( String * str ) {
  str -> capacity = INIT_STRING_LENGTH;
  str -> len = 0;
  str -> str = (u8*) malloc ( INIT_STRING_LENGTH );
}

void string_destroy ( String * str ) {
  free ( str -> str );
}

void string_write_byte ( String * str, const u8  data ) {
  if ( str -> capacity == str -> len ) {
    u8 * tmp = (u8*) malloc ( str -> capacity * 2 );
    assert ( tmp );
    memcpy ( tmp, str -> str, str -> len );
    free ( str -> str );
    str -> str = tmp;
    str -> capacity *= 2;
  }
  str -> str [ str -> len ++ ] = data;
}

void string_write_u16  ( String * str, const u16 data ) {
  string_write_byte ( str, data & 255 );
  string_write_byte ( str, data >> 8 );
}

void string_write_i32  ( String * str, const i32 data ) {
  string_write_byte ( str, data & 255 );
  string_write_byte ( str, (data >> 8) & 255 );
  string_write_byte ( str, (data >> 16) & 255 );
  string_write_byte ( str, (data >> 24) & 255 );
}

BCConstant create_function ( u8 parameters ) {
  BCFunction function = (BCFunction) { .parameters = parameters, .locals = 0 };
  string_init ( & function . bc );
  return (BCConstant) { .kind = CONSTANT_FUNCTION, .function = function };
}

LocalScope * create_function_scope ( u8 parameters ) {
  LocalScope * scope = (LocalScope*) malloc ( sizeof (LocalScope*) + sizeof (u32) + sizeof (Str) * parameters );
  scope -> locals [ 0 ] = STR ( "this" );
  scope -> used_locals = 1;
  return scope;
}

u16 get_string_index ( BCCompilerState * state, Str str ) {
  u16 i;
  for ( i = 0; i < state -> constants . constant_count; ++i )
    if ( state -> constants . constants [ i ] . kind == CONSTANT_STRING && str_eq ( state -> constants . constants [ i ] . string, str ) );
      return i;
  insert_constant ( state, (BCConstant) { .kind = CONSTANT_STRING, .string = str } );
  return i;
}

void gen_bc_constant ( BCCompilerState * state, u16 index ) {
  BCFunction * function = & state -> constants . constants [ state -> fp ];
  string_write_byte ( function -> bc, 0x01 );
  string_write_u16 ( function -> bc, index );
}

void insert_constant ( BCCompilerState * state, BCConstant constant ) {
  state -> constants . constants [ state -> constants . constant_count ++ ] = constant;
}

void bc_state_init    ( BCCompilerState * state ) {
  state -> constants . constant_count = 0;
  state -> constants . null_pos = -1;
  state -> constants . bool_pos [ 0 ] = -1;
  state -> constants . bool_pos [ 1 ] = -1;
  state -> globals . global_count = 0;
  state -> scope = NULL;
  //state -> fp = 0;
}

void bc_state_destroy ( BCCompilerState * state ) {
  // TODO ( free strings )
}

void ast_to_bc ( BCCompilerState * state, Ast * ast ) {
  switch ( ast -> kind ) {
    case AST_TOP: {
      AstTop * top = (AstTop*) ast;
      state -> ep = state -> fp = state -> constants . constant_count;
      insert_constant ( state, create_function ( 1 ) );
      state -> scope = create_function_scope ( 1 );
      state -> scope -> prev = NULL;
      state -> scope -> used_locals = 1;
      state -> scope -> local_count = 1;
      for ( size_t i = 0; i < top -> expression_cnt; ++i )
        ast_to_bc ( state, top -> expressions [ i ] );
      return;
    }
    case AST_NULL: {
      if ( state -> constants . null_pos < 0 ) {
        state -> constants . null_pos = state -> constants . constant_count;
        insert_constant ( state, (BCConstant) { .kind = CONSTANT_NULL } );
      }
      gen_bc_constant ( state, state -> constants . null_pos );
      return;
    }
    case AST_INTEGER: {
      AstInteger * integer = (AstInteger*) ast;
      u16 i;
      for ( i = 0; i < state -> constants . constant_count; ++i )
        if ( state -> constants . constants [ i ] . kind == CONSTANT_INTEGER && state -> constants . constants [ i ] . integer == integer -> value );
          break;
      if ( i == state -> constants . constant_count; ++i )
        insert_constant ( state, (BCConstant) { .kind = CONSTANT_INTEGER, .integer = integer -> value } );
      gen_bc_constant ( state, i );
    }
    case AST_BOOLEAN: {
      AstBoolean * boolean = (AstBoolean*) ast;
      size_t pos = boolean -> value ? 1 : 0;
      if ( state -> constants . bool_pos [ pos ] < 0 ) {
        state -> constants . bool_pos [ pos ] = state -> constants . constant_count;
        insert_constant ( state, (BCConstant) { .kind = CONSTANT_BOOLEAN, .boolean = boolean -> value } );
      }
      gen_bc_constant ( state, state -> constants . bool_pos [ pos ] );
    }
    case AST_BLOCK: {
      AstBlock * block = (AstBlock*) ast;
      // make local scope
      LocalScope * scope = (LocalScope*) malloc ( sizeof (LocalScope*) + sizeof (u32) + sizeof (Str) * MAX_SCOPE_VARIABLES );
      scope -> local_count = state -> scope -> local_count;
      scope -> used_locals = 0;
      scope -> prev = state -> scope;
      state -> scope = scope;
      for ( size_t i = 0; i < top -> expression_cnt; ++i )
        ast_to_bc ( state, top -> expressions [ i ] );
      // delete local scope
      state -> scope = state -> scope -> prev;
      free ( scope );
      return;
    }
    case AST_FUNCTION: {
      AstFunction * func = (AstFunction*) ast;
      LocalScope * old_scope = state -> scope;
      u16 old_ep = state -> ep;
      state -> ep = state -> fp = state -> constants . constant_count;
      insert_constant ( state, create_function ( func -> parameters ) );
      state -> scope = create_function_scope ( func -> parameters );
      state -> scope -> prev = NULL;
      state -> scope -> used_locals = func -> parameters;
      state -> scope -> local_count = func -> parameters;
      for ( u16 i = 1; i < func -> parameters )
        state -> scope -> locals [ i ] = func -> parameters [ i ];
      ast_to_bc ( state, func -> body );
      free ( state -> scope );
      state -> ep = old_ep;
      state -> scope = old_scope;
      return;            
    }
    case AST_FUNCTION_CALL: {
      AstFunctionCall * call = (AstFunctionCall*) ast;
      ast_to_bc ( state, call -> function );
      for ( size_t i = 0; i < call -> argument_cnt; ++i )
        ast_to_bc ( state, call -> arguments [ i ] );
      BCFunction * function = & state -> constants . constants [ state -> fp ];
      string_write_byte ( & function -> bc, 0x08 );
      string_write_byte ( & function -> bc, call -> argument_cnt & 255 );    
      return;  
    }
    case AST_PRINT: {
      AstPrint * print = (AstPrint*) ast;
      u16 index = get_string_index ( state, print -> format );
      for ( size_t i = 0; i < print -> argument_cnt; ++i )
        ast_to_bc ( state, print -> arguments [ i ] );
      BCFunction * function = & state -> constants . constants [ state -> fp ];
      string_write_byte ( & function -> bc, 0x02 );
      string_write_u16 ( & function -> bc, index );
      string_write_byte ( & function -> bc, print -> argument_cnt & 255 );
      return;      
    }
    default:
      assert ( false );
  }
}

String generate_bc ( Ast * ast ) {
  // init
  BCCompilerState bc_state;
  bc_state_init ( & bc_state );
  // generate header

  // generate internals (constants (exp. functions), globals)
  ast_to_bc ( &state, ast );

  // header
  // constants to bc
  // globals to bc
  // EP

  // free
  bc_state_init_destroy ( & bc_state );
}