
#include "animate.h"
#include "../scene/renderer.h"
#include "manager.h"

#include <tuple>

namespace Gui {

Camera Anim_Camera::at(float t) const {
    Camera ret(dim);
    auto [p, r, f, a] = splines.at(t);
    Vec3 dir = r.rotate(Vec3{0.0f, 0.0f, -1.0f});
    ret.look_at(p + dir, p);
    ret.set_fov(f);
    ret.set_ar(a);
    return ret;
}

void Anim_Camera::set(float t, const Camera &cam) {
    splines.set(t, cam.pos(), Quat::euler(Mat4::rotate_z_to(cam.center() - cam.pos()).to_euler()),
                cam.get_fov(), cam.get_ar());
}

void Animate::update_dim(Vec2 dim) { ui_camera.dim(dim); }

bool Animate::keydown(Widgets &widgets, Undo &undo, Scene_ID sel, SDL_Keysym key) {

#ifdef __APPLE__
    if (key.sym == SDLK_BACKSPACE && key.mod & KMOD_GUI) {
#else
    if (key.sym == SDLK_DELETE) {
#endif
        if (joint_select)
            undo.del_bone(sel, joint_select);
        else if (sel)
            undo.del_obj(sel);
    }

    if (key.sym == SDLK_SPACE) {
        playing = !playing;
        last_frame = SDL_GetPerformanceCounter();
        return true;
    }

    return false;
}

void Animate::render(Scene &scene, Scene_Maybe obj_opt, Widgets &widgets, Camera &user_cam) {

    Mat4 view = user_cam.get_view();
    auto &R = Renderer::get();

    ui_camera.render(view);

    if (visualize_splines)
        for (auto &e : spline_cache)
            R.lines(e.second, view);

    if (!obj_opt.has_value())
        return;
    Scene_Item &item = obj_opt.value();

    if (item.is<Scene_Light>()) {
        Scene_Light &light = item.get<Scene_Light>();
        if (light.is_env())
            return;
    }

    Pose &pose = item.pose();
    float scale = std::min((user_cam.pos() - pose.pos).norm() / 5.5f, 10.0f);

    item.render(view);

    if (item.is<Scene_Object>() && item.get<Scene_Object>().armature.has_bones()) {

        Scene_Object &obj = item.get<Scene_Object>();
        joint_id_offset = scene.used_ids();
        obj.armature.render(view * obj.pose.transform(), joint_select, false, true,
                            joint_id_offset);

        if (!joint_select) {
            R.begin_outline();
            BBox box = obj.bbox();
            obj.render(view, false, true);
            obj.armature.outline(view * obj.pose.transform(), joint_select, false, true, box,
                                 joint_id_offset);
            R.end_outline(view, box);
        } else {
            widgets.active = Widget_Type::rotate;
        }

    } else {
        R.outline(view, item);
    }

    if (joint_select) {

        Scene_Object &obj = item.get<Scene_Object>();
        widgets.render(view, pose.transform() * obj.armature.posed_base_of(joint_select), scale);

    } else {
        widgets.render(view, pose.pos, scale);
    }
}

void Animate::make_spline(Scene_ID id, const Anim_Pose &pose) {

    if (!pose.splines.any())
        return;

    auto entry = spline_cache.find(id);
    if (entry == spline_cache.end()) {
        std::tie(entry, std::ignore) = spline_cache.insert({id, GL::Lines()});
    }

    GL::Lines &lines = entry->second;
    lines.clear();

    Vec3 prev = pose.at(0.0f).pos;
    for (int i = 1; i < max_frame; i++) {

        float f = (float)i;
        float c = (float)(i % 20) / 19.0f;
        Vec3 cur = pose.at(f).pos;
        lines.add(prev, cur, Vec3{c, c, 1.0f});
        prev = cur;
    }
}

void Animate::camera_spline() {

    auto entry = spline_cache.find(0);
    if (entry == spline_cache.end()) {
        std::tie(entry, std::ignore) = spline_cache.insert({0, GL::Lines()});
    }

    GL::Lines &lines = entry->second;
    lines.clear();

    Vec3 prev = anim_camera.at(0.0f).pos();
    for (int i = 1; i < max_frame; i++) {
        float f = (float)i;
        float c = (float)(i % 20) / 19.0f;
        Vec3 cur = anim_camera.at(f).pos();
        lines.add(prev, cur, Vec3{c, c, 1.0f});
        prev = cur;
    }
}

void Animate::UIsidebar(Manager &manager, Undo &undo, Scene_Maybe obj_opt, Camera &user_cam) {

    if (joint_select) {
        ImGui::Text("Edit Joint");
        if (ImGui::DragFloat3("Pose", joint_select->pose.data, 1.0f, 0.0f, 0.0f, "%.2f"))
            obj_opt.value().get().get<Scene_Object>().set_pose_dirty();
        if (ImGui::IsItemActivated())
            old_euler = joint_select->pose;
        if (ImGui::IsItemDeactivatedAfterEdit() && old_euler != joint_select->pose) {
            joint_select->pose = joint_select->pose.range(0.0f, 360.0f);
            undo.pose_bone(obj_opt.value().get().id(), joint_select, old_euler);
        }
        ImGui::Separator();
    }

    if (ui_camera.UI(undo, user_cam)) {
        camera_selected = true;
        if (obj_opt.has_value())
            prev_selected = obj_opt.value().get().id();
    }
}

void Animate::end_transform(Undo &undo, Scene_Item &obj) {
    if (joint_select) {
        undo.pose_bone(obj.id(), joint_select, old_euler);
    } else {
        undo.update_pose(obj.id(), old_pose);
    }
    old_pose = {};
    old_p_to_j = Mat4::I;
}

Vec3 Animate::selected_pos(Scene_Item &item) {
    if (joint_select) {
        return item.pose().transform() *
               item.get<Scene_Object>().armature.posed_base_of(joint_select);
    }
    return item.pose().pos;
}

void Animate::apply_transform(Widgets &widgets, Scene_Item &item) {
    if (joint_select) {
        Scene_Object &obj = item.get<Scene_Object>();
        Vec3 euler = widgets.apply_action(old_pose).euler;
        joint_select->pose = (old_p_to_j * Mat4::euler(euler)).to_euler();
        obj.set_pose_dirty();
    } else {
        item.pose() = widgets.apply_action(old_pose);
    }
}

bool Animate::select(Scene &scene, Widgets &widgets, Scene_ID selected, Scene_ID id, Vec3 cam,
                     Vec2 spos, Vec3 dir) {

    if (widgets.want_drag()) {

        if (joint_select) {

            Scene_Object &obj = scene.get_obj(selected);
            Vec3 base = obj.pose.transform() * obj.armature.posed_base_of(joint_select);
            widgets.start_drag(base, cam, spos, dir);

            Mat4 j_to_p = obj.pose.transform() * obj.armature.joint_to_posed(joint_select);
            old_euler = joint_select->pose;
            old_pose.euler = j_to_p.to_euler();
            j_to_p = j_to_p * Mat4::euler(old_euler).inverse();
            old_p_to_j = j_to_p.inverse();

        } else {

            Scene_Item &item = scene.get(selected).value();
            Pose &pose = item.pose();
            widgets.start_drag(pose.pos, cam, spos, dir);
            old_pose = pose;
        }

        return false;
    }

    Scene_Maybe mb = scene.get(selected);
    if (mb.has_value()) {

        Scene_Item &item = mb.value().get();
        if (item.is<Scene_Object>()) {

            Scene_Object &obj = item.get<Scene_Object>();
            if (id >= joint_id_offset && obj.armature.has_bones()) {

                Scene_ID j_id = id - joint_id_offset;
                joint_select = obj.armature.get_joint(j_id);
                if (joint_select) {
                    widgets.active = Widget_Type::rotate;
                }
                return false;
            }
        }
    }

    if (id >= n_Widget_IDs) {
        if (id == selected) {
            joint_select = nullptr;
        }
        return true;
    }

    joint_select = nullptr;
    return false;
}

void Animate::timeline(Manager &manager, Undo &undo, Scene &scene, Scene_Maybe obj,
                       Camera &user_cam) {

    // NOTE(max): this is pretty messy
    //      Would be good to add the ability to set per-component keyframes
    //      ^ I started with that but it was hard to make work with assimp and
    //        generally made everything a lot messier.

    ImVec2 size = ImGui::GetWindowSize();

    ImGui::Columns(2);
    ImGui::SetColumnWidth(0, 150.0f);

    if (!playing) {
        if (ImGui::Button("Play")) {
            playing = true;
            last_frame = SDL_GetPerformanceCounter();
        }
    } else {
        if (ImGui::Button("Stop")) {
            playing = false;
        }
    }

    ImGui::SameLine();
    if (ImGui::Button("Render")) {
        ui_render.open();
    }
    ui_render.animate(scene, ui_camera, user_cam, max_frame);

    ImGui::PushStyleVar(ImGuiStyleVar_ItemSpacing, {0.0f, 0.0f});
    ImGui::Dummy({1.0f, 4.0f});
    ImGui::PopStyleVar();

    if (ImGui::Button("Add Frames")) {
        max_frame += 3 * frame_rate;
    }
    if (ImGui::Button("Crop End")) {
        max_frame = current_frame + 1;
        current_frame = std::min(current_frame, max_frame - 1);
        scene.for_items([this](Scene_Item &item) {
            item.animation().splines.crop((float)max_frame);
            if (item.is<Scene_Object>())
                item.get<Scene_Object>().armature.crop((float)max_frame);
            else if (item.is<Scene_Light>())
                item.get<Scene_Light>().lanim.splines.crop((float)max_frame);
            make_spline(item.id(), item.animation());
        });
        set_time(scene, (float)current_frame);
    }
    ImGui::SliderInt("Rate", &frame_rate, 1, 240);
    frame_rate = clamp(frame_rate, 1, 240);

    ImGui::Checkbox("Draw Splines", &visualize_splines);

    ImGui::NextColumn();

    Scene_Item *select = nullptr;
    if (obj.has_value()) {
        Scene_Item &item = obj.value();
        select = &item;
        if (item.id() != prev_selected) {
            camera_selected = false;
            prev_selected = item.id();
        }
    }

    bool frame_changed = false;

    ImGui::Text("Keyframe:");
    ImGui::SameLine();
    if (ImGui::Button("Clear")) {
        if (camera_selected) {
            anim_camera.splines.erase((float)current_frame);
            camera_spline();
        } else if (select) {
            select->animation().splines.erase((float)current_frame);
            make_spline(select->id(), select->animation());
        }
    }

    auto set_item = [&, this](Scene_Item &item) {
        if (item.is<Scene_Object>()) {
            undo.anim_pose_bones(item.id(), (float)current_frame);
        }
        if (item.is<Scene_Light>()) {
            undo.anim_light(item.id(), (float)current_frame);
        }
        make_spline(item.id(), item.animation());
    };

    ImGui::SameLine();
    if (ImGui::Button("Set")) {
        if (camera_selected) {
            undo.anim_camera(anim_camera, (float)current_frame, ui_camera.get());
            camera_spline();
        } else if (select) {
            set_item(*select);
        }
    }
    ImGui::SameLine();
    if (ImGui::Button("Set All")) {
        undo.anim_camera(anim_camera, (float)current_frame, ui_camera.get());
        camera_spline();
        scene.for_items(set_item);
    }

    ImGui::Separator();
    ImGui::Dummy({74.0f, 1.0f});
    ImGui::SameLine();
    if (ImGui::SliderInt("Frame", &current_frame, 0, max_frame - 1)) {
        frame_changed = true;
    }
    ImGui::BeginChild("Timeline", {size.x - 20.0f, size.y - 80.0f}, false,
                      ImGuiWindowFlags_HorizontalScrollbar);

    ImGui::PushStyleVar(ImGuiStyleVar_ItemSpacing, {0.0f, 0.0f});

    const int name_chars = 6;
    std::vector<bool> frames;
    std::vector<Scene_ID> live_ids;

    {
        frames.clear();
        frames.resize(max_frame);

        std::string name = "Camera";
        ImVec2 sz = ImGui::CalcTextSize(name.c_str());
        if (camera_selected)
            ImGui::TextColored({Color::outline.x, Color::outline.y, Color::outline.z, 1.0f}, "%s",
                               name.c_str());
        else
            ImGui::Text("%s", name.c_str());
        ImGui::SameLine();
        ImGui::Dummy({80.0f - sz.x, 1.0f});
        ImGui::SameLine();
        ImGui::PushID("##CAMERA_");

        std::set<float> keys = anim_camera.splines.keys();

        for (float f : keys) {
            int frame = (int)std::round(f);
            if (frame >= 0 && frame < max_frame)
                frames[frame] = true;
        }

        for (int i = 0; i < max_frame; i++) {
            if (i > 0)
                ImGui::SameLine();
            ImGui::PushID(i);

            bool color = false;
            std::string label = "_";

            if (i == current_frame) {
                ImGui::PushStyleColor(ImGuiCol_Button, ImGui::GetColorU32(ImGuiCol_ButtonActive));
                color = true;
            }

            if (frames[i]) {
                label = "*";
                if (i != current_frame) {
                    color = true;
                    ImGui::PushStyleColor(ImGuiCol_Button,
                                          ImGui::GetColorU32(ImGuiCol_ButtonHovered));
                }
            }
            if (ImGui::SmallButton(label.c_str())) {
                current_frame = i;
                frame_changed = true;
                camera_selected = true;
            }
            if (color)
                ImGui::PopStyleColor();
            ImGui::PopID();
        }
        ImGui::PopID();
    }

    scene.for_items([&, this](Scene_Item &item) {
        frames.clear();
        frames.resize(max_frame);

        std::string name = const_cast<const Scene_Item &>(item).name();
        name.resize(name_chars);

        ImVec2 size = ImGui::CalcTextSize(name.c_str());
        if (!camera_selected && select && item.id() == select->id())
            ImGui::TextColored({Color::outline.x, Color::outline.y, Color::outline.z, 1.0f}, "%s",
                               name.c_str());
        else
            ImGui::Text("%s", name.c_str());
        ImGui::SameLine();
        ImGui::Dummy({80.0f - size.x, 1.0f});
        ImGui::SameLine();

        ImGui::PushID(item.id());

        Anim_Pose animation = item.animation();
        std::set<float> keys = animation.splines.keys();
        if (item.is<Scene_Light>()) {
            std::set<float> more_keys = item.get<Scene_Light>().lanim.splines.keys();
            keys.insert(more_keys.begin(), more_keys.end());
        }
        if (item.is<Scene_Object>()) {
            std::set<float> more_keys = item.get<Scene_Object>().armature.keys();
            keys.insert(more_keys.begin(), more_keys.end());
        }

        for (float f : keys) {
            int frame = (int)std::round(f);
            if (frame >= 0 && frame < max_frame)
                frames[frame] = true;
        }

        for (int i = 0; i < max_frame; i++) {
            if (i > 0)
                ImGui::SameLine();
            ImGui::PushID(i);

            bool color = false;
            std::string label = "_";

            if (i == current_frame) {
                ImGui::PushStyleColor(ImGuiCol_Button, ImGui::GetColorU32(ImGuiCol_ButtonActive));
                color = true;
            }

            if (frames[i]) {
                label = "*";
                if (i != current_frame) {
                    color = true;
                    ImGui::PushStyleColor(ImGuiCol_Button,
                                          ImGui::GetColorU32(ImGuiCol_ButtonHovered));
                }
            }

            if (ImGui::SmallButton(label.c_str())) {
                current_frame = i;
                frame_changed = true;
                camera_selected = false;
                manager.set_select(item.id());
            }
            if (color)
                ImGui::PopStyleColor();
            ImGui::PopID();
        }

        live_ids.push_back(item.id());

        ImGui::SameLine();
        ImGui::Dummy({142.0f, 1.0f});
        ImGui::PopID();
    });

    ImGui::PopStyleVar();
    ImGui::EndChild();

    std::unordered_map<Scene_ID, GL::Lines> new_cache;
    for (Scene_ID i : live_ids) {
        auto entry = spline_cache.find(i);
        if (entry != spline_cache.end()) {
            new_cache[i] = std::move(entry->second);
        }
    }
    spline_cache = std::move(new_cache);

    if (frame_changed)
        update(scene);
}

Camera Animate::set_time(Scene &scene, float time) {

    current_frame = (int)time;

    scene.for_items([time](Scene_Item &item) { item.set_time(time); });

    Camera cam = anim_camera.at(time);
    if (anim_camera.splines.any()) {
        ui_camera.load(cam.center(), cam.pos(), cam.get_ar(), cam.get_fov());
    } else {
        cam = ui_camera.get();
    }
    return cam;
}

void Animate::set(int n_frames, int fps) {
    max_frame = n_frames;
    frame_rate = fps;
    current_frame = std::min(current_frame, max_frame - 1);
}

Anim_Camera &Animate::camera() { return anim_camera; }

const Anim_Camera &Animate::camera() const { return anim_camera; }

float Animate::fps() const { return (float)frame_rate; }

int Animate::n_frames() const { return max_frame; }

std::string Animate::pump_output(Scene &scene) { return ui_render.step(*this, scene); }

void Animate::refresh(Scene &scene) { set_time(scene, (float)current_frame); }

void Animate::clear() {
    anim_camera.splines.clear();
    joint_select = nullptr;
}

void Animate::update(Scene &scene) {

    Uint64 time = SDL_GetPerformanceCounter();

    if (playing) {
        if ((time - last_frame) * frame_rate / SDL_GetPerformanceFrequency()) {
            if (current_frame == max_frame - 1) {
                playing = false;
                current_frame = 0;
            } else {
                current_frame++;
            }
            last_frame = time;
        }
    }

    if (displayed_frame != current_frame) {
        set_time(scene, (float)current_frame);
        displayed_frame = current_frame;
    }
}

} // namespace Gui
