////////////////////////////////////////////////////////////////////////////////////////
//
//  Copyright 2026 OVITO GmbH, Germany
//
//  This file is part of OVITO (Open Visualization Tool).
//
//  OVITO is free software; you can redistribute it and/or modify it either under the
//  terms of the GNU General Public License version 3 as published by the Free Software
//  Foundation (the "GPL") or, at your option, under the terms of the MIT License.
//  If you do not alter this notice, a recipient may use your version of this
//  file under either the GPL or the MIT License.
//
//  You should have received a copy of the GPL along with this program in a
//  file LICENSE.GPL.txt.  You should have received a copy of the MIT License along
//  with this program in a file LICENSE.MIT.txt
//
//  This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY KIND,
//  either express or implied. See the GPL or the MIT License for the specific language
//  governing rights and limitations.
//
////////////////////////////////////////////////////////////////////////////////////////

#include <ovito/gui/desktop/GUI.h>
#include <ovito/gui/desktop/widgets/general/SpinnerWidget.h>
#include <ovito/gui/desktop/mainwin/MainWindow.h>
#include <ovito/gui/desktop/properties/PropertiesPanel.h>
#include <ovito/gui/desktop/properties/NumericalParameterUI.h>
#include <ovito/gui/desktop/dialogs/AnimationSettingsDialog.h>
#include <ovito/gui/base/actions/ActionManager.h>
#include <ovito/core/dataset/animation/controller/KeyframeController.h>
#include <ovito/core/utilities/units/UnitsManager.h>
#include "AnimationKeyEditorDialog.h"

namespace Ovito {

class NumericalItemDelegate : public QStyledItemDelegate
{
    ParameterUnit* _units;
    FloatType _minValue, _maxValue;

public:

    NumericalItemDelegate(QObject* parent, ParameterUnit* units, FloatType minValue, FloatType maxValue) : QStyledItemDelegate(parent), _units(units), _minValue(minValue), _maxValue(maxValue) {}

    virtual QString displayText(const QVariant& value, const QLocale& locale) const override {
        if(_units)
            return _units->formatValue(_units->nativeToUser(value.value<FloatType>()));
        return value.toString();
    }

    virtual QWidget* createEditor(QWidget* parent, const QStyleOptionViewItem& option, const QModelIndex& index) const override {
        QWidget* container = new QWidget(parent);
        QHBoxLayout* layout = new QHBoxLayout(container);
        layout->setSpacing(0);
        layout->setContentsMargins(0,0,0,0);
        QLineEdit* edit = new QLineEdit();
        edit->setFrame(false);
        layout->addWidget(edit, 1);
        container->setFocusProxy(edit);
        SpinnerWidget* spinner = new SpinnerWidget(nullptr, edit);
        if(_units) spinner->setUnit(_units);
        spinner->setMinValue(_minValue);
        spinner->setMaxValue(_maxValue);
        layout->addWidget(spinner);
        connect(spinner, &SpinnerWidget::valueChanged, [this,container]() {
            Q_EMIT const_cast<NumericalItemDelegate*>(this)->commitData(container);
        });
        return container;
    }

    virtual void setEditorData(QWidget* editor, const QModelIndex& index) const override {
        SpinnerWidget* spinner = editor->findChild<SpinnerWidget*>();
        QVariant data = index.data(Qt::EditRole);
        if(data.userType() == qMetaTypeId<FloatType>())
            spinner->setFloatValue(data.value<FloatType>());
        else if(data.userType() == qMetaTypeId<int>())
            spinner->setFloatValue(data.toInt());
    }

    virtual void setModelData(QWidget* editor, QAbstractItemModel* model, const QModelIndex& index) const override {
        SpinnerWidget* spinner = editor->findChild<SpinnerWidget*>();
        QVariant data = index.data(Qt::EditRole);
        if(data.userType() == qMetaTypeId<FloatType>())
            model->setData(index, QVariant::fromValue(spinner->floatValue()));
        else if(data.userType() == qMetaTypeId<int>())
            model->setData(index, QVariant::fromValue(spinner->intValue()));
    }

