#include "../util/typedef.h"
#include "../util/string.h"
#include "../util/vector.h"
#include "../util/file.h"

#include <assert.h>

struct MODEL_DATA {
  VEC3 position;
  VEC2 uv;
  VEC3 normal;
};

struct MODEL {
  LIST<MODEL_DATA> vertices;
  LIST<U32> indices;
};

static inline char* skip_space( char* s ) {
  while( *s == ' ' || *s == '\t' ) ++s;
  return s;
}

static inline VEC3 read_vec3( char* s ) {
  VEC3 temp{};
  s = skip_space( s );
  temp.x = strtof( s, &s );
  s = skip_space( s );
  temp.y = strtof( s, &s );
  s = skip_space( s );
  temp.z = strtof( s, &s );
  return temp;
}

namespace MTL {
  enum LINE_TYPES : U16 {
    comment = ( '#' | ( ' ' << 8 ) ),
    newmtl  = ( 'n' | ( 'e' << 8 ) ),
    Ns      = ( 'N' | ( 's' << 8 ) ),
    Ka      = ( 'K' | ( 'a' << 8 ) ),
    Kd      = ( 'K' | ( 'd' << 8 ) ),
    Ks      = ( 'K' | ( 's' << 8 ) ),
    Ke      = ( 'K' | ( 'e' << 8 ) ),
    Ni      = ( 'N' | ( 'i' << 8 ) ),
    d       = ( 'd' | ( ' ' << 8 ) ),
    illum   = ( 'i' | ( 'l' << 8 ) ),
    map_Kd  = ( 'm' | ( 'a' << 8 ) )
  };

  struct MTLDATA {
    STR name;
    F32 Ns;
    VEC3 Ka;
    VEC3 Kd;
    VEC3 Ks;
    VEC3 Ke;
    F32 Ni;
    F32 d;
    I32 illum;
    STR map_Kd;
  };

  struct MTLLIB {
    STR name;
    LIST<MTLDATA> data;
  };

  inline MTLLIB mtl_from_file( const char* path, const char* name ) {
    U64 mtl_size = 0;
    STR mtl_path( path );
    mtl_path += name;
    char* f_mtl = (char*)file_read( mtl_path.data, &mtl_size );
    if( !f_mtl ) {
      dlog( "mtl_from_file() : failed to read mtl file %s\n", mtl_path.data );
      return {};
    }

    MTLLIB mtl;
    mtl.name = name;
    for( U64 i = 0; i < mtl_size - 1 && f_mtl[i]; ++i ) {
      if( f_mtl[i] == '\n' )
        continue;

      I32 c0 = f_mtl[i];
      U16 c1 = f_mtl[i + 1];
      I32 first_two = c0 | c1 << 8;

      switch( first_two ) {
      case LINE_TYPES::newmtl: {
        char* start = f_mtl + i + 7;
        char* end = start;
        while( *end && *end != '\n' ) ++end;
        U32 len = end - start;
        STR temp;
        temp.reserve( len + 1 );
        memcpy( temp, start, len );
        temp.data[len] = 0;
        mtl.data.resize( mtl.data.size + 1 );
        mtl.data[mtl.data.size - 1].name = temp.data;
      } break;
      case LINE_TYPES::d:
      case LINE_TYPES::Ni:
      case LINE_TYPES::Ns: {
        char* s = f_mtl + i + 2;
        s = skip_space( s );
        F32 temp = strtof( s, &s );
        switch( first_two ) {
        case LINE_TYPES::d:
          mtl.data[mtl.data.size - 1].d  = temp;
          break;
        case LINE_TYPES::Ni:
          mtl.data[mtl.data.size - 1].Ni = temp;
          break;
        case LINE_TYPES::Ns:
          mtl.data[mtl.data.size - 1].Ns = temp;
          break;
        }
      } break;
      case LINE_TYPES::illum: {
        char* s = f_mtl + i + 6;
        mtl.data[mtl.data.size - 1].illum = (I32)strtol( s, &s, 10 );
      } break;
      case LINE_TYPES::Ka:
      case LINE_TYPES::Kd:
      case LINE_TYPES::Ke:
      case LINE_TYPES::Ks: {
        char* s = f_mtl + i + 3;
        switch( first_two ) {
        case LINE_TYPES::Ka:
          mtl.data[mtl.data.size - 1].Ka = read_vec3( s );
          break;
        case LINE_TYPES::Kd:
          mtl.data[mtl.data.size - 1].Kd = read_vec3( s );
          break;
        case LINE_TYPES::Ke:
          mtl.data[mtl.data.size - 1].Ke = read_vec3( s );
          break;
        case LINE_TYPES::Ks:
          mtl.data[mtl.data.size - 1].Ks = read_vec3( s );
          break;
        }
      } break;
      case LINE_TYPES::map_Kd: {
        char* start = f_mtl + i + 7;
        char* end = start;
        while( *end && *end != '\n' ) ++end;
        U32 len = end - start;
        STR temp;
        temp.reserve( len + 1 );
        memcpy( temp, start, len );
        temp.data[len] = 0;
        mtl.data[mtl.data.size - 1].map_Kd = temp.data;
      } break;
      case LINE_TYPES::comment: break;
      default: dlog( "mtl_from_file() : unhandled char combo %c%c\n", c0, c1 ); break;
      }

      while( f_mtl[i] && f_mtl[i] != '\n' ) ++i;
    }

    free( (void*)f_mtl );
    return mtl;
  }
};

