/*
 * CS5600 University of Utah
 * Charles McGarvey
 * mcgarvey@eng.utah.edu
 */

#include <errno.h>

#include "array.h"
#include "tri.h"
#include "model.h"


// create an interface for a vector array
DEFINE_ARRAY_TYPE(vec);


struct model
{
    list_t*     triangles;
    mat_t       model;
    color_t     specular;
    scal_t      shininess;
    char*       name;
    int         count;
};


static int _model_read_raw(model_t* m, const char* filename);
static int _model_read_obj(model_t* m, const char* filename);

static int _model_try_read_cache(const char* filename, list_t** l);
static void _model_write_cache(const char* filename, int count, list_t* l);
    

model_t* model_alloc(const char* filename)
{
    int     type = 0;
    char*   ext = strrchr(filename, '.');
    if (ext == NULL) {
        goto fail;
    }
    ++ext;

    if (strcmp(ext, "raw") == 0) {
        return model_alloc2(filename, MODEL_TYPE_RAW);
    }
    if (strcmp(ext, "obj") == 0) {
        return model_alloc2(filename, MODEL_TYPE_OBJ);
    }

fail:
    fprintf(stderr, "Unknown file type: %s", filename);
    return NULL;
}

model_t* model_alloc2(const char* filename, int type)
{
    model_t* m = (model_t*)mem_alloc(sizeof(model_t));
    m->triangles = NULL;
    m->model = MAT_IDENTITY;
    m->specular = COLOR_WHITE;
    m->shininess = S(64.0);
    m->name = mem_strdup(filename);
    m->count = 0;

#if CACHE_GEOMETRY
    int count = _model_try_read_cache(filename, &m->triangles);
    if (0 < count) {
        return m;
    }
#endif

    int load;
    switch (type) {
    case MODEL_TYPE_RAW:
        load = _model_read_raw(m, filename);
        break;
    case MODEL_TYPE_OBJ:
        load = _model_read_obj(m, filename);
        break;
    }
    if (load != 0) {
        model_destroy(m);
        return NULL;
    }

#if CACHE_GEOMETRY
    _model_write_cache(filename, m->count, m->triangles);
#endif

    return m;
}

void model_destroy(model_t* m)
{
    list_destroy(&m->triangles);
    mem_free(m->name);
    mem_free(m);
}


const list_t* model_geometry(const model_t* m)
{
    return m->triangles;
}

int model_size(const model_t* m)
{
    return m->count;
}

const char* model_name(const model_t* m)
{
    return m->name;
}

color_t model_specular(const model_t* m)
{
    return m->specular;
}

scal_t model_shininess(const model_t* m)
{
    return m->shininess;
}

void model_transformation(const model_t* m, mat_t* transform)
{
    *transform = m->model;
}


void model_transform(model_t* m, const mat_t* transform)
{
    m->model = mat_mult(m->model, *transform);
}

void model_material(model_t* m, color_t specular, scal_t shininess)
{
    m->specular  = specular;
    m->shininess = shininess;
}


#if CALC_NORMALS

#include "map.h"
DEFINE_MAP_TYPE3(vec_t, list_t*, vnorm, vec_compare(*a, *b));


/*
 * Associate a triangle with one of its vertices.
 */
static void _find_normals_add_vertex(map_t* m, vec_t v, tri_t* t)
{
    list_t** l = map_vnorm_search(m, v);
    if (l == NULL) {
        map_vnorm_data_t* d = map_vnorm_insert(m, v, NULL);
        l = &d->val;
    }
    list_push(l, t);
}

/*
 * Associate a triangle with all of its vertices.
 */
static void _find_normals_add_triangle(map_t* m, tri_t* t)
{
    _find_normals_add_vertex(m, t->a.v, t);
    _find_normals_add_vertex(m, t->b.v, t);
    _find_normals_add_vertex(m, t->c.v, t);
}

/*
 * Calculate an averaged normal from a list of triangles that share a common
 * vertex.
 */
