#include "TimelinePanel.h"

#include <pxr/usd/usd/editContext.h>
#include <imgui.h>

#include <algorithm>
#include <cmath>

namespace UsdLayerManager {

// ---------------------------------------------------------------------------
void TimelinePanel::SetStage(pxr::UsdStageRefPtr stage)
{
    m_stage   = stage;
    m_playing = false;
    m_reversed = false;

    m_startFrame = 1.0;
    m_endFrame   = 100.0;
    m_fps        = 24.0;
    if (stage) {
        if (stage->HasAuthoredTimeCodeRange()) {
            m_startFrame = stage->GetStartTimeCode();
            m_endFrame   = stage->GetEndTimeCode();
        }
        double tps = stage->GetTimeCodesPerSecond();
        if (tps > 0.0) m_fps = tps;
    }
    if (m_endFrame < m_startFrame) m_endFrame = m_startFrame;

    m_currentFrame = m_startFrame;
    NotifyTimeChanged();
}

// ---------------------------------------------------------------------------
void TimelinePanel::Update(float deltaTime)
{
    if (!m_playing) return;

    double playStart = m_startFrame;
    double playEnd   = m_endFrame;
    double sign      = m_reversed ? -1.0 : 1.0;
    double next      = m_currentFrame + sign * double(deltaTime) * m_fps;

    if (m_reversed && next < playStart) {
        switch (m_loopMode) {
        case LoopMode::Loop: {
            double span = playEnd - playStart;
            next = (span > 0.0)
                ? playEnd - std::fmod(playStart - next, span)
                : playEnd;
            break;
        }
        case LoopMode::Bounce:
            next      = playStart + (playStart - next); // reflect
            m_reversed = false;
            break;
        default:
            next      = playStart;
            m_playing = false;
            break;
        }
    } else if (!m_reversed && next > playEnd) {
        switch (m_loopMode) {
        case LoopMode::Loop: {
            double span = playEnd - playStart;
            next = (span > 0.0)
                ? playStart + std::fmod(next - playStart, span)
                : playStart;
            break;
        }
        case LoopMode::Bounce:
            next      = playEnd - (next - playEnd); // reflect
            m_reversed = true;
            break;
        default:
            next      = playEnd;
            m_playing = false;
            break;
        }
    }

    SetCurrentFrame(next);
}

// ---------------------------------------------------------------------------
void TimelinePanel::Render()
{
    // Shared colours for active (pressed-in) toggle buttons.
    const ImVec4 kActive  (0.26f, 0.59f, 0.98f, 1.00f);
    const ImVec4 kActiveHv(0.36f, 0.69f, 1.00f, 1.00f);

    // Plain icon button (no active-state colouring).
    auto iconBtn = [&](const char* id, Icon icon, const char* fallback) -> bool {
        if (m_iconManager) {
            ImTextureID tex = m_iconManager->Get(icon);
            return ImGui::ImageButton(id, ImTextureRef(tex), ImVec2(18.f, 18.f));
        }
        return ImGui::Button(fallback);
    };

    // Icon button that highlights blue when `active` is true.
    auto iconToggle = [&](const char* id, Icon icon, const char* fallback,
                          bool active) -> bool {
        if (active) {
            ImGui::PushStyleColor(ImGuiCol_Button,        kActive);
            ImGui::PushStyleColor(ImGuiCol_ButtonHovered, kActiveHv);
        }
        bool clicked = iconBtn(id, icon, fallback);
        if (active) ImGui::PopStyleColor(2);
        return clicked;
    };

    // ---- Transport buttons ----
    if (iconBtn("##skipback", Icon::SkipBack, "|<")) {
        m_playing = false;
        m_reversed = false;
        SetCurrentFrame(m_startFrame);
    }
    ImGui::SameLine();

    if (iconBtn("##stepback", Icon::StepBack, "<")) {
        m_playing = false;
        SetCurrentFrame(std::floor(m_currentFrame) - 1.0);
    }
    ImGui::SameLine();

    // Play backward (rewind): clicking starts reverse play; clicking again pauses.
    {
        bool isPlayingBack = m_playing && m_reversed;
        if (iconToggle("##playback",
                       isPlayingBack ? Icon::Pause : Icon::PlayBack,
                       isPlayingBack ? "||" : "<|",
                       isPlayingBack)) {
            if (isPlayingBack) {
                m_playing = false;
            } else {
                m_reversed = true;
                m_playing  = true;
            }
        }
    }
    ImGui::SameLine();

    // Play forward: clicking starts forward play; clicking again pauses.
    {
        bool isPlayingFwd = m_playing && !m_reversed;
        if (iconToggle("##playfwd",
                       isPlayingFwd ? Icon::Pause : Icon::Play,
                       isPlayingFwd ? "||" : "|>",
                       isPlayingFwd)) {
            if (isPlayingFwd) {
                m_playing = false;
            } else {
                m_reversed = false;
                m_playing  = true;
            }
        }
    }
    ImGui::SameLine();

    if (iconBtn("##stepfwd", Icon::StepForward, ">")) {
        m_playing = false;
        SetCurrentFrame(std::floor(m_currentFrame) + 1.0);
    }
    ImGui::SameLine();

    if (iconBtn("##skipend", Icon::SkipEnd, ">|")) {
        m_playing  = false;
        m_reversed = false;
        SetCurrentFrame(m_endFrame);
    }
    ImGui::SameLine();

    // ---- Loop mode button — click to open Loop / Bounce / Off submenu ----
    ImGui::TextDisabled("|");
    ImGui::SameLine();
    {
        bool  isActive  = (m_loopMode != LoopMode::None);
        Icon  modeIcon  = (m_loopMode == LoopMode::Bounce) ? Icon::Bounce : Icon::Loop;
        const char* tip = (m_loopMode == LoopMode::Bounce) ? "Bounce (click to change)"
                        : (m_loopMode == LoopMode::Loop)   ? "Loop (click to change)"
                        :                                     "No loop (click to change)";
        if (iconToggle("##loopModeBtn", modeIcon, "Loop", isActive))
            ImGui::OpenPopup("##loopModePopup");
        if (ImGui::IsItemHovered()) ImGui::SetTooltip("%s", tip);

        if (ImGui::BeginPopup("##loopModePopup")) {
            if (ImGui::Selectable("Loop",   m_loopMode == LoopMode::Loop))
                m_loopMode = LoopMode::Loop;
            if (ImGui::Selectable("Bounce", m_loopMode == LoopMode::Bounce))
                m_loopMode = LoopMode::Bounce;
            ImGui::Separator();
            if (ImGui::Selectable("Off",    m_loopMode == LoopMode::None))
                m_loopMode = LoopMode::None;
            ImGui::EndPopup();
        }
    }
    ImGui::SameLine();

    // ---- FPS + Auto-Key ----
    ImGui::TextDisabled("|");
    ImGui::SameLine();
    ImGui::Text("%.6g fps", m_fps);
    ImGui::SameLine();
    if (ImGui::Checkbox("Auto-Key", &m_autoKey)) NotifyTimeChanged();

    // ---- Slider row: [start] [----slider----] [end] ----
    {
        const float kNumWidth = 50.f;
        const float kGap      =  4.f;

        int start = int(std::lround(m_startFrame));
        int end   = int(std::lround(m_endFrame));

        ImGui::SetNextItemWidth(kNumWidth);
        bool rangeEdited = ImGui::DragInt("##start", &start, 1.0f, 0, 0, "%d");
        if (ImGui::IsItemHovered()) ImGui::SetTooltip("Start frame");
        ImGui::SameLine(0.f, kGap);

        float sliderWidth = ImGui::GetContentRegionAvail().x - kNumWidth - kGap;
        float frame = float(m_currentFrame);
        ImGui::SetNextItemWidth(sliderWidth);
        if (ImGui::SliderFloat("##frame", &frame,
                               float(m_startFrame), float(m_endFrame), "%.0f")) {
            m_playing = false;
            SetCurrentFrame(std::round(double(frame)));
        }
        ImGui::SameLine(0.f, kGap);

        ImGui::SetNextItemWidth(kNumWidth);
        rangeEdited |= ImGui::DragInt("##end", &end, 1.0f, 0, 0, "%d");
        if (ImGui::IsItemHovered()) ImGui::SetTooltip("End frame");

        if (rangeEdited) {
            if (end < start) end = start;
            m_startFrame = double(start);
            m_endFrame   = double(end);
            if (m_stage) {
                pxr::UsdEditContext ec(m_stage, m_stage->GetRootLayer());
                m_stage->SetStartTimeCode(m_startFrame);
                m_stage->SetEndTimeCode(m_endFrame);
            }
            SetCurrentFrame(m_currentFrame);
        }
    }
}

// ---------------------------------------------------------------------------
void TimelinePanel::SetCurrentFrame(double frame)
{
    frame = std::min(std::max(frame, m_startFrame), m_endFrame);
    if (frame != m_currentFrame) {
        m_currentFrame = frame;
        NotifyTimeChanged();
    }
}

void TimelinePanel::NotifyTimeChanged()
{
    if (OnTimeChanged)
        OnTimeChanged(pxr::UsdTimeCode(m_currentFrame), EditTime());
}

pxr::UsdTimeCode TimelinePanel::EditTime() const
{
    return m_autoKey ? pxr::UsdTimeCode(m_currentFrame)
                     : pxr::UsdTimeCode::Default();
}

} // namespace UsdLayerManager