namespace OBJ {
  enum LINE_TYPES : U16 {
    comment = ( '#' | ( ' ' << 8 ) ),
    f       = ( 'f' | ( ' ' << 8 ) ),
    mtllib  = ( 'm' | ( 't' << 8 ) ),
    o       = ( 'o' | ( ' ' << 8 ) ),
    s       = ( 's' | ( ' ' << 8 ) ),
    v       = ( 'v' | ( ' ' << 8 ) ),
    vn      = ( 'v' | ( 'n' << 8 ) ),
    vt      = ( 'v' | ( 't' << 8 ) ),
    usemtl  = ( 'u' | ( 's' << 8 ) )
  };

  enum FACE_TYPES : U8 {
    V     = 0,
    VVT   = 1,
    VVN   = 2,
    VVTVN = 3
  };

  constexpr U32 STR_LEN = 0x40;

  struct OBJFACE {
    OBJ::FACE_TYPES type;
    LIST<U32> indices;
  };

  struct OBJDATA {
    LIST<VEC2> uvs;
    LIST<VEC3> normals;
    LIST<VEC3> vertices;

    LIST<MTL::MTLLIB> mtllib;

    struct UseMtl {
      LIST<U32> face_start;
      LIST<STR> mtl_name;
    } usemtl;

    LIST<OBJFACE> faces;
  };

  static inline U32 face_stride( FACE_TYPES t ) {
    switch( t ) {
    case FACE_TYPES::V:     return 1;
    case FACE_TYPES::VVT:   return 2;
    case FACE_TYPES::VVN:   return 2;
    case FACE_TYPES::VVTVN: return 3;
    }
    return 0;
  }

