/* This source file is part of the Tomviz project, https://tomviz.org/.
   It is released under the 3-Clause BSD License, see "LICENSE". */

#include "PtychoRunner.h"

#include "CameraReaction.h"
#include "DataSource.h"
#include "LoadDataReaction.h"
#include "ProgressDialog.h"
#include "PtychoDialog.h"
#include "PythonUtilities.h"
#include "Utilities.h"

#include <QDebug>
#include <QDir>
#include <QFileInfo>
#include <QFutureWatcher>
#include <QMessageBox>
#include <QPointer>
#include <QtConcurrent>

namespace tomviz {

class PtychoRunner::Internal : public QObject
{
public:
  QPointer<PtychoRunner> parent;
  QPointer<QWidget> parentWidget;
  QPointer<PtychoDialog> ptychoDialog;
  QPointer<ProgressDialog> progressDialog;
  QFutureWatcher<bool> ptychoFutureWatcher;

  // Python modules and functions
  Python::Module ptychoModule;
  Python::Function stackPtychoFunc;

  // Ptycho options
  QString ptychoDirectory;
  QString outputDirectory;
  bool rotateDatasets = true;

  // These will be generated by calling functions
  QStringList versionList;
  QList<long> sidList;
  QList<double> angleList;

  QStringList outputFiles;

  // Recon options
  QStringList selectedArrays;

  bool autoLoadFinalData = true;

  Internal(PtychoRunner* p) : parent(p)
  {
    setParent(p);
    parentWidget = mainWidget();

    progressDialog = new ProgressDialog(parentWidget);
    progressDialog->setWindowTitle("Tomviz");
    progressDialog->showOutputWidget(true);
    progressDialog->resize(progressDialog->width(), 500);

    setupConnections();
  }

  void setupConnections()
  {
    connect(&ptychoFutureWatcher, &QFutureWatcher<bool>::finished, this,
            &Internal::stackPtychoFinished);
  }

  void importModule()
  {
    Python python;

    if (ptychoModule.isValid()) {
      return;
    }

    ptychoModule = python.import("tomviz.ptycho");
    if (!ptychoModule.isValid()) {
      qCritical() << "Failed to import \"tomviz.ptycho\" module";
    }
  }

  bool isInstalled()
  {
    importModule();

    Python python;

    auto installed = ptychoModule.findFunction("installed");
    if (!installed.isValid()) {
      qCritical() << "Failed to import \"tomviz.ptycho.installed\"";
      return false;
    }

    auto res = installed.call();

    if (!res.isValid()) {
      qCritical() << "Error calling \"tomviz.ptycho.installed\"";
      return false;
    }

    return res.toBool();
  }

  bool onFailure(const QString& msg)
  {
    qCritical() << msg;
    QMessageBox::critical(parentWidget, "Tomviz", msg);
    return false;
  }

  QString importError()
  {
    importModule();

    Python python;

    auto func = ptychoModule.findFunction("import_error");
    if (!func.isValid()) {
      qCritical() << "Failed to import \"tomviz.ptycho.import_error\"";
      return "import_error not found";
    }

    auto res = func.call();

    if (!res.isValid()) {
      qCritical() << "Error calling \"tomviz.ptycho.import_error\"";
      return "import_error not found";
    }

    return res.toString();
  }

  void importFunctions()
  {
    importModule();

    Python python;

    if (!stackPtychoFunc.isValid()) {
      stackPtychoFunc = ptychoModule.findFunction("load_stack_ptycho");
      if (!stackPtychoFunc.isValid()) {
        qCritical() << "Failed to find function \"load_stack_ptycho\"";
      }
    }
  }

  template <typename T>
  void clearWidget(QPointer<T> d)
  {
    if (!d) {
      return;
    }

    d->hide();
    d->deleteLater();
    d.clear();
  }

  void clear()
  {
    clearWidget(ptychoDialog);
  }

  void start()
  {
    clear();
    importFunctions();
    showPtychoDialog();
  }

  void showPtychoDialog()
  {
    clearWidget(ptychoDialog);

    ptychoDialog = new PtychoDialog(parentWidget);
    connect(ptychoDialog.data(), &QDialog::accepted, this,
            &Internal::ptychoDialogAccepted);
    ptychoDialog->show();
  }

