/*
 * Piece Model File viewer using OpenGL and FreeGLUT.
 * Copyright (C) 2025 Kostas Michalopoulos.
 *
 * This software is provided 'as-is', without any express or implied
 * warranty.  In no event will the authors be held liable for any damages
 * arising from the use of this software.
 *
 * Permission is granted to anyone to use this software for any purpose,
 * including commercial applications, and to alter it and redistribute it
 * freely, subject to the following restrictions:
 *
 * 1. The origin of this software must not be misrepresented; you must not
 *    claim that you wrote the original software. If you use this software
 *    in a product, an acknowledgment in the product documentation would be
 *    appreciated but is not required.
 * 2. Altered source versions must be plainly marked as such, and must not be
 *    misrepresented as being the original software.
 * 3. This notice may not be removed or altered from any source distribution.
 *
 * Kostas Michalopoulos <badsector@runtimeterror.com>
 *
 ******************************************************************************
 * Compile with:
 *
 *  gcc -o pmfview pmfview.c -lGL -lGLU -lglut -lm
 *
 ******************************************************************************
 */

#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <limits.h>
#include <float.h>
#include <math.h>
#include <GL/gl.h>
#include <GL/freeglut.h>

///////// The following functions could be reused by copy/pasting them /////////

//// PMF types, defines and declarations (can be placed in a separate H file)

// Frame flags
#define FRAMEFLAG_NONE 0x00
#define FRAMEFLAG_KEYFRAME 0x01
#define FRAMEFLAG_EVENT 0x02

// Animation flags
#define ANIMFLAG_NONE 0x00
#define ANIMFLAG_LOOP 0x01

// Contains a single vertex
typedef struct {
    float p[3]; // Position
    float n[3]; // Normal
    float tc[2]; // Texture coordinates
} Vertex;

// Contains a single piece node and pointers to its children
typedef struct Piece {
    char name[32]; // Piece name
    float pos[3]; // Piece position
    Vertex* vertex; // Vertex data
    struct Piece** child; // Piece children
    uint32_t vertices; // Number of vertices (divisible by 3)
    uint32_t children; // Number of piece children
} Piece;

// A single animation frame made up of root position and piece rotations
typedef struct {
    float rootpos[3]; // Root node position
    float* rotations; // Node rotations as a series of X,Y,Z triplets
    uint8_t flags; // Frame flags
} Frame;

// A single animation
typedef struct {
    char name[32]; // Animation name
    Frame* frame; // Animation frames
    uint16_t frames; // Number of animation frames
    uint8_t fps; // Playback speed in frames per second
    uint8_t flags; // Animation flags
} Animation;

// A PMF model, contains the piece hierarchy and animations
typedef struct {
    Piece* root; // Root piece node
    Animation* animation; // Animations
    uint32_t animations; // Number of animations
    uint32_t pieces; // Number of pieces in the animation
} Model;

// Model animation state, stored separately from the model to allow multiple
// models with different states
typedef struct {
    float rootpos[3]; // Root position
    float* rotations; // Rotations
    uint32_t anim; // Animation index
    float timebase; // Animation start time base
} ModelState;

//// Function prototypes
Model* load_pmf_model(const char* filename);
void free_pmf_model(Model* pmf);
ModelState* create_pmf_model_state(Model* pmf);
void free_pmf_model_state(ModelState* ms);
void blend_pmf_model_state_transformations(Model* pmf, ModelState* ms, const float* sourcepos, const float* sourcerots, const float* targetpos, const float* targetrots, float blend);
void calc_pmf_model_state_transformations_at(Model* pmf, ModelState* ms, float seconds);
int set_pmf_model_state_animation(Model* pmf, ModelState* ms, const char* name, float timebase);

//// Function implementations (can be placed in a separate C file)

// Forward declarations
static void free_pmf_piece(Piece* piece);

// Load a name string stored as a Pascal shortstring (1 byte for len + 31 chars)
static int load_pmf_name(FILE* f, char name[32]) {
    if (fread(name, 32, 1, f) != 1) return 0;
    int len = name[0];
    memmove(name, name + 1, 31);
    name[len] = 0;
    return 1;
}

