#include "TimelinePanel.h"

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

#include <algorithm>
#include <cmath>
#include <cstring>

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_isPlayblasting) return;  // time is stepped externally during capture
    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();
    ImGui::SameLine();
    ImGui::TextDisabled("|");
    ImGui::SameLine();

    // ---- Playblast button / progress ----
    if (m_isPlayblasting) {
        int total = m_playblastSettings.endFrame - m_playblastSettings.startFrame + 1;
        ImGui::Text("Capturing %d / %d", m_playblastFrameIdx, total);
        ImGui::SameLine();
        if (ImGui::Button("Cancel##pb")) AbortPlayblast();
    } else {
        if (ImGui::Button("Playblast")) {
            m_playblastSettings.startFrame = int(m_startFrame);
            m_playblastSettings.endFrame   = int(m_endFrame);
            m_showPlayblastDialog = true;
        }
    }

    // ---- Playblast settings dialog ----
    if (m_showPlayblastDialog) {
        ImGui::OpenPopup("Playblast##dlg");
        m_showPlayblastDialog = false;
    }
    if (ImGui::BeginPopupModal("Playblast##dlg", nullptr,
                               ImGuiWindowFlags_AlwaysAutoResize)) {
        // ── Output path ──────────────────────────────────────────────────────
        ImGui::Text("Output directory:");
        ImGui::SetNextItemWidth(340.f);
        ImGui::InputText("##outdir", m_playblastSettings.outputDir,
                         sizeof(m_playblastSettings.outputDir));
        ImGui::SameLine();
        if (ImGui::Button("Browse...##outdir")) {
            if (OnBrowseFolder) {
                std::string dir = OnBrowseFolder();
                if (!dir.empty()) {
                    strncpy(m_playblastSettings.outputDir, dir.c_str(),
                            sizeof(m_playblastSettings.outputDir) - 1);
                    m_playblastSettings.outputDir[sizeof(m_playblastSettings.outputDir) - 1] = '\0';
                }
            }
        }

        ImGui::Text("File prefix:");
        ImGui::SetNextItemWidth(200.f);
        ImGui::InputText("##prefix", m_playblastSettings.filePrefix,
                         sizeof(m_playblastSettings.filePrefix));

        // ── Frame range ───────────────────────────────────────────────────────
        ImGui::Separator();
        ImGui::SetNextItemWidth(90.f);
        ImGui::InputInt("Start frame##pb", &m_playblastSettings.startFrame);
        ImGui::SetNextItemWidth(90.f);
        ImGui::InputInt("End frame##pb",   &m_playblastSettings.endFrame);
        if (m_playblastSettings.endFrame < m_playblastSettings.startFrame)
            m_playblastSettings.endFrame = m_playblastSettings.startFrame;

        // ── Capture resolution ────────────────────────────────────────────────
        ImGui::Separator();
        ImGui::Text("Capture resolution:");
        ImGui::SetNextItemWidth(80.f);
        ImGui::InputInt("W##pbw", &m_playblastSettings.captureWidth);
        if (m_playblastSettings.captureWidth  < 16) m_playblastSettings.captureWidth  = 16;
        ImGui::SameLine();
        ImGui::Text("x");
        ImGui::SameLine();
        ImGui::SetNextItemWidth(80.f);
        ImGui::InputInt("H##pbh", &m_playblastSettings.captureHeight);
        if (m_playblastSettings.captureHeight < 16) m_playblastSettings.captureHeight = 16;

        // Common presets
        ImGui::SameLine();
        if (ImGui::Button("HD"))  { m_playblastSettings.captureWidth = 1280; m_playblastSettings.captureHeight = 720;  }
        ImGui::SameLine();
        if (ImGui::Button("FHD")) { m_playblastSettings.captureWidth = 1920; m_playblastSettings.captureHeight = 1080; }
        ImGui::SameLine();
        if (ImGui::Button("4K"))  { m_playblastSettings.captureWidth = 3840; m_playblastSettings.captureHeight = 2160; }

        // ── Movie export ──────────────────────────────────────────────────────
        ImGui::Separator();
        ImGui::Checkbox("Export movie (H.264 MP4 via libavcodec)", &m_playblastSettings.exportMovie);
        if (m_playblastSettings.exportMovie) {
            ImGui::Indent();
            ImGui::Checkbox("Also save BMP frames alongside movie", &m_playblastSettings.keepFrames);
            ImGui::Unindent();
        }

        ImGui::Separator();
        ImGui::TextDisabled("Writes: <dir>\\<prefix>.NNNN.bmp%s",
                            m_playblastSettings.exportMovie ? " → <prefix>.mp4" : "");

        // ── Buttons ───────────────────────────────────────────────────────────
        ImGui::Spacing();
        if (ImGui::Button("Cancel##dlg")) ImGui::CloseCurrentPopup();
        ImGui::SameLine();
        bool canStart = m_playblastSettings.outputDir[0] != '\0' &&
                        m_playblastSettings.filePrefix[0] != '\0';
        ImGui::BeginDisabled(!canStart);
        if (ImGui::Button("Start##dlg")) {
            ImGui::CloseCurrentPopup();
            m_playblastWasPlaying        = m_playing;
            m_playing                    = false;
            m_playblastSettings.fps      = m_fps;
            m_playblastFrame             = double(m_playblastSettings.startFrame);
            m_playblastFrameIdx          = 0;
            m_playblastSkipFirst         = true;
            m_isPlayblasting             = true;
            SetCurrentFrame(m_playblastFrame);
        }
        ImGui::EndDisabled();
        ImGui::EndPopup();
    }

    // ---- 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::SetCurrentFrameExternal(double frame)
{
    SetCurrentFrame(frame);
}

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

// ---------------------------------------------------------------------------
bool TimelinePanel::AdvancePlayblast()
{
    m_playblastFrame += 1.0;
    m_playblastFrameIdx++;
    if (m_playblastFrame > double(m_playblastSettings.endFrame)) {
        AbortPlayblast();
        return false;
    }
    SetCurrentFrame(m_playblastFrame);
    return true;
}

void TimelinePanel::AbortPlayblast()
{
    m_isPlayblasting = false;
    m_playing        = m_playblastWasPlaying;
}

} // namespace UsdLayerManager