static void _find_normals_average(const vec_t* v, list_t** l)
{
    // first, compute the average normal
    vec_t n = VEC_ZERO;
    for (list_t* i = *l; i; i = i->link) {
        tri_t* t = (tri_t*)i->val;
        n = vec_add(n, tri_normal(*t));
    }
    n = vec_normalize(n);

    // set the normal on each triangle's vertex that is shared
    while (*l) {
        tri_t* t = (tri_t*)(*l)->val;
        if (vec_isequal(*v, t->a.v)) {
            t->a.n = n;
        }
        else if (vec_isequal(*v, t->b.v)) {
            t->b.n = n;
        }
        else if (vec_isequal(*v, t->c.v)) {
            t->c.n = n;
        }
        list_pop(l);
    }
}

#endif // CALC_NORMALS


static int _model_read_raw(model_t* m, const char* filename)
{
    FILE* file = fopen(filename, "r");
    if (file == NULL) {
        fprintf(stderr, "Cannot read %s: %s\n", filename, strerror(errno));
        return -1;
    }

#if CALC_NORMALS
    map_t* nlookup = map_vnorm_alloc();
#endif

    double x1, y1, z1, x2, y2, z2, x3, y3, z3;
    while (fscanf(file, " %lf %lf %lf %lf %lf %lf %lf %lf %lf",
                  &x1, &y1, &z1, &x2, &y2, &z2, &x3, &y3, &z3) == 9) {
        tri_t* t = tri_alloc(
            vert_new2((scal_t)x1, (scal_t)y1, (scal_t)z1),
            vert_new2((scal_t)x2, (scal_t)y2, (scal_t)z2),
            vert_new2((scal_t)x3, (scal_t)y3, (scal_t)z3)
        );
        list_push2(&m->triangles, t, mem_free);
        ++m->count;

#if CALC_NORMALS
        _find_normals_add_triangle(nlookup, t);
#else
        vec_t n = vec_normalize(tri_normal(*t));
        t->a.n = n;
        t->b.n = n;
        t->c.n = n;
#endif
    }

#if CALC_NORMALS
    map_vnorm_call(nlookup, _find_normals_average);
    rbtree_destroy(nlookup);
#endif

    fclose(file);
    if (m->triangles == NULL) {
        fprintf(stderr, "No triangles read from %s\n", filename);
        return -1;
    }
    return 0;
}