    virtual void updateEditorGeometry(QWidget* editor, const QStyleOptionViewItem& option, const QModelIndex& index) const override {
        editor->setGeometry(option.rect);
    }
};

class AnimationKeyModel : public QAbstractTableModel
{
public:

    /// Constructor.
    AnimationKeyModel(AnimationKeyEditorDialog& dialog, KeyframeController* ctrl, const PropertyFieldDescriptor* propertyField, QObject* parent = nullptr)
        : QAbstractTableModel(parent), _dialog(dialog), _animationSettings(dialog.activeAnimationSettings()), _propertyField(propertyField) {
        _ctrl.setTarget(ctrl);
        _ctrlType = ctrl->controllerType();
        _keys.setTargets(ctrl->keys());
        _ctrl.connect(this, &AnimationKeyModel::onCtrlEvent);
        _keys.connect(this, &AnimationKeyModel::onKeyEvent);
    }

    /// Returns the animation controller being edited.
    KeyframeController* ctrl() const { return _ctrl.target(); }

    /// Returns the list of animation keys.
    const auto& keys() const { return _keys.targets(); }

    /// Returns the number of rows.
    virtual int rowCount(const QModelIndex& parent = QModelIndex()) const override {
        if(parent.isValid()) return 0;
        return keys().size();
    }

    /// Returns the number of columns.
    virtual int columnCount(const QModelIndex& parent = QModelIndex()) const override {
        if(parent.isValid()) return 0;
        else if(_ctrlType == Controller::ControllerTypeFloat || _ctrlType == Controller::ControllerTypeInt) return 1;
        else if(_ctrlType == Controller::ControllerTypeVector3 || _ctrlType == Controller::ControllerTypePosition) return 3;
        else if(_ctrlType == Controller::ControllerTypeRotation) return 4;
        else return 0;
    }

    /// Returns the data stored.
    virtual QVariant data(const QModelIndex& index, int role = Qt::DisplayRole) const override {
        if(!index.isValid()) return QVariant();
        if(role == Qt::DisplayRole || role == Qt::EditRole) {
            int row = index.row();
            if(row >= 0 && row < keys().size()) {
                const AnimationKey* key = keys()[row];
                if(_ctrlType == Controller::ControllerTypeFloat || _ctrlType == Controller::ControllerTypeInt) {
                    return key->valueQVariant();
                }
                else if(_ctrlType == Controller::ControllerTypeVector3) {
                    Vector3 v = static_object_cast<Vector3AnimationKey>(key)->value();
                    return QVariant::fromValue(v[index.column()]);
                }
                else if(_ctrlType == Controller::ControllerTypePosition) {
                    Vector3 v = static_object_cast<PositionAnimationKey>(key)->value();
                    return QVariant::fromValue(v[index.column()]);
                }
                else if(_ctrlType == Controller::ControllerTypeRotation) {
                    Rotation r = static_object_cast<RotationAnimationKey>(key)->value();
                    if(index.column() < 3)
                        return QVariant::fromValue(r.axis()[index.column()]);
                    else
                        return QVariant::fromValue(r.angle());
                }
            }
        }
        return QVariant();
    }

    /// Returns the item flags for the given index.
    virtual Qt::ItemFlags flags(const QModelIndex& index) const override {
        if(!index.isValid())
            return QAbstractTableModel::flags(index);
        return Qt::ItemFlags(Qt::ItemIsSelectable | Qt::ItemIsEnabled | Qt::ItemIsEditable);
    }