  inline OBJDATA obj_from_file( const char* path, const char* name ) {
    char* f_obj = nullptr;
    STR obj_path( path );
    U64 obj_size = 0;
    obj_path += name;
    obj_path += ".obj";

    f_obj = (char*)file_read( obj_path.data, &obj_size );
    if( !f_obj ) {
      dlog( "obj_from_file() : failed to read obj file %s\n", obj_path.data );
      return {};
    }

    OBJDATA obj;
    for( U64 i = 0; i < obj_size - 1 && f_obj[i]; ++i ) {
      if( f_obj[i] == (I32)'#' ) {
        while( f_obj[i] && f_obj[i] != '\n' ) ++i;
        continue;
      }

      I32 c0 = f_obj[i];
      U16 c1 = f_obj[i + 1];
      I32 first_two = c0 | c1 << 8;

      switch( first_two ) {
      case LINE_TYPES::f: {
        FACE_TYPES type = (FACE_TYPES)0;
        char* s = f_obj + i + 2;
        bool type_set = 0;
        OBJFACE temp{};

        s = skip_space( s );
        while( *s && *s != '\n' ) {
          bool has_vt = 0;
          bool has_vn = 0;

          U32 v = strtoul( s, &s, 10 );
          temp.indices.push( --v );

          if( *s == '/' ) {
            ++s;

            if( *s != '/' ) {
              U32 vt = strtoul( s, &s, 10 );
              temp.indices.push( --vt );
              has_vt = 1;
            }

            if( *s == '/' ) {
              ++s;
              U32 vn = strtoul( s, &s, 10 );
              temp.indices.push( --vn );
              has_vn = 1;
            }
          }

          if( !type_set ) {
            type_set = 1;
            if( has_vt && has_vn ) type = FACE_TYPES::VVTVN;
            else if( has_vt )      type = FACE_TYPES::VVT;
            else if( has_vn )      type = FACE_TYPES::VVN;
            else                   type = FACE_TYPES::V;
          }

          while( *s && *s != ' ' && *s != '\n' ) ++s;
          s = skip_space( s );
        }
        temp.type = type;
        U32 stride = face_stride( type );
        U32 count = temp.indices.size / stride;
        if ( count <= 3 )
          obj.faces.push( temp );
        else {
          for( U32 idx = 1; idx < count - 1; ++idx ) {
            OBJFACE tri{};
            tri.type = type;

            for( U32 ind = 0; ind < stride; ++ind )
              tri.indices.push( temp.indices.data[ind] );
            for( U32 ind = 0; ind < stride; ++ind )
              tri.indices.push( temp.indices.data[idx * stride + ind] );
            for( U32 ind = 0; ind < stride; ++ind )
              tri.indices.push( temp.indices.data[( idx + 1 ) * stride + ind] );
            obj.faces.push( tri );
          }
        }
      } break;
      case LINE_TYPES::v:
      case LINE_TYPES::vn: {
        bool is_v = first_two == LINE_TYPES::v;
        U32 start = i + ( is_v ? 2 : 3 );
        char* s = f_obj + start;
        if( is_v )
          obj.vertices.push( read_vec3( s ) );
        else
          obj.normals.push( read_vec3( s ) );
      } break;
      case LINE_TYPES::vt: {
        char* s = f_obj + i + 3;
        VEC2 temp{};
        s = skip_space( s );
        temp.x = strtof( s, &s );
        s = skip_space( s );
        temp.y = strtof( s, &s );
        obj.uvs.push( temp );
      } break;
      case LINE_TYPES::mtllib:
      case LINE_TYPES::usemtl: {
        bool is_mtllib = first_two == LINE_TYPES::mtllib;
        char* start = f_obj + i + 7;
        char* end = start;
        while( *end && *end != '\n' ) ++end;
        U32 len = end - start;
        STR temp;
        temp.reserve( len + 1 );
        memcpy( temp, start, len );
        temp.data[len] = 0;
        if( !is_mtllib ) {
          if( obj.usemtl.mtl_name.size )
            obj.usemtl.face_start.push( obj.faces.size );
          obj.usemtl.mtl_name.push( temp.data );
        } else
          obj.mtllib.push( MTL::mtl_from_file( path, temp.data ) );
      } break;
      case LINE_TYPES::comment:
      case LINE_TYPES::o:
      case LINE_TYPES::s: break;
      default: dlog( "obj_from_file() : unhandled char combo %c%c\n", c0, c1 ); break;
      }

      while( f_obj[i] && f_obj[i] != '\n' ) ++i;
    }
    obj.usemtl.face_start.push( obj.faces.size );
    if ( !obj.usemtl.mtl_name.size )
      obj.usemtl.mtl_name.push( "default" );
    free( (void*)f_obj );
    return obj;
  }

  inline MODEL obj_to_model( OBJDATA obj ) {
    struct KEY { U32 v, vt, vn; };
    struct ENTRY { KEY key; U32 idx; };

    MODEL model;
    LIST<ENTRY> table;
    table.reserve( obj.faces.size * 3 );
    for( U32 f = 0; f < obj.faces.size; ++f ) {
      OBJFACE& face = obj.faces.data[f];
      U32 stride = face_stride( face.type );
      U32 count = face.indices.size / stride;

      for( U32 v = 0; v < count; ++v ) {
        KEY key = {
          0xFFFFFFFF,
          0xFFFFFFFF,
          0xFFFFFFFF
        };
        U32 base = v * stride;
        key.v = face.indices.data[base + 0];
        if( face.type == FACE_TYPES::VVT || face.type == FACE_TYPES::VVTVN )
          key.vt = face.indices.data[base + 1];
        if( face.type == FACE_TYPES::VVN )
          key.vn = face.indices.data[base + 1];
        if( face.type == FACE_TYPES::VVTVN )
          key.vn = face.indices.data[base + 2];

        U32 idx = 0;
        bool found = 0;
        // this should be fine until ~1m triangles, right?
        for( U32 i = 0; i < table.size; ++i ) {
          ENTRY& e = table.data[i];
          if( e.key.vn != key.vn
           || e.key.vt != key.vt
           || e.key.v != key.v
          ) continue;

          idx = e.idx;
          found = 1;
          break;
        }

        if( !found ) {
          MODEL_DATA data{};
          data.position = obj.vertices.data[key.v];
          if( face.type == FACE_TYPES::VVTVN
           || face.type == FACE_TYPES::VVT )
            data.uv = obj.uvs.data[key.vt];
          if( face.type == FACE_TYPES::VVTVN
           || face.type == FACE_TYPES::VVN )
            data.normal = obj.normals.data[key.vn];
          idx = model.vertices.size;
          model.vertices.push( data );
          table.push({ key, idx });
        }
        model.indices.push( idx );
      }
    }
    return model;
  }
};

inline MODEL model_from_obj( const I8* path, const I8* name ) {
  OBJ::OBJDATA obj = OBJ::obj_from_file( path, name );
  return OBJ::obj_to_model( obj );
}