static int _model_read_obj(model_t* m, const char* filename)
{
    FILE* file = fopen(filename, "r");
    if (file == NULL) {
        fprintf(stderr, "Cannot read %s: %s\n", filename, strerror(errno));
        return -1;
    }

#if CALC_NORMALS
    map_t* nlookup = map_vnorm_alloc();
#endif

    array_t* v  = array_vec_alloc();
    array_t* vt = array_vec_alloc();
    array_t* vn = array_vec_alloc();

    array_vec_push(v, VEC_ZERO);
    array_vec_push(vt, VEC_ZERO);
    array_vec_push(vn, VEC_ZERO);

    char line[4096];
    while (fgets(line, 4096, file)) {
        char name[4096];
        double f1, f2, f3, f4;
        int meh;
        int i1, i2, i3, i4, i5, i6, i7, i8, i9;
        if (sscanf(line, "v %lf %lf %lf %lf ", &f1, &f2, &f3, &f4) == 4) {
            array_vec_push(v, vec_new2((scal_t)f1, (scal_t)f2, (scal_t)f3, (scal_t)f4));
        }
        else if (sscanf(line, "v %lf %lf %lf ", &f1, &f2, &f3) == 3) {
            array_vec_push(v, vec_new((scal_t)f1, (scal_t)f2, (scal_t)f3));
        }
        else if (sscanf(line, "vt %lf %lf ", &f1, &f2) == 2) {
            array_vec_push(vt, vec_new((scal_t)f1, (scal_t)f2, S(0.0)));
        }
        else if (sscanf(line, "vn %lf %lf %lf ", &f1, &f2, &f3) == 3) {
            array_vec_push(vn, vec_new((scal_t)f1, (scal_t)f2, (scal_t)f3));
        }
        else if (sscanf(line, "f %d %d %d ", &i1, &i2, &i3) == 3) {
            tri_t* t = tri_alloc(
                vert_new(*array_vec_index(v, i1)),
                vert_new(*array_vec_index(v, i2)),
                vert_new(*array_vec_index(v, i3))
            );
            list_push2(&m->triangles, t, mem_free);
            ++m->count;

#if CALC_NORMALS
            _find_normals_add_triangle(nlookup, t);
#else
            vec_t n = vec_normalize(tri_normal(*t));
            t->a.n = n;
            t->b.n = n;
            t->c.n = n;
#endif
        }
        else if (sscanf(line, "f %d/%d %d/%d %d/%d ", &i1, &i4, &i2, &i4, &i3, &i4) == 6) {
            tri_t* t = tri_alloc(
                vert_new(*array_vec_index(v, i1)),
                vert_new(*array_vec_index(v, i2)),
                vert_new(*array_vec_index(v, i3))
            );
            list_push2(&m->triangles, t, mem_free);
            ++m->count;

#if CALC_NORMALS
            _find_normals_add_triangle(nlookup, t);
#else
            vec_t n = vec_normalize(tri_normal(*t));
            t->a.n = n;
            t->b.n = n;
            t->c.n = n;
#endif
        }
        else if (sscanf(line, "f %d//%d %d//%d %d//%d ",
                        &i1, &i2, &i3, &i4, &i5, &i6) == 6) {
            tri_t* t = tri_alloc(
                vert_new(*array_vec_index(v, i1)),
                vert_new(*array_vec_index(v, i3)),
                vert_new(*array_vec_index(v, i5))
            );
            list_push2(&m->triangles, t, mem_free);
            ++m->count;
            t->a.n = *array_vec_index(vn, i2);
            t->b.n = *array_vec_index(vn, i4);
            t->c.n = *array_vec_index(vn, i6);
        }
        else if (sscanf(line, "f %d/%d/%d %d/%d/%d %d/%d/%d ",
                        &i1, &meh, &i2, &i3, &meh, &i4, &i5, &meh, &i6) == 9) {
            tri_t* t = tri_alloc(
                vert_new(*array_vec_index(v, i1)),
                vert_new(*array_vec_index(v, i3)),
                vert_new(*array_vec_index(v, i5))
            );
            list_push2(&m->triangles, t, mem_free);
            ++m->count;
            t->a.n = *array_vec_index(vn, i2);
            t->b.n = *array_vec_index(vn, i4);
            t->c.n = *array_vec_index(vn, i6);
        }
        // f 1/2 3/4 5/6
        // f 1/2/3 4/5/6 7/8/9
    }

    array_destroy(v);
    array_destroy(vt);
    array_destroy(vn);

#if CALC_NORMALS
    map_vnorm_call(nlookup, _find_normals_average);
    rbtree_destroy(nlookup);
#endif

    fclose(file);
    if (m->triangles == NULL) {
        fprintf(stderr, "No triangles read from %s\n", filename);
        return -1;
    }
    return 0;
}


#define _CHECK_IO(X) if ((X) <= 0) goto fail

/*
 * Try to read the triangle geometry from the cache file.
 */