// Recursive function to call piece nodes
static Piece* load_pmf_model_piece(FILE* f, uint32_t* pieces) {
    // Create the piece
    Piece* piece = calloc(1, sizeof(Piece));
    if (!piece) return NULL;
    // Load the piece's name
    if (!load_pmf_name(f, piece->name)) goto fail;
    // Load the piece position
    if (fread(piece->pos, 12, 1, f) != 1) goto fail;
    // Load the number of vertices
    if (fread(&piece->vertices, 4, 1, f) != 1) goto fail;
    // Load vertex data
    piece->vertex = malloc(sizeof(Vertex) * piece->vertices);
    if (!piece->vertex) goto fail;
    for (uint32_t i = 0; i < piece->vertices; i++) {
        if (fread(piece->vertex[i].p, 12, 1, f) != 1) goto fail;
        if (fread(piece->vertex[i].n, 12, 1, f) != 1) goto fail;
        if (fread(piece->vertex[i].tc, 8, 1, f) != 1) goto fail;
    }
    // Load the number of children
    if (fread(&piece->children, 4, 1, f) != 1) goto fail;
    // Load child piece nodes
    piece->child = calloc(1, sizeof(Piece*) * piece->children);
    if (!piece->child) goto fail;
    for (uint32_t i = 0; i < piece->children; i++) {
        piece->child[i] = load_pmf_model_piece(f, pieces);
        if (!piece->child[i]) goto fail;
    }
    // Increase the PMF piece counter
    (*pieces)++;
    return piece;
fail:
    free_pmf_piece(piece);
    return NULL;
}

// Load a PMF model file
Model* load_pmf_model(const char* filename) {
    // Attempt to open the file
    char magic[4];
    FILE* f = fopen(filename, "rb");
    if (!f) return 0;
    // Check magic identifier
    if (fread(magic, 4, 1, f) != 1) {
        fclose(f);
        return NULL;
    }
    if (memcmp(magic, "PMF1", 4)) {
        fclose(f);
        return NULL;
    }

    // Create PMF object
    Model* pmf = calloc(1, sizeof(Model));
    if (!pmf) goto fail;
    // Load piece node hierarchy
    pmf->pieces = 0;
    pmf->root = load_pmf_model_piece(f, &pmf->pieces);
    if (!pmf->root) goto fail;
    // Load number of animations
    if (fread(&pmf->animations, 4, 1, f) != 1) goto fail;
    // Load animation data
    pmf->animation = calloc(1, sizeof(Animation) * pmf->animations);
    if (!pmf->animation) goto fail;
    for (uint32_t i = 0; i < pmf->animations; i++) {
        Animation* a = pmf->animation + i;
        // Load animation name
        if (!load_pmf_name(f, a->name)) goto fail;
        // Load animation info (flags, FPS, number of frames)
        if (fread(&a->flags, 1, 1, f) != 1) goto fail;
        if (fread(&a->fps, 1, 1, f) != 1) goto fail;
        if (fread(&a->frames, 2, 1, f) != 1) goto fail;
        // Load frame data
        a->frame = calloc(1, sizeof(Frame) * a->frames);
        if (!a->frame) goto fail;
        for (uint32_t j = 0; j < a->frames; j++) {
            // Load frame flags
            if (fread(&a->frame[j].flags, 1, 1, f) != 1) goto fail;
            // Load root piece position for this frame
            if (fread(a->frame[j].rootpos, 12, 1, f) != 1) goto fail;
            // Load piece rotations for this frame
            a->frame[j].rotations = malloc(12 * pmf->pieces);
            if (!a->frame[j].rotations) goto fail;
            if (fread(a->frame[j].rotations, 12, pmf->pieces, f) != pmf->pieces) goto fail;
        }
    }
    // Done
    fclose(f);
    return pmf;
fail:
    fclose(f);
    free_pmf_model(pmf);
    return NULL;
}

