#include "StageRuntime.hpp"

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <limits>

namespace openroller::psp {
namespace {

constexpr std::size_t kMaximumStageBytes = 4u * 1024u * 1024u;
constexpr float kPi = 3.14159265358979323846f;

void setError(char* output, std::size_t capacity, const char* message) {
    if (!output || capacity == 0) return;
    std::snprintf(output, capacity, "%s", message);
}

template <typename T>
bool sectionValid(const StagePackageHeader& header, PackageSection section) {
    if ((section.offset & 15u) != 0 || section.offset < header.headerSize) return false;
    if (section.count > std::numeric_limits<std::uint32_t>::max() / sizeof(T)) return false;
    const std::uint32_t bytes = section.count * static_cast<std::uint32_t>(sizeof(T));
    return section.offset <= header.fileSize && bytes <= header.fileSize - section.offset;
}

template <typename T>
const T* sectionPointer(const void* storage, PackageSection section) {
    const auto* bytes = static_cast<const std::uint8_t*>(storage);
    return reinterpret_cast<const T*>(bytes + section.offset);
}

bool rangeValid(PackageRange range, std::uint32_t count) {
    return range.first <= count && range.count <= count - range.first;
}

Vec3 add(Vec3 a, Vec3 b) { return {a.x + b.x, a.y + b.y, a.z + b.z}; }
Vec3 subtract(Vec3 a, Vec3 b) { return {a.x - b.x, a.y - b.y, a.z - b.z}; }
Vec3 multiply(Vec3 a, float value) { return {a.x * value, a.y * value, a.z * value}; }
float dot(Vec3 a, Vec3 b) { return a.x * b.x + a.y * b.y + a.z * b.z; }
Vec3 cross(Vec3 a, Vec3 b) {
    return {
        a.y * b.z - a.z * b.y,
        a.z * b.x - a.x * b.z,
        a.x * b.y - a.y * b.x,
    };
}
float length(Vec3 value) { return std::sqrt(dot(value, value)); }
Vec3 normalize(Vec3 value, Vec3 fallback = {0.0f, 0.0f, 0.0f}) {
    const float magnitude = length(value);
    return magnitude > 1.0e-6f ? multiply(value, 1.0f / magnitude) : fallback;
}
Vec3 mix(Vec3 a, Vec3 b, float u) { return add(a, multiply(subtract(b, a), u)); }
float mix(float a, float b, float u) { return a + (b - a) * u; }

Vec3 fromArray(const float value[3]) { return {value[0], value[1], value[2]}; }

std::uint32_t lowerTrackIndex(const StageView& stage, float timeMs) {
    const std::uint32_t count = stage.header->track.count;
    std::uint32_t first = 0;
    std::uint32_t last = count;
    while (first < last) {
        const std::uint32_t middle = first + (last - first) / 2;
        if (static_cast<float>(stage.track[middle].timeMs) <= timeMs) first = middle + 1;
        else last = middle;
    }
    return first == 0 ? 0 : first - 1;
}

std::uint32_t lowerCameraIndex(const StageView& stage, float timeMs) {
    const std::uint32_t count = stage.header->cameras.count;
    std::uint32_t first = 0;
    std::uint32_t last = count;
    while (first < last) {
        const std::uint32_t middle = first + (last - first) / 2;
        if (static_cast<float>(stage.cameras[middle].timeMs) <= timeMs) first = middle + 1;
        else last = middle;
    }
    return first == 0 ? 0 : first - 1;
}

Vec3 cameraOrbit(const PackageCameraPoint& camera) {
    const float a = (-camera.rotationA[1] * kPi / 180.0f) * 0.5f;
    const float b = (camera.rotationA[0] * kPi / 180.0f) * 0.5f;
    const float c = 0.0f;
    const float ca = std::cos(a), cb = std::cos(b), cc = std::cos(c);
    const float sa = std::sin(a), sb = std::sin(b), sc = std::sin(c);
    const float qw = sc * sa * sb + cc * ca * cb;
    const float qx = sc * ca * sb + cc * sa * cb;
    const float qy = cc * ca * sb - sc * sa * cb;
    const float qz = cc * sa * sb - sc * ca * cb;
    return multiply({
        2.0f * (qz * qx + qw * qy),
        2.0f * (qy * qz - qw * qx),
        1.0f - 2.0f * (qx * qx + qy * qy),
    }, camera.distance);
}

Vec3 rotateAround(Vec3 value, Vec3 axis, float degrees) {
    const float radians = degrees * kPi / 180.0f;
    const float cosine = std::cos(radians);
    const float sine = std::sin(radians);
    return add(
        add(multiply(value, cosine), multiply(cross(axis, value), sine)),
        multiply(axis, dot(axis, value) * (1.0f - cosine)));
}

void adjustCameraUp(CameraState* camera, float rollDegrees) {
    const Vec3 view = normalize(subtract(camera->target, camera->eye));
    if (dot(view, view) < 1.0e-10f) {
        camera->up = {0.0f, 1.0f, 0.0f};
        return;
    }
    Vec3 reference{0.0f, 1.0f, 0.0f};
    Vec3 projected = subtract(reference, multiply(view, dot(reference, view)));
    if (dot(projected, projected) < 1.0e-8f) {
        reference = {0.0f, 0.0f, 1.0f};
        projected = subtract(reference, multiply(view, dot(reference, view)));
    }
    camera->up = normalize(projected, {0.0f, 1.0f, 0.0f});
    if (rollDegrees != 0.0f) camera->up = rotateAround(camera->up, view, rollDegrees);
}

PackageCameraPoint mixCamera(const PackageCameraPoint& a, const PackageCameraPoint& b, float u) {
    PackageCameraPoint output = a;
    output.distance = mix(a.distance, b.distance, u);
    for (int i = 0; i < 2; ++i) output.rotationA[i] = mix(a.rotationA[i], b.rotationA[i], u);
    for (int i = 0; i < 3; ++i) {
        output.originOffset[i] = mix(a.originOffset[i], b.originOffset[i], u);
        output.fieldFar[i] = mix(a.fieldFar[i], b.fieldFar[i], u);
        output.fieldNear[i] = mix(a.fieldNear[i], b.fieldNear[i], u);
    }
    output.rotationB = mix(a.rotationB, b.rotationB, u);
    return output;
}

CameraState evaluateCameraInternal(
    const StageView& stage,
    float timeMs,
    bool interpolate,
    int depth) {
    CameraState state{};
    if (stage.header->cameras.count == 0 || depth > 8) {
        state.target = trackPositionAt(stage, timeMs);
        state.eye = add(state.target, {0.0f, 0.0f, 10.0f});
        state.projectionBlend = 1.0f;
        return state;
    }

    const std::uint32_t index = lowerCameraIndex(stage, timeMs);
    const PackageCameraPoint* key = &stage.cameras[index];
    PackageCameraPoint mixed{};
    const bool between = interpolate && key->fMode != 0 &&
        timeMs > static_cast<float>(key->timeMs) && index + 1 < stage.header->cameras.count &&
        timeMs < static_cast<float>(stage.cameras[index + 1].timeMs);
    if (between) {
        const PackageCameraPoint& following = stage.cameras[index + 1];
        const float span = static_cast<float>(following.timeMs - key->timeMs);
        const float u = span > 0.0f ? (timeMs - static_cast<float>(key->timeMs)) / span : 0.0f;
        if (key->fMode == 1) {
            CameraState from = evaluateCameraInternal(stage, static_cast<float>(key->timeMs), true, depth + 1);
            CameraState to = evaluateCameraInternal(stage, static_cast<float>(following.timeMs), true, depth + 1);
            adjustCameraUp(&from, 0.0f);
            adjustCameraUp(&to, 0.0f);
            state.eye = mix(from.eye, to.eye, u);
            state.target = mix(from.target, to.target, u);
            state.up = normalize(mix(from.up, to.up, u), {0.0f, 1.0f, 0.0f});
            state.projectionBlend = mix(from.projectionBlend, to.projectionBlend, u);
            const float roll = mix(key->rotationB, following.rotationB, u);
            const Vec3 axis = normalize(subtract(state.target, state.eye));
            if (roll != 0.0f && dot(axis, axis) > 0.0f) state.up = rotateAround(state.up, axis, roll);
            return state;
        }
        if (key->fMode == 2) {
            mixed = mixCamera(*key, following, u);
            key = &mixed;
        }
    }

    const Vec3 orbit = cameraOrbit(*key);
    const Vec3 currentTrack = trackPositionAt(stage, timeMs);
    const Vec3 origin = fromArray(key->originOffset);
    switch (key->aMode) {
        case 0:
            state.target = add(currentTrack, origin);
            state.eye = add(state.target, orbit);
            break;
        case 1:
            state.target = add(trackPositionAt(stage, static_cast<float>(key->timeMs)), origin);
            state.eye = add(state.target, orbit);
            break;
        case 2:
            if (index > 1) {
                return evaluateCameraInternal(
                    stage, static_cast<float>(stage.cameras[index].timeMs) - 1.0f, false, depth + 1);
            }
            state.target = add(currentTrack, origin);
            state.eye = add(state.target, orbit);
            break;
        case 3:
            state.target = add(currentTrack, origin);
            state.eye = add(add(trackPositionAt(stage, static_cast<float>(key->timeMs)), origin), orbit);
            break;
        case 4:
            state.target = fromArray(key->fieldFar);
            state.eye = fromArray(key->fieldNear);
            break;
        case 5:
            state.target = add(currentTrack, origin);
            state.eye = fromArray(key->fieldNear);
            break;
        case 6:
            state.target = fromArray(key->fieldFar);
            state.eye = add(currentTrack, orbit);
            break;
        default:
            state.target = add(currentTrack, origin);
            state.eye = add(state.target, orbit);
            break;
    }
    adjustCameraUp(&state, key->rotationB);
    state.projectionBlend = key->projectionType ? 1.0f : 0.0f;
    return state;
}

std::uint8_t colorChannel(std::uint32_t color, int shift) {
    return static_cast<std::uint8_t>((color >> shift) & 0xffu);
}

std::uint32_t mixColor(std::uint32_t a, std::uint32_t b, float u) {
    std::uint32_t output = 0;
    for (int shift = 0; shift <= 24; shift += 8) {
        const float value = mix(
            static_cast<float>(colorChannel(a, shift)),
            static_cast<float>(colorChannel(b, shift)), u);
        output |= static_cast<std::uint32_t>(std::clamp(value, 0.0f, 255.0f) + 0.5f) << shift;
    }
    return output;
}

} // namespace

bool loadStagePackage(const char* path, StageView* stage, char* error, std::size_t errorCapacity) {
    if (!path || !stage) {
        setError(error, errorCapacity, "invalid stage loader arguments");
        return false;
    }
    unloadStagePackage(stage);

    std::FILE* file = std::fopen(path, "rb");
    if (!file) {
        setError(error, errorCapacity, "stage.orps not found");
        return false;
    }
    if (std::fseek(file, 0, SEEK_END) != 0) {
        std::fclose(file);
        setError(error, errorCapacity, "could not seek stage.orps");
        return false;
    }
    const long length = std::ftell(file);
    if (length < static_cast<long>(sizeof(StagePackageHeader)) ||
        length > static_cast<long>(kMaximumStageBytes) ||
        std::fseek(file, 0, SEEK_SET) != 0) {
        std::fclose(file);
        setError(error, errorCapacity, "stage.orps has an invalid size");
        return false;
    }

    void* storage = std::malloc(static_cast<std::size_t>(length));
    if (!storage) {
        std::fclose(file);
        setError(error, errorCapacity, "not enough memory for stage.orps");
        return false;
    }
    const std::size_t read = std::fread(storage, 1, static_cast<std::size_t>(length), file);
    std::fclose(file);
    if (read != static_cast<std::size_t>(length)) {
        std::free(storage);
        setError(error, errorCapacity, "could not read complete stage.orps");
        return false;
    }

    const auto* header = static_cast<const StagePackageHeader*>(storage);
    const bool headerValid =
        std::memcmp(header->magic, kStagePackageMagic, sizeof(header->magic)) == 0 &&
        header->version == kStagePackageVersion &&
        header->headerSize == sizeof(StagePackageHeader) &&
        header->fileSize == static_cast<std::uint32_t>(length);
    const bool sectionsValid = headerValid &&
        sectionValid<PackageTrackPoint>(*header, header->track) &&
        sectionValid<PackageNote>(*header, header->notes) &&
        sectionValid<PackageCameraPoint>(*header, header->cameras) &&
        sectionValid<PackageDrawDistancePoint>(*header, header->drawDistances) &&
        sectionValid<PackageBackgroundColorPoint>(*header, header->backgroundColors) &&
        sectionValid<PackageBackgroundModel>(*header, header->backgroundModels) &&
        sectionValid<PackageBackgroundVertex>(*header, header->backgroundVertices) &&
        sectionValid<PackageBackgroundObject>(*header, header->backgroundObjects) &&
        sectionValid<PackageVisibilityKey>(*header, header->visibilityKeys) &&
        sectionValid<PackageTransformKey>(*header, header->transformKeys) &&
        sectionValid<PackageObjectColorKey>(*header, header->objectColorKeys) &&
        sectionValid<PackageParticlePoint>(*header, header->particles) &&
        sectionValid<PackageVisualizerPoint>(*header, header->visualizer) &&
        sectionValid<PackageBpmPoint>(*header, header->bpmChanges);
    bool contentsValid = sectionsValid;
    if (contentsValid) {
        const auto* models = sectionPointer<PackageBackgroundModel>(storage, header->backgroundModels);
        for (std::uint32_t i = 0; i < header->backgroundModels.count; ++i) {
            contentsValid = contentsValid &&
                rangeValid(models[i].triangles, header->backgroundVertices.count) &&
                rangeValid(models[i].solidLines, header->backgroundVertices.count) &&
                rangeValid(models[i].wireframeLines, header->backgroundVertices.count);
        }
        const auto* objects = sectionPointer<PackageBackgroundObject>(storage, header->backgroundObjects);
        for (std::uint32_t i = 0; i < header->backgroundObjects.count; ++i) {
            const bool parentValid = objects[i].parentIndex < 0 ||
                static_cast<std::uint32_t>(objects[i].parentIndex) < header->backgroundObjects.count;
            contentsValid = contentsValid &&
                (objects[i].model == std::numeric_limits<std::uint32_t>::max() ||
                 objects[i].model < header->backgroundModels.count) && parentValid &&
                rangeValid(objects[i].visibility, header->visibilityKeys.count) &&
                rangeValid(objects[i].movement, header->transformKeys.count) &&
                rangeValid(objects[i].scaling, header->transformKeys.count) &&
                rangeValid(objects[i].rotations, header->transformKeys.count) &&
                rangeValid(objects[i].colorChanges, header->objectColorKeys.count);
        }
    }
    if (!contentsValid || header->track.count < 2 ||
        header->backgroundObjects.count > kMaximumPackageBackgroundObjects) {
        std::free(storage);
        setError(error, errorCapacity, "stage.orps header or sections are invalid");
        return false;
    }

    stage->storage = storage;
    stage->storageSize = static_cast<std::size_t>(length);
    stage->header = header;
    stage->track = sectionPointer<PackageTrackPoint>(storage, header->track);
    stage->notes = sectionPointer<PackageNote>(storage, header->notes);
    stage->cameras = sectionPointer<PackageCameraPoint>(storage, header->cameras);
    stage->drawDistances = sectionPointer<PackageDrawDistancePoint>(storage, header->drawDistances);
    stage->backgroundColors = sectionPointer<PackageBackgroundColorPoint>(storage, header->backgroundColors);
    stage->backgroundModels = sectionPointer<PackageBackgroundModel>(storage, header->backgroundModels);
    stage->backgroundVertices = sectionPointer<PackageBackgroundVertex>(storage, header->backgroundVertices);
    stage->backgroundObjects = sectionPointer<PackageBackgroundObject>(storage, header->backgroundObjects);
    stage->visibilityKeys = sectionPointer<PackageVisibilityKey>(storage, header->visibilityKeys);
    stage->transformKeys = sectionPointer<PackageTransformKey>(storage, header->transformKeys);
    stage->objectColorKeys = sectionPointer<PackageObjectColorKey>(storage, header->objectColorKeys);
    stage->particles = sectionPointer<PackageParticlePoint>(storage, header->particles);
    stage->visualizer = sectionPointer<PackageVisualizerPoint>(storage, header->visualizer);
    stage->bpmChanges = sectionPointer<PackageBpmPoint>(storage, header->bpmChanges);
    setError(error, errorCapacity, "ok");
    return true;
}

void unloadStagePackage(StageView* stage) {
    if (!stage) return;
    std::free(stage->storage);
    *stage = {};
}

Vec3 trackPositionAt(const StageView& stage, float timeMs) {
    if (!stage.header || stage.header->track.count == 0) return {};
    if (timeMs <= static_cast<float>(stage.track[0].timeMs)) {
        return {stage.track[0].x, stage.track[0].y, stage.track[0].z};
    }
    const std::uint32_t last = stage.header->track.count - 1;
    if (timeMs >= static_cast<float>(stage.track[last].timeMs)) {
        return {stage.track[last].x, stage.track[last].y, stage.track[last].z};
    }
    const std::uint32_t index = lowerTrackIndex(stage, timeMs);
    const PackageTrackPoint& a = stage.track[index];
    const PackageTrackPoint& b = stage.track[index + 1];
    const float span = static_cast<float>(b.timeMs - a.timeMs);
    const float u = span > 0.0f ? (timeMs - static_cast<float>(a.timeMs)) / span : 0.0f;
    return mix({a.x, a.y, a.z}, {b.x, b.y, b.z}, u);
}

Vec3 trackTangentAt(const StageView& stage, float timeMs) {
    const float first = static_cast<float>(stage.track[0].timeMs);
    const float last = static_cast<float>(stage.track[stage.header->track.count - 1].timeMs);
    const Vec3 before = trackPositionAt(stage, std::max(first, timeMs - 4.0f));
    const Vec3 after = trackPositionAt(stage, std::min(last, timeMs + 4.0f));
    return normalize(subtract(after, before), {0.0f, 0.0f, 1.0f});
}

CameraState evaluateCamera(const StageView& stage, float timeMs) {
    return evaluateCameraInternal(stage, timeMs, true, 0);
}

BackgroundColors evaluateBackground(const StageView& stage, float timeMs) {
    if (!stage.header || stage.header->backgroundColors.count == 0) return {};
    const std::uint32_t count = stage.header->backgroundColors.count;
    std::uint32_t index = 0;
    while (index + 1 < count && static_cast<float>(stage.backgroundColors[index + 1].timeMs) <= timeMs) {
        ++index;
    }
    const PackageBackgroundColorPoint& a = stage.backgroundColors[index];
    if (index + 1 >= count) {
        return {a.topRightRgba, a.topLeftRgba, a.bottomRightRgba, a.bottomLeftRgba};
    }
    const PackageBackgroundColorPoint& b = stage.backgroundColors[index + 1];
    const float span = static_cast<float>(b.timeMs - a.timeMs);
    const float u = span > 0.0f
        ? std::clamp((timeMs - static_cast<float>(a.timeMs)) / span, 0.0f, 1.0f)
        : 0.0f;
    return {
        mixColor(a.topRightRgba, b.topRightRgba, u),
        mixColor(a.topLeftRgba, b.topLeftRgba, u),
        mixColor(a.bottomRightRgba, b.bottomRightRgba, u),
        mixColor(a.bottomLeftRgba, b.bottomLeftRgba, u),
    };
}

float evaluateDrawAhead(const StageView& stage, float timeMs) {
    if (!stage.header || stage.header->drawDistances.count == 0) {
        return stage.header ? stage.header->forwardDrawDistance : 7.0f;
    }
    const std::uint32_t count = stage.header->drawDistances.count;
    std::uint32_t index = 0;
    while (index + 1 < count && static_cast<float>(stage.drawDistances[index + 1].timeMs) <= timeMs) ++index;
    const PackageDrawDistancePoint& a = stage.drawDistances[index];
    if (index + 1 >= count) return std::max(0.0f, a.distance);
    const PackageDrawDistancePoint& b = stage.drawDistances[index + 1];
    const float span = static_cast<float>(b.timeMs - a.timeMs);
    const float u = span > 0.0f
        ? std::clamp((timeMs - static_cast<float>(a.timeMs)) / span, 0.0f, 1.0f)
        : 0.0f;
    return std::max(0.0f, mix(a.distance, b.distance, u));
}

float evaluateBeatDurationMs(const StageView& stage, float timeMs) {
    if (!stage.header || stage.header->bpmChanges.count == 0) return 500.0f;
    std::uint32_t active = 0;
    for (std::uint32_t i = 1; i < stage.header->bpmChanges.count; ++i) {
        if (static_cast<float>(stage.bpmChanges[i].timeMs) > timeMs) break;
        active = i;
    }
    const std::uint32_t bpm = stage.bpmChanges[active].bpm;
    return bpm > 0 ? 60000.0f / static_cast<float>(bpm) : 500.0f;
}

} // namespace openroller::psp