static int _model_try_read_cache(const char* filename, list_t** l)
{
    int     count = 0;
    char*   cachename  = mem_strcat(".", filename);
    FILE*   file = fopen(cachename, "rb");
    if (file == NULL) {
        goto fail;
    }

    _CHECK_IO(fread(&count, sizeof(count), 1, file));

    float x1, y1, z1, x2, y2, z2, x3, y3, z3;
    for (int i = 0; i < count; ++i) {
        _CHECK_IO(fread(&x1, sizeof(float), 1, file));
        _CHECK_IO(fread(&y1, sizeof(float), 1, file));
        _CHECK_IO(fread(&z1, sizeof(float), 1, file));
        _CHECK_IO(fread(&x2, sizeof(float), 1, file));
        _CHECK_IO(fread(&y2, sizeof(float), 1, file));
        _CHECK_IO(fread(&z2, sizeof(float), 1, file));
        _CHECK_IO(fread(&x3, sizeof(float), 1, file));
        _CHECK_IO(fread(&y3, sizeof(float), 1, file));
        _CHECK_IO(fread(&z3, sizeof(float), 1, file));
        tri_t* t = tri_alloc(
            vert_new2((scal_t)x1, (scal_t)y1, (scal_t)z1),
            vert_new2((scal_t)x2, (scal_t)y2, (scal_t)z2),
            vert_new2((scal_t)x3, (scal_t)y3, (scal_t)z3)
        );
        list_push2(l, t, mem_free);
        _CHECK_IO(fread(&x1, sizeof(float), 1, file));
        _CHECK_IO(fread(&y1, sizeof(float), 1, file));
        _CHECK_IO(fread(&z1, sizeof(float), 1, file));
        _CHECK_IO(fread(&x2, sizeof(float), 1, file));
        _CHECK_IO(fread(&y2, sizeof(float), 1, file));
        _CHECK_IO(fread(&z2, sizeof(float), 1, file));
        _CHECK_IO(fread(&x3, sizeof(float), 1, file));
        _CHECK_IO(fread(&y3, sizeof(float), 1, file));
        _CHECK_IO(fread(&z3, sizeof(float), 1, file));
        t->a.n = vec_new((scal_t)x1, (scal_t)y1, (scal_t)z1);
        t->b.n = vec_new((scal_t)x2, (scal_t)y2, (scal_t)z2);
        t->c.n = vec_new((scal_t)x3, (scal_t)y3, (scal_t)z3);
    }

fail:
    mem_free(cachename);
    if (file != NULL) {
        fclose(file);
    }

    return count;
}

/*
 * Write the triangle data to the cache.
 */
static void _model_write_cache(const char* filename, int count, list_t* l)
{
    char*   cachename  = mem_strcat(".", filename);
    FILE*   file = fopen(cachename, "wb");
    if (file == NULL) {
        fprintf(stderr, "Cannot write %s: %s\n", cachename, strerror(errno));
        goto fail;
    }

    _CHECK_IO(fwrite(&count, sizeof(count), 1, file));
    for (list_t* i = l; i; i = i->link) {
        tri_t* t = (tri_t*)i->val;
        _CHECK_IO(fwrite(&t->a.v.x, sizeof(float), 1, file));
        _CHECK_IO(fwrite(&t->a.v.y, sizeof(float), 1, file));
        _CHECK_IO(fwrite(&t->a.v.z, sizeof(float), 1, file));
        _CHECK_IO(fwrite(&t->b.v.x, sizeof(float), 1, file));
        _CHECK_IO(fwrite(&t->b.v.y, sizeof(float), 1, file));
        _CHECK_IO(fwrite(&t->b.v.z, sizeof(float), 1, file));
        _CHECK_IO(fwrite(&t->c.v.x, sizeof(float), 1, file));
        _CHECK_IO(fwrite(&t->c.v.y, sizeof(float), 1, file));
        _CHECK_IO(fwrite(&t->c.v.z, sizeof(float), 1, file));
        _CHECK_IO(fwrite(&t->a.n.x, sizeof(float), 1, file));
        _CHECK_IO(fwrite(&t->a.n.y, sizeof(float), 1, file));
        _CHECK_IO(fwrite(&t->a.n.z, sizeof(float), 1, file));
        _CHECK_IO(fwrite(&t->b.n.x, sizeof(float), 1, file));
        _CHECK_IO(fwrite(&t->b.n.y, sizeof(float), 1, file));
        _CHECK_IO(fwrite(&t->b.n.z, sizeof(float), 1, file));
        _CHECK_IO(fwrite(&t->c.n.x, sizeof(float), 1, file));
        _CHECK_IO(fwrite(&t->c.n.y, sizeof(float), 1, file));
        _CHECK_IO(fwrite(&t->c.n.z, sizeof(float), 1, file));
    }

fail:
    mem_free(cachename);
    if (file != NULL) {
        fclose(file);
    }
}

#undef _CHECK_IO