    /// Sets the role data for the item at index to value.
    virtual bool setData(const QModelIndex& index, const QVariant& value, int role) override {
        if(index.isValid() && role == Qt::EditRole) {
            int row = index.row();
            if(row >= 0 && row < keys().size()) {
                return _dialog.performTransaction(tr("Change animation key value"), [&]() {
                    AnimationKey* key = keys()[row];
                    if(_ctrlType == Controller::ControllerTypeFloat || _ctrlType == Controller::ControllerTypeInt) {
                        key->setValueQVariant(value);
                    }
                    else if(_ctrlType == Controller::ControllerTypeVector3) {
                        Vector3 vec = static_object_cast<Vector3AnimationKey>(key)->value();
                        vec[index.column()] = value.value<FloatType>();
                        static_object_cast<Vector3AnimationKey>(key)->setValue(vec);
                    }
                    else if(_ctrlType == Controller::ControllerTypePosition) {
                        Vector3 vec = static_object_cast<PositionAnimationKey>(key)->value();
                        vec[index.column()] = value.value<FloatType>();
                        static_object_cast<PositionAnimationKey>(key)->setValue(vec);
                    }
                    else if(_ctrlType == Controller::ControllerTypeRotation) {
                        Rotation r = static_object_cast<RotationAnimationKey>(key)->value();
                        if(index.column() < 3) {
                            Vector3 axis = r.axis();
                            axis[index.column()] = value.value<FloatType>();
                            axis.normalizeSafely();
                            r.setAxis(axis);
                        }
                        else if(index.column() == 3) {
                            r.setAngle(value.value<FloatType>());
                        }
                        static_object_cast<RotationAnimationKey>(key)->setValue(r);
                    }
                });
            }
        }
        return false;
    }

    /// Returns the data for the given role and section in the header with the specified orientation.
    virtual QVariant headerData(int section, Qt::Orientation orientation, int role) const override {
        if(orientation == Qt::Horizontal && role == Qt::DisplayRole) {
            if(_ctrlType == Controller::ControllerTypeFloat || _ctrlType == Controller::ControllerTypeInt) {
                return _propertyField->displayName();
            }
            else if(_ctrlType == Controller::ControllerTypeVector3 || _ctrlType == Controller::ControllerTypePosition) {
                if(section == 0) return _propertyField->displayName() + QString(" (X)");
                else if(section == 1) return _propertyField->displayName() + QString(" (Y)");
                else if(section == 2) return _propertyField->displayName() + QString(" (Z)");
            }
            else if(_ctrlType == Controller::ControllerTypeRotation) {
                if(section == 0) return _propertyField->displayName() + QString(" (Axis X)");
                else if(section == 1) return _propertyField->displayName() + QString(" (Axis Y)");
                else if(section == 2) return _propertyField->displayName() + QString(" (Axis Z)");
                else if(section == 3) return _propertyField->displayName() + QString(" (Angle)");
            }
        }
        else if(orientation == Qt::Vertical && role == Qt::DisplayRole) {
            if(section >= 0 && section < keys().size() && _animationSettings) {
                const AnimationKey* key = keys()[section];
                return QVariant::fromValue(tr("Frame: %1").arg(_animationSettings->timeToString(key->time())));
            }
        }
        return QAbstractTableModel::headerData(section, orientation, role);
    }