// Free a piece and its children
static void free_pmf_piece(Piece* piece) {
    if (!piece) return;
    for (uint32_t i = 0; i < piece->children; i++) {
        free_pmf_piece(piece->child[i]);
    }
    free(piece->child);
    free(piece->vertex);
    free(piece);
}

// Free a PMF model
void free_pmf_model(Model* pmf) {
    if (!pmf) return;
    // Free piece hierarchy
    free_pmf_piece(pmf->root);
    // Free animation data
    for (uint32_t i = 0; i < pmf->animations; i++) {
        Animation* a = pmf->animation + i;
        for (uint32_t j = 0; j < a->frames; j++) {
            free(a->frame[j].rotations);
        }
    }
    free(pmf->animation);
    free(pmf);
}

// Create a new model animation state for the given model
ModelState* create_pmf_model_state(Model* pmf) {
    // Create the state object
    ModelState* ms = calloc(1, sizeof(ModelState));
    if (!ms) return NULL;
    // Allocate zeroed memory for rotations
    ms->rotations = calloc(1, pmf->pieces * 12);
    if (!ms->rotations) {
        free(ms);
        return NULL;
    }
    return ms;
}

// Destroy the given model animation state
void free_pmf_model_state(ModelState* ms) {
    if (!ms) return;
    free(ms->rotations);
    free(ms);
}

// Blend the model transformations (position and rotations) using the given
// blending factor in the model state.  Both pairs for source and target values
// can be either frame data or the current model state data.
void blend_pmf_model_state_transformations(Model* pmf, ModelState* ms, const float* sourcepos, const float* sourcerots, const float* targetpos, const float* targetrots, float blend) {
    for (uint32_t i = 0; i < 3; i++) {
        ms->rootpos[i] = blend * targetpos[i] + (1.0f - blend) * sourcepos[i];
    }
    for (uint32_t i = 0; i < pmf->pieces * 3; i++) {
        ms->rotations[i] = blend * targetrots[i] + (1.0f - blend) * sourcerots[i];
    }
}

// Calculate the model state transformation for the given time in seconds
void calc_pmf_model_state_transformations_at(Model* pmf, ModelState* ms, float seconds) {
    if (ms->anim >= pmf->animations) return;
    Animation* a = pmf->animation + ms->anim;
    if (!a->frames) return;
    // Calculate current and next frame with sub-frame blending
    // (note: this actually runs a frame behind, but it looks fine in practice)
    float framef = (seconds - ms->timebase) * a->fps;
    uint32_t frame = (uint32_t)framef, next;
    float blend = framef - frame;
    if (a->flags & ANIMFLAG_LOOP) {
        frame = frame % a->frames;
        next = (frame + 1) % a->frames;
    } else {
        if (frame >= a->frames) {
            frame = a->frames - 1;
            next = frame;
        } else if (frame + 1 == a->frames) {
            next = frame;
        } else {
            next = frame + 1;
        }
    }
    // Blend the two frames
    blend_pmf_model_state_transformations(pmf, ms, a->frame[frame].rootpos, a->frame[frame].rotations, a->frame[next].rootpos, a->frame[next].rotations, blend);
}

// Set the model state transformation for the animation of the given name,
// shifting its time by the given time base (the timebase can be used to either
// start an animation from the beginning by setting it to the current seconds
// or by switching to another animation by using the current modelstate's
// timebase)
int set_pmf_model_state_animation(Model* pmf, ModelState* ms, const char* name, float timebase) {
    for (size_t i = 0; i < pmf->animations; i++) {
        if (!strcmp(name, pmf->animation[i].name)) {
            ms->anim = i;
            ms->timebase = timebase;
            if (pmf->animation[i].frames > 0) {
                memcpy(ms->rootpos, pmf->animation[i].frame[0].rootpos, 12);
                memcpy(ms->rotations, pmf->animation[i].frame[0].rotations, pmf->pieces * 12);
            } else {
                memset(ms->rootpos, 0, 12);
                memset(ms->rotations, 0, pmf->pieces * 12);
            }
            return 1;
        }
    }
    return 0;
}

