#include "SS13ArcadeEnv.h"
#include <stdlib.h>

int clamp(int value, int min, int max) {
    if(value < min) return min;
    if(value > max) return max;
    return value;
}

void CreateEnvironment(struct SS13ArcadeEnv* env) {
    env->player_hp = 30;
    env->player_mp = 10;
    env->enemy_hp = 45;
    env->enemy_mp = 20;
    env->steps = 0;
}

void Step(struct SS13ArcadeEnv* instance, const int action, struct Observation* obs) {
    int reward = 0;
    bool terminated = false;
    bool truncated = false;

    switch (action) {
        case 0: // Attack
            instance->enemy_hp -= (rand() % 5) + 2; // 2-6
            break;
        case 1: // Heal
            instance->player_hp += (rand() % 3) + 6; // 6-8
            instance->player_mp -= (rand() % 3) + 1; // 1-3
            break;
        case 2: // Charge
            instance->player_mp += (rand() % 4) + 4; // 4-7
            break;
    }

    if(instance->enemy_hp <= 0 || instance->enemy_mp <= 0) { // Enemy Defeated
        reward = 1;
        terminated = true;
    } else if(instance->enemy_mp <= 5 && (rand() % 10) + 1 >= 7) { // Enemy Drain Player MP 1-10
        instance->player_mp -= (rand() % 2) + 2; // 2-3
    } else if(instance->enemy_hp <= 10 && instance->enemy_mp > 4) { // Enemy Heal
        instance->enemy_hp += 4;
        instance->enemy_mp -= 4;
    } else {
        instance->player_hp -= (rand() % 4) + 3; // 3-6
    }

    if(instance->player_hp <= 0 || instance->player_mp <= 0) {
        reward = -1;
        terminated = true;
    } else if(instance->steps == 200) {
        truncated = true;
    }
    instance->steps += 1;

    instance->player_hp = clamp(instance->player_hp, 0, 100);
    instance->player_mp = clamp(instance->player_mp, 0, 100);
    instance->enemy_hp = clamp(instance->enemy_hp, 0, 100);

    struct State state;
    GetState(instance, &state);

    obs->state = state;
    obs->reward = reward;
    obs->terminated = terminated;
    obs->truncated = truncated;
}

void GetState(struct SS13ArcadeEnv* instance, struct State* state) {
    state->player_hp = instance->player_hp;
    state->player_mp = instance->player_mp;
    state->enemy_hp = instance->enemy_hp;
    state->enemy_mp = instance->enemy_mp;
    state->steps = instance->steps;
}

void Reset(struct SS13ArcadeEnv* instance, struct State* state) {
    CreateEnvironment(instance);
    GetState(instance, state);
}