    /// Is called when the animation controller generates a notification event.
    void onCtrlEvent(RefTarget* source, const ReferenceEvent& event) {
        if(event.type() == ReferenceEvent::ReferenceRemoved) {
            const ReferenceFieldEvent& refEvent = static_cast<const ReferenceFieldEvent&>(event);
            if(refEvent.field() == PROPERTY_FIELD(KeyframeController::keys)) {
                int index = keys().indexOf(static_object_cast<AnimationKey>(refEvent.oldTarget()));
                if(index >= 0) {
                    beginRemoveRows(QModelIndex(), index, index);
                    _keys.remove(index);
                    endRemoveRows();
                }
                OVITO_ASSERT(keys().size() == ctrl()->keys().size());
            }
        }
        else if(event.type() == ReferenceEvent::ReferenceAdded) {
            const ReferenceFieldEvent& refEvent = static_cast<const ReferenceFieldEvent&>(event);
            if(refEvent.field() == PROPERTY_FIELD(KeyframeController::keys)) {
                OVITO_ASSERT(keys().size() == ctrl()->keys().size() - 1);
                beginInsertRows(QModelIndex(), refEvent.index(), refEvent.index());
                _keys.insert(refEvent.index(), static_object_cast<AnimationKey>(refEvent.newTarget()));
                endInsertRows();
            }
        }
        else if(event.type() == ReferenceEvent::ReferenceChanged) {
            const ReferenceFieldEvent& refEvent = static_cast<const ReferenceFieldEvent&>(event);
            if(refEvent.field() == PROPERTY_FIELD(KeyframeController::keys)) {
                OVITO_ASSERT(keys().size() == ctrl()->keys().size());
                _keys.set(refEvent.index(), static_object_cast<AnimationKey>(refEvent.newTarget()));
                Q_EMIT dataChanged(createIndex(refEvent.index(), 0), createIndex(refEvent.index(), columnCount() - 1));
                Q_EMIT headerDataChanged(Qt::Vertical, refEvent.index(), refEvent.index());
            }
        }
    }

    /// Is called when an animation key generates a notification event.
    void onKeyEvent(RefTarget* source, const ReferenceEvent& event) {
        if(event.type() == ReferenceEvent::TargetChanged) {
            int index = keys().indexOf(static_object_cast<AnimationKey>(source));
            OVITO_ASSERT(index >= 0);
            Q_EMIT dataChanged(createIndex(index, 0), createIndex(index, columnCount() - 1));
            Q_EMIT headerDataChanged(Qt::Vertical, index, index);
        }
        else if(event.type() == ReferenceEvent::TargetDeleted) {
            int index = keys().indexOf(static_object_cast<AnimationKey>(source));
            OVITO_ASSERT(index >= 0);
            beginRemoveRows(QModelIndex(), index, index);
            _keys.remove(index);
            endRemoveRows();
        }
    }

private:

    /// The parent dialog window.
    AnimationKeyEditorDialog& _dialog;

    /// The active animation settings.
    OORef<AnimationSettings> _animationSettings;

    /// The controller whose keys are being edited.
    RefTargetListener<KeyframeController> _ctrl;

    /// The list of keys.
    VectorRefTargetListener<AnimationKey> _keys;

    /// The type of controller being edited.
    Controller::ControllerType _ctrlType;