////////// The following functions are specific to this OpenGL viewer //////////

// Utility function that calls set_pmf_model_state_animation using an animation
// index and the current time as reported by FreeGLUT to start an animation
// from the beginning
static void set_pmf_model_state_animation_by_index_for_glut(Model* pmf, ModelState* ms, int index) {
    if (pmf->animations > 0) {
        set_pmf_model_state_animation(pmf, ms, pmf->animation[index].name, (float)glutGet(GLUT_ELAPSED_TIME) / 1000.0f);
    } else {
        fprintf(stderr, "pmfview: warning: no animations in the model file\n");
    }
}

// Render a single model piece using OpenGL
static void render_pmf_model_piece(Piece* piece, ModelState* ms, const float* pos, const float** rotbase) {
    // Push the current transformation matrix so that the altered transformation
    // does not affect sibling pieces
    glPushMatrix();
    // Apply piece translation and rotation
    glTranslatef(pos[0], pos[1], pos[2]);
    glRotatef((*rotbase)[0], 1, 0, 0);
    glRotatef((*rotbase)[1], 0, 1, 0);
    glRotatef((*rotbase)[2], 0, 0, 1);
    // Advance the piece rotation pointer
    *rotbase += 3;
    // Render the piece
    glBegin(GL_TRIANGLES);
    for (uint32_t i = 0; i < piece->vertices; i++) {
        // Note: texture coordinates are ignored since the viewer does not load
        // any texture image data
        glNormal3fv(piece->vertex[i].n);
        glVertex3fv(piece->vertex[i].p);
    }
    glEnd();
    // Render children
    for (uint32_t i = 0; i < piece->children; i++) {
        render_pmf_model_piece(piece->child[i], ms, piece->child[i]->pos, rotbase);
    }
    // Pop the transformation matrix
    glPopMatrix();
}

// Render the given PMF model using OpenGL
static void render_pmf_model(Model* pmf, ModelState* ms) {
    // Define a piece rotation pointer that points to the current state's
    // rotations.  The pointer will be advanced linearly for each piece
    const float* rotbase = ms->rotations;
    // Render the root node.  Notice this uses the modelstate's root position
    // for `pos` instead of using the piece's own position like in recursive
    // calls inside render_pmf_model_piece as the root position can be animated
    render_pmf_model_piece(pmf->root, ms, ms->rootpos, &rotbase);
}

// Global variables for the viewer
static Model* pmf; // Model
static ModelState* ms; // Animation state
static float camrot = 45; // Camera rotation
static float camdist = 10; // Camera distance
static float camheight = 2; // Camera height
static int keystate[255]; // Keyboard up(1) or down(0) state
static int wireframe; // Display wireframe