  void ptychoDialogAccepted()
  {
    // Gather the settings and decide what to do
    ptychoDirectory = ptychoDialog->ptychoDirectory();
    sidList = ptychoDialog->selectedSids();
    versionList = ptychoDialog->selectedVersions();
    angleList = ptychoDialog->selectedAngles();
    outputDirectory = ptychoDialog->outputDirectory();
    rotateDatasets = ptychoDialog->rotateDatasets();

    runStackPtycho();
  }

  void runStackPtycho()
  {
    progressDialog->clearOutputWidget();
    progressDialog->setText("Stacking ptychography datasets...");
    progressDialog->show();
    auto future = QtConcurrent::run(std::bind(&Internal::_runStackPtycho, this));
    ptychoFutureWatcher.setFuture(future);
  }

  bool _runStackPtycho()
  {
    Python python;

    // Reset the output files
    outputFiles.clear();

    if (!stackPtychoFunc.isValid()) {
      qCritical() << "Failed to find function \"load_stack_ptycho\"";
      return false;
    }

    Python::Dict kwargs;
    kwargs.set("version_list", versionList);
    kwargs.set("sid_list", sidList);
    kwargs.set("angle_list", angleList);
    kwargs.set("ptycho_dir", ptychoDirectory);
    kwargs.set("output_dir", outputDirectory);
    kwargs.set("rotate_datasets", rotateDatasets);
    auto result = stackPtychoFunc.call(kwargs);

    if (!result.isValid() || !result.isList()) {
      qCritical() << "Error calling tomviz.ptycho.load_stack_ptycho";
      return false;
    }

    auto resultList = result.toList();
    for (int i = 0; i < resultList.length(); ++i) {
      outputFiles.append(resultList[i].toString());
    }

    return true;
  }

  void stackPtychoFinished()
  {
    progressDialog->accept();

    auto success = ptychoFutureWatcher.result();
    if (!success || !validateResult()) {
      QString msg = "Stack ptycho failed";
      onFailure(msg);
      // Show the dialog again
      showPtychoDialog();
      return;
    }

    qDebug() << "Results written to:";
    for (const auto& filePath : outputFiles) {
      qDebug() << "  " << filePath;
    }

    if (autoLoadFinalData) {
      loadOutputFiles();
    }
  }

  bool validateResult() const
  {
    if (outputFiles.isEmpty()) {
      return false;
    }

    for (auto& filePath: outputFiles) {
      if (!QFile::exists(filePath)) {
        return false;
      }
    }

    return true;
  }

  void loadOutputFiles()
  {
    // Convert sidList to QVector<int> for scan IDs
    QVector<int> scanIDs;
    scanIDs.reserve(sidList.size());
    for (auto& sid : sidList) {
      scanIDs.push_back(static_cast<int>(sid));
    }

    for (auto& filePath: outputFiles) {
      auto* dataSource = LoadDataReaction::loadData(filePath);
      if (!dataSource || !dataSource->imageData()) {
        qCritical() << "Failed to load file:" << filePath;
        return;
      }
      if (!scanIDs.isEmpty()) {
        dataSource->setScanIDs(scanIDs);
      }
    }

    // Automatically update camera to BNL convention
    CameraReaction::resetPositiveZ();
    CameraReaction::rotateCamera(-90);

    QString title = "Loading ptycho data complete";
    auto text =
      QString("Ptycho data in \"%1\" was written and loaded into Tomviz")
        .arg(outputDirectory);
    QMessageBox::information(parentWidget, title, text);
  }
};

PtychoRunner::PtychoRunner(QObject* parent)
  : QObject(parent), m_internal(new Internal(this))
{
}

PtychoRunner::~PtychoRunner() = default;

bool PtychoRunner::isInstalled()
{
  return m_internal->isInstalled();
}

QString PtychoRunner::importError()
{
  return m_internal->importError();
}

void PtychoRunner::start()
{
  m_internal->start();
}

void PtychoRunner::setAutoLoadFinalData(bool b)
{
  m_internal->autoLoadFinalData = b;
}

} // namespace tomviz
