#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 done = false;

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

    if(instance->enemy_hp <= 0 || instance->enemy_mp <= 0) { // Enemy Defeated
        reward = 1;
        done = true;
    } else if(instance->enemy_mp <= 5 && rand() % 1 + 9 >= 7) { // Enemy Drain Player MP
        instance->player_mp -= rand() % 2 + 1;
    } 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() % 3 + 3;
    }

    if(instance->player_hp <= 0 || instance->player_mp <= 0 || instance->steps == 200) {
        reward = -1;
        done = 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->done = done;
}

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);
}