// FreeGLUT callback for displaying the current frame
static void display() {
    // Clear color and depth buffers
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

    // Set the projection matrix for 3D rendering
    glMatrixMode(GL_PROJECTION);
    glLoadIdentity();
    gluPerspective(45.0f, (float)glutGet(GLUT_WINDOW_WIDTH) / (float)glutGet(GLUT_WINDOW_HEIGHT), 0.01f, 100.0f);

    // Set the modelview matrix for 3D rendering
    glMatrixMode(GL_MODELVIEW);
    glLoadIdentity();
    glTranslatef(0, -camheight, -camdist);
    glRotatef(camrot, 0, 1, 0);

    // Setup state
    glEnable(GL_DEPTH_TEST);

    // Draw origin
    glPointSize(4);
    glBegin(GL_POINTS);
    glColor3f(0.9f, 0.7f, 0.6f);
    glVertex3f(0, 0, 0);
    glEnd();
    glPointSize(1);

    // Draw XYZ axes
    glBegin(GL_LINES);
    glColor3f(1, 0, 0);
    glVertex3f(0, 0, 0);
    glVertex3f(1000, 0, 0);
    glColor3f(0.5f, 0, 0);
    glVertex3f(0, 0, 0);
    glVertex3f(-1000, 0, 0);
    glColor3f(0, 1, 0);
    glVertex3f(0, 0, 0);
    glVertex3f(0, 1000, 0);
    glColor3f(0, 0.5f, 0);
    glVertex3f(0, 0, 0);
    glVertex3f(0, -1000, 0);
    glColor3f(0, 0, 1);
    glVertex3f(0, 0, 0);
    glVertex3f(0, 0, 1000);
    glColor3f(0, 0, 0.5f);
    glVertex3f(0, 0, 0);
    glVertex3f(0, 0, -1000);
    glEnd();
    glColor3f(1, 1, 1);

    // Enable lighting
    glEnable(GL_LIGHTING);
    glEnable(GL_LIGHT0);

    // Render the model
    if (wireframe) glPolygonMode(GL_FRONT_AND_BACK, GL_LINE);
    render_pmf_model(pmf, ms);
    if (wireframe) glPolygonMode(GL_FRONT_AND_BACK, GL_FILL);

    // Disable lighting
    glDisable(GL_LIGHTING);

    // Reset the projection and modelview matrices for the 2D text overlay
    glMatrixMode(GL_PROJECTION);
    glLoadIdentity();
    glMatrixMode(GL_MODELVIEW);
    glLoadIdentity();

    // Draw the text overlay
    glTranslatef(-1, -1, 0);
    glScalef(0.0004f / ((float)glutGet(GLUT_WINDOW_WIDTH) / (float)glutGet(GLUT_WINDOW_HEIGHT)), 0.0005f, 0.0005f);
    glTranslatef(0, glutStrokeHeight(GLUT_STROKE_ROMAN) * 0.25f, 0);
    glPushMatrix();
    if (pmf->animations > 0) {
        static char tmp[256];
        const Animation* a = pmf->animation + ms->anim;
        sprintf(tmp, "Current animation: %s (%d FPS, %d frames, ", a->name, a->fps, a->frames);
        glutStrokeString(GLUT_STROKE_ROMAN, (unsigned char*)tmp);
        if (a->flags & ANIMFLAG_LOOP)
            glutStrokeString(GLUT_STROKE_ROMAN, (unsigned char*)"looping)");
        else
            glutStrokeString(GLUT_STROKE_ROMAN, (unsigned char*)"non-looping)");
    } else {
        glutStrokeString(GLUT_STROKE_ROMAN, (unsigned char*)"No animations");
    }
    glPopMatrix();
    glTranslatef(0, glutStrokeHeight(GLUT_STROKE_ROMAN), 0);
    glutStrokeString(GLUT_STROKE_ROMAN, (unsigned char*)"ADWSQE OPIK+- adjust camera   NM cycle   R restart   V wireframe   Right button menu");

    // Swap buffers to show the result
    glutSwapBuffers();
}

// FreeGLUT callback for window shape change
static void reshape(int width, int height) {
    // Resize the viewport to match the window size
    glViewport(0, 0, width, height);
}

// FreeGLUT callback for key press
static void keydown(unsigned char key, int x, int y) {
    keystate[key] = 1;
    switch (key) {
    case 'N':
    case 'n':
        set_pmf_model_state_animation_by_index_for_glut(pmf, ms, (pmf->animations + ms->anim - 1) % pmf->animations);
        break;
    case 'M':
    case 'm':
        set_pmf_model_state_animation_by_index_for_glut(pmf, ms, (pmf->animations + ms->anim + 1) % pmf->animations);
        break;
    case 'V':
    case 'v':
        wireframe = !wireframe;
        break;
    case 'r':
    case 'R':
        set_pmf_model_state_animation_by_index_for_glut(pmf, ms, ms->anim);
        break;
    case 27: // Exit when Escape is pressed
        glutLeaveMainLoop();
        break;
    }
}

// FreeGLUT callback for key release
static void keyup(unsigned char key, int x, int y) {
    keystate[key] = 0;
}