    /// The property field being animated.
    const PropertyFieldDescriptor* _propertyField;
};

/******************************************************************************
* The constructor of the dialog widget.
******************************************************************************/
AnimationKeyEditorDialog::AnimationKeyEditorDialog(KeyframeController* ctrl, const PropertyFieldDescriptor* propertyField, QWidget* parent, MainWindowUI& ui) :
    QDialog(parent),
    UserInterfaceComponent<MainWindowUI>(ui),
    UndoableTransaction(ui, tr("Edit animatable parameter")),
    UndoSuspender(operation())
{
    setWindowTitle(tr("Parameter animation: %1").arg(propertyField->displayName()));
    _ctrl.setTarget(ctrl);

    // Make sure the controller has at least one animation key.
    if(ctrl->keys().empty()) {
        performActions(*this, [&] {
            try { ctrl->createKey(AnimationTime(0)); }
            catch(const Exception&) {}
        });
    }

    QVBoxLayout* mainLayout = new QVBoxLayout(this);

    mainLayout->addWidget(new QLabel(tr("Animation keys:")));
    _tableWidget = new QTableView();
    AnimationKeyModel* model = new AnimationKeyModel(*this, ctrl, propertyField, _tableWidget);
    _model = model;
    QItemSelectionModel* selModel = _tableWidget->selectionModel();
    _tableWidget->setModel(_model);
    delete selModel;
    _tableWidget->setSelectionBehavior(QAbstractItemView::SelectRows);
    _tableWidget->setEditTriggers(QAbstractItemView::AllEditTriggers);

    mainLayout->addStrut(model->columnCount() * 120 + 100);

    ParameterUnit* units = nullptr;
    FloatType minValue = FLOATTYPE_MIN;
    FloatType maxValue = FLOATTYPE_MAX;
    if(propertyField->numericalParameterInfo()) {
        minValue = propertyField->numericalParameterInfo()->minValue;
        maxValue = propertyField->numericalParameterInfo()->maxValue;
        if(propertyField->numericalParameterInfo()->unitType != nullptr)
            units = unitsManager().getUnit(propertyField->numericalParameterInfo()->unitType);
    }

    if(ctrl->controllerType() != Controller::ControllerTypeRotation) {
        NumericalItemDelegate* numericalDelegate = new NumericalItemDelegate(_tableWidget, units, minValue, maxValue);
        for(int col = 0; col < model->columnCount(); col++)
            _tableWidget->setItemDelegateForColumn(col, numericalDelegate);
    }
    else {
        NumericalItemDelegate* axisDelegate = new NumericalItemDelegate(_tableWidget, unitsManager().worldUnit(), FLOATTYPE_MIN, FLOATTYPE_MAX);
        NumericalItemDelegate* angleDelegate = new NumericalItemDelegate(_tableWidget, unitsManager().angleUnit(), FLOATTYPE_MIN, FLOATTYPE_MAX);
        _tableWidget->setItemDelegateForColumn(0, axisDelegate);
        _tableWidget->setItemDelegateForColumn(1, axisDelegate);
        _tableWidget->setItemDelegateForColumn(2, axisDelegate);
        _tableWidget->setItemDelegateForColumn(3, angleDelegate);
    }

    _tableWidget->resizeColumnsToContents();

    QHBoxLayout* hlayout = new QHBoxLayout();
    hlayout->setContentsMargins(0,0,0,0);
    hlayout->setSpacing(0);
    hlayout->addWidget(_tableWidget, 1);

    QToolBar* toolbar = new QToolBar();
    toolbar->setOrientation(Qt::Vertical);
    toolbar->setFloatable(false);
    _addKeyAction = toolbar->addAction(QIcon::fromTheme("animation_add_key"), tr("Create animation key"));
    connect(_addKeyAction, &QAction::triggered, this, &AnimationKeyEditorDialog::onAddKey);
    _deleteKeyAction = toolbar->addAction(QIcon::fromTheme("animation_delete_key"), tr("Delete animation key"));
    _deleteKeyAction->setEnabled(false);
    connect(_deleteKeyAction, &QAction::triggered, this, &AnimationKeyEditorDialog::onDeleteKey);

    toolbar->addSeparator();
    QAction* animSettingsAction = toolbar->addAction(QIcon::fromTheme("animation_settings"), tr("Animation settings..."));
    connect(animSettingsAction, &QAction::triggered, this, [this]() {
        AnimationSettingsDialog(this->ui(), this).exec();
    });

    hlayout->addWidget(toolbar);
    mainLayout->addLayout(hlayout);

    _keyPropPanel = new PropertiesPanel(ui, this);
    mainLayout->addWidget(_keyPropPanel);

    QDialogButtonBox* buttonBox = new QDialogButtonBox(QDialogButtonBox::Ok | QDialogButtonBox::Cancel | QDialogButtonBox::Help, Qt::Horizontal, this);
    QPushButton* animSettingsButton = buttonBox->addButton(tr("Animation settings..."), QDialogButtonBox::ActionRole);
    connect(buttonBox, &QDialogButtonBox::accepted, this, &AnimationKeyEditorDialog::onOk);
    connect(buttonBox, &QDialogButtonBox::rejected, this, &AnimationKeyEditorDialog::reject);

    // Handler for Animation Settings... button.
    connect(animSettingsButton, &QPushButton::clicked, actionManager()->getAction(ACTION_ANIMATION_SETTINGS), &QAction::trigger);

    // Handler for Help button.
    connect(buttonBox, &QDialogButtonBox::helpRequested, this, [this]() {
        actionManager()->openHelpTopic(QStringLiteral("manual:usage.animation"));
    });

    mainLayout->addWidget(buttonBox);

    connect(_tableWidget->selectionModel(), &QItemSelectionModel::selectionChanged, [this]() {
        QModelIndexList selection = _tableWidget->selectionModel()->selectedRows();
        if(_model->rowCount() > 1 && selection.empty() == false) {
            _deleteKeyAction->setEnabled(true);
            const QModelIndex& index = selection.first();
            OVITO_ASSERT(index.row() >= 0 && index.row() < this->ctrl()->keys().size());
            _keyPropPanel->setEditObject(this->ctrl()->keys()[index.row()]);
        }
        else {
            _deleteKeyAction->setEnabled(false);
            _keyPropPanel->setEditObject(nullptr);
        }
    });
    if(_model->rowCount() >= 0)
        _tableWidget->selectRow(_model->rowCount() - 1);
}

/******************************************************************************
* Event handler for the Ok button.
******************************************************************************/
void AnimationKeyEditorDialog::onOk()
{
    setFocus(); // Remove focus from child widgets to commit newly entered values in text widgets etc.
    commit();
    accept();
}

/******************************************************************************
* Handles the 'Add key' button.
******************************************************************************/
void AnimationKeyEditorDialog::onAddKey()
{
    QDialog dlg(this);
    dlg.setWindowTitle(tr("New animation key"));
    QVBoxLayout* mainLayout = new QVBoxLayout(&dlg);
    QHBoxLayout* subLayout = new QHBoxLayout();
    subLayout->setContentsMargins(0,0,0,0);
    subLayout->setSpacing(0);
    subLayout->addWidget(new QLabel(tr("Create key at animation time:")));
    subLayout->addSpacing(4);
    QLineEdit* timeEdit = new QLineEdit();
    subLayout->addWidget(timeEdit);
    SpinnerWidget* timeSpinner = new SpinnerWidget();
    timeSpinner->setTextBox(timeEdit);
    timeSpinner->setUnit(unitsManager().integerIdentityUnit());
    if(AnimationSettings* anim = activeAnimationSettings()) {
        timeSpinner->setIntValue(anim->currentFrame());
        timeSpinner->setMinValue(anim->firstFrame());
        timeSpinner->setMaxValue(anim->lastFrame());
    }
    else return;
    subLayout->addWidget(timeSpinner);
    mainLayout->addLayout(subLayout);
    QDialogButtonBox* buttonBox = new QDialogButtonBox(QDialogButtonBox::Ok | QDialogButtonBox::Cancel, Qt::Horizontal, &dlg);
    connect(buttonBox, &QDialogButtonBox::accepted, &dlg, &QDialog::accept);
    connect(buttonBox, &QDialogButtonBox::rejected, &dlg, &QDialog::reject);
    mainLayout->addWidget(buttonBox);
    if(dlg.exec() == QDialog::Accepted) {
        performActions(*this, [&]() {
            int index = ctrl()->createKey(AnimationTime::fromFrame(timeSpinner->intValue()));
            _tableWidget->selectRow(index);
        });
    }
}

/******************************************************************************
* Handles the 'Delete key' button.
******************************************************************************/
void AnimationKeyEditorDialog::onDeleteKey()
{
    QModelIndexList selection = _tableWidget->selectionModel()->selectedRows();
    performActions(*this, [&]() {
        QVector<OORef<AnimationKey>> keysToDelete;
        for(const QModelIndex& index : selection) {
            OVITO_ASSERT(index.row() >= 0 && index.row() < ctrl()->keys().size());
            keysToDelete.push_back(ctrl()->keys()[index.row()]);
        }
        ctrl()->deleteKeys(keysToDelete);
    });
}

}   // End of namespace