// FreeGLUT timer tick callback (called only once but rescheduled inside)
static void tick(int value) {
    // Update the camera based on key state
    if (keystate['a'] || keystate['A'] || keystate['o'] || keystate['O'])
        camrot -= 1;
    if (keystate['d'] || keystate['D'] || keystate['p'] || keystate['P'])
        camrot += 1;
    if (keystate['w'] || keystate['W'] || keystate['i'] || keystate['I'])
        camheight += 0.05f;
    if (keystate['s'] || keystate['S'] || keystate['k'] || keystate['K'])
        camheight -= 0.05f;
    if (keystate['q'] || keystate['Q'] || keystate['+'])
        camdist -= 0.05f;
    if (keystate['e'] || keystate['E'] || keystate['-'])
        camdist += 0.05f;

    // Calculate model animation state for the current time
    float current_time = (float)glutGet(GLUT_ELAPSED_TIME) / 1000.0f;
    calc_pmf_model_state_transformations_at(pmf, ms, current_time);

    // Cause a display refresh
    glutPostRedisplay();

    // Reapply the timer for ~60Hz
    glutTimerFunc(16, tick, 0);
}

// FreeGLUT callback for menu entry
static void menu(int anim) {
    // The menu contains only animations so if a user selects a menu entry it
    // is assumed to be an animation index
    set_pmf_model_state_animation_by_index_for_glut(pmf, ms, anim);
}

// Display an error message in stderr and exit
static void fatal(const char* msg) {
    fprintf(stderr, "pmfview: %s\n", msg);
    exit(1);
}

int main(int argc, char** argv) {
    const char* filename = NULL;
    const char* animname = NULL;

    // Parse the arguments
    for (int i = 1; i < argc; i++) {
        if (!strcmp(argv[i], "--help")) {
            puts("pmfview [-a animation] path-to-pmf-file");
            puts("Load and play the first animation in the given PMF file");
            puts("If -a animation is given, the animation with the given name will be played");
            return 0;
        } else if (argv[i][0] == '-') {
            if (!strcmp(argv[i], "-a")) {
                ++i;
                if (i >= argc) fatal("Missing animation name from -a");
                animname = argv[i];
            } else
                fatal("Unknown option");
        } else
            filename = argv[i];
    }
    if (!filename) {
        fprintf(stderr, "Missing filename. Use --help for help on how to use pmfview\n");
        return 1;
    }

    // Load the PMF model
    pmf = load_pmf_model(filename);
    if (!pmf) {
        fprintf(stderr, "Failed to load PMF file '%s'\n", filename);
        return 1;
    }
    ms = create_pmf_model_state(pmf);
    if (!ms) {
        fprintf(stderr, "Failed to create PMF model state\n");
        free_pmf_model(pmf);
        return 1;
    }

    // Initialize GLUT
    glutInit(&argc, argv);
    glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGB | GLUT_DEPTH);
    glutInitWindowSize(800, 600);
    glutCreateWindow("Piece Model File Viewer");
    glutGet(GLUT_ELAPSED_TIME);

    // Register callback functions
    glutDisplayFunc(display);
    glutReshapeFunc(reshape);
    glutKeyboardFunc(keydown);
    glutKeyboardUpFunc(keyup);

    // Start the timer for ~60Hz
    glutTimerFunc(16, tick, 0);

    // Create menu with animations
    glutCreateMenu(menu);
    for (uint32_t i = 0; i < pmf->animations; i++) {
        glutAddMenuEntry(pmf->animation[i].name, i);
    }
    glutAttachMenu(2);

    // Start the first animation
    if (animname) {
        if (!set_pmf_model_state_animation(pmf, ms, animname, 0)) {
            fprintf(stderr, "pmfview: warning: unknown animation '%s'\n", animname);
            set_pmf_model_state_animation_by_index_for_glut(pmf, ms, 0);
        }
    } else if (pmf->animations > 0) {
        set_pmf_model_state_animation_by_index_for_glut(pmf, ms, 0);
    }

    // Main loop
    glutMainLoop();

    // Release PMF data
    free_pmf_model_state(ms);
    free_pmf_model(pmf);

    return 0;
}