/*****************************************************************************
 * $CAMITK_LICENCE_BEGIN$
 *
 * CamiTK - Computer Assisted Medical Intervention ToolKit
 * (c) 2001-2025 Univ. Grenoble Alpes, CNRS, Grenoble INP - UGA, TIMC, 38000 Grenoble, France
 *
 * Visit http://camitk.imag.fr for more information
 *
 * This file is part of CamiTK.
 *
 * CamiTK is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License version 3
 * only, as published by the Free Software Foundation.
 *
 * CamiTK is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License version 3 for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * version 3 along with CamiTK.  If not, see <http://www.gnu.org/licenses/>.
 *
 * $CAMITK_LICENCE_END$
 ****************************************************************************/
#include "TotalSegmentator.h"

// CamiTK includes
#include <Property.h>
#include <ImageComponent.h>
#include <Application.h>
#include <Log.h>

// For the wizard and external process, Qt stuff
#include <QProcess>
#include <QTemporaryFile>
#include <QDir>
#include <QFileDialog>
#include <QFileInfo>
#include <QSettings>
#include <QWizard>
#include <QLabel>
#include <QVBoxLayout>
#include <QString>
#include <QLineEdit>
#include <QPushButton>
#include <QStyle>

using namespace camitk;

// --------------- defineTotalSegmentatorExecutable -------------------
/**
 * Create a wizard to select the executable TotalSegmentator and check if it is valid
 */
void TotalSegmentator::defineTotalSegmentatorExecutable() {
    //- Installation wizard
    QWizard wizard;
    wizard.setWindowTitle("Setting the TotalSegmentator Executable Manually");
    wizard.setMinimumSize(400, 300);

    //-- Introduction Page
    QWizardPage* introPage = new QWizardPage();
    introPage->setTitle("TotalSegmentator Executable Not Found");
    introPage->setSubTitle("Problem");

    QLabel* introLabel = new QLabel("<p>TotalSegmentator Executable was not found in the system and installation path. Please select the TotalSegmentator executable on your computer (in the 'bin' directory after installation).</p><p><b>Note:</b> you can download the latest TotalSegmentator version <a href=\"https://github.com/wasserth/TotalSegmentator\">from the official TotalSegmentator github pages</a>.", introPage);
    introLabel->setOpenExternalLinks(true);
    introLabel->setWordWrap(true);
    QVBoxLayout* introLayout = new QVBoxLayout();
    introLayout->addWidget(introLabel);
    introPage->setLayout(introLayout);

    wizard.addPage(introPage);

    //-- File Selection Page
    QWizardPage* fileSelectionPage = new QWizardPage();
    fileSelectionPage->setTitle("Select TotalSegmentator Executable");
    fileSelectionPage->setSubTitle("Please find and select the TotalSegmentator executable");

    QVBoxLayout* fileLayout = new QVBoxLayout();
    QLabel* fileLabel = new QLabel("Select the file 'TotalSegmentator':", fileSelectionPage);

    QString fileFilter;
#ifdef Q_OS_WIN
    fileFilter = "Executable Files (*.exe)";
#else
    fileFilter = "All Files (*)";
#endif
    QHBoxLayout* lineEditLayout = new QHBoxLayout();
    QLineEdit* fileLineEdit = new QLineEdit(fileSelectionPage);
    QPushButton* browseButton = new QPushButton("Browse", fileSelectionPage);

    //-- verification page labels and layout (updated during file selection)
    QLabel* executableLabel = new QLabel("TotalSegmentator Executable: undefined");
    QLabel* versionLabel = new QLabel("TotalSegmentator Version: not found");
    QHBoxLayout* warningLayout = new QHBoxLayout();
    QHBoxLayout* okLayout = new QHBoxLayout();
    QVBoxLayout* verificationLayout = new QVBoxLayout();

    // Connect the Browse button to a lambda opening a file selection dialog and checking the executable version
    connect(browseButton, &QPushButton::clicked, [this, &wizard, fileLineEdit, fileFilter, executableLabel, versionLabel, warningLayout, okLayout, verificationLayout]() {
        QString fileName = QFileDialog::getOpenFileName(&wizard, "Select TotalSegmentator Executable", QDir::homePath(), fileFilter);
        if (!fileName.isEmpty()) {
            // Set waiting cursor
            QApplication::setOverrideCursor(QCursor(Qt::WaitCursor));

            fileLineEdit->setText(fileName);
            executableLabel->setText(QString("TotalSegmentator Executable: %1").arg(fileName));
            QSettings& settings = Application::getSettings();
            settings.beginGroup("TotalSegmentator");
            settings.setValue("ExecutableFilePath", fileName);
            // run TotalSegmentator to check if everything is setup properly
            QProcess totalSegmentatorProcess;
            // ⚠ setting working directory is mandatory when using a python executable
            // without it, some python package display cryptic errors about missing libraries or files.
            totalSegmentatorProcess.setWorkingDirectory(QFileInfo(fileName).dir().absolutePath());
            totalSegmentatorProcess.setProgram(fileName);
            totalSegmentatorProcess.setArguments(QStringList{"--version"});
            totalSegmentatorProcess.start();

            if (totalSegmentatorProcess.waitForFinished()) {
                QString errors = totalSegmentatorProcess.readAllStandardError();
                if (errors.isEmpty()) {
                    // TotalSegmentator ran without error
                    versionLabel->setText(QString("TotalSegmentator Version: %1").arg(QString(totalSegmentatorProcess.readAllStandardOutput())));
                    verificationLayout->addLayout(okLayout);
                }
                else {
                    CAMITK_TRACE("Error calling TotalSegmentator --version : " + errors + "\n output: \n" + totalSegmentatorProcess.readAllStandardOutput())
                    // an error occurred
                    settings.setValue("ExecutableFilePath", "");
                    verificationLayout->addLayout(warningLayout);
                }
            }
            settings.endGroup();

            Application::restoreOverrideCursor();
        }
    });
    lineEditLayout->addWidget(fileLineEdit);
    lineEditLayout->addWidget(browseButton);
    fileLayout->addWidget(fileLabel);
    fileLayout->addLayout(lineEditLayout);
    fileSelectionPage->setLayout(fileLayout);

    wizard.addPage(fileSelectionPage);

    //-- Verification page
    QWizardPage* verificationPage = new QWizardPage();
    verificationPage->setTitle("TotalSegmentator Executable Verification");
    verificationLayout->addWidget(executableLabel);
    verificationLayout->addWidget(versionLabel);

    //-- Add message in case of error
    QLabel* warningIconLabel = new QLabel();
    warningIconLabel->setPixmap(Application::style()->standardIcon(QStyle::SP_MessageBoxWarning).pixmap(32, 32));
    QLabel* warningTextLabel = new QLabel("<p>An error occurred when testing the TotalSegmentator executable.</p>");
    warningTextLabel->setOpenExternalLinks(true);
    warningTextLabel->setWordWrap(true);
    warningLayout->addWidget(warningIconLabel);
    warningLayout->addWidget(warningTextLabel);
    warningIconLabel->setFixedWidth(32);

    //-- Add reinsuring label in case everything looks good
    QLabel* okIconLabel = new QLabel();
    okIconLabel->setPixmap(Application::style()->standardIcon(QStyle::SP_MessageBoxInformation).pixmap(32, 32));
    QLabel* okTextLabel = new QLabel("TotalSegmentator executable is functional.");
    okTextLabel->setWordWrap(true);
    okLayout->addWidget(okIconLabel);
    okLayout->addWidget(okTextLabel);
    warningIconLabel->setFixedWidth(32);

    // add general layout (which will be updated during the file selection
    // page depending on the TotalSegmentator execution results)
    verificationPage->setLayout(verificationLayout);

    // show the dialog to the used
    wizard.addPage(verificationPage);
    wizard.exec();
}

// -------------------- init --------------------
void TotalSegmentator::init() {
    // Set default parameter values
    setParameterValue("Task", 0);
    setParameterValue("Gpu", 1);
    setParameterValue("Fast", 1);
}

// -------------------- process --------------------
Action::ApplyStatus TotalSegmentator::process() {

    // Check whether we can find the executable in the settings
    QSettings& settings = Application::getSettings();
    settings.beginGroup("TotalSegmentator");
    QString executablePath = settings.value("ExecutableFilePath", QString()).toString();
    settings.endGroup();
    // If not, ask for it
    if (executablePath.isEmpty()) {
        // Ask for the executable location and add the settings for this Action
        defineTotalSegmentatorExecutable();
        settings.beginGroup("TotalSegmentator");
        // Now that the user may have entered it, read it again
        executablePath = settings.value("ExecutableFilePath", QString()).toString();
        // Store the current parameters in settings
        settings.setValue("Gpu", getParameterValue("Gpu"));
        settings.setValue("Fast", getParameterValue("Fast"));
        settings.setValue("Task", getParameterValue("Task"));
        settings.setValue("Generate Meshes", getParameterValue("Generate Meshes"));
        settings.endGroup();
    }
    // Even if there is a path, check that there is a file there
    if (!QFile(executablePath).exists()) {
        // Reset the setting
        settings.beginGroup("TotalSegmentator");
        settings.remove("ExecutableFilePath");
        settings.endGroup();
        return ERROR;
    }

    // Input ImageComponent
    Component* processedImage = getTargets().last();
    QString originalFilename = processedImage->getFileName();
    QString niftiFileName;
    bool deleteNifti;

    // Create a temporary nifti file if necessary and a temporary output file
    QTemporaryFile niftiFile(QDir::tempPath() + QDir::separator() + QFileInfo(originalFilename).baseName() + "-XXXXXX.nii");
    // we will remove it manually
    niftiFile.setAutoRemove(false);
    QTemporaryFile outputFile(QDir::tempPath() + QDir::separator() + QFileInfo(originalFilename).baseName() + "-out-XXXXXX.nii");
    outputFile.setAutoRemove(false);
    if (!outputFile.open()) {
        CAMITK_WARNING("Could not create temporary output file !")
        return ERROR;
    }

    // Check if the input is already a nifti, otherwise convert it
    if (!originalFilename.endsWith(".nii", Qt::CaseInsensitive) && !originalFilename.endsWith(".nii.gz", Qt::CaseInsensitive)) {
        // Rename the component to the temporary nifti path (but keep the old name)
        if (niftiFile.open()) {
            niftiFileName = niftiFile.fileName();
            deleteNifti = true;
            processedImage->setFileName(niftiFileName);
            if (!Application::save(processedImage)) {
                CAMITK_WARNING_ALT("Cannot save temporary Nifti file - Cannot run TotalSegmentator");
                // restore the correct path in the Component
                processedImage->setFileName(originalFilename);
                return Action::ApplyStatus::ERROR;
            }
            else {
                // Now that the nifti file is saved, restore the correct path in the Component
                processedImage->setFileName(originalFilename);
                if (!niftiFile.exists()) {
                    CAMITK_WARNING_ALT("Temporary Nifti file was not created - Cannot run TotalSegmentator");
                    return Action::ApplyStatus::ERROR;
                }
                CAMITK_TRACE_ALT(QString("Saved temporary Nifti copy %1").arg(niftiFileName));
            }
        }
        else {
            CAMITK_WARNING_ALT(QString("Cannot create temporary Nifti file %1 - Cannot run TotalSegmentator").arg(niftiFile.fileName()));
            return Action::ApplyStatus::ERROR;
        }
    }
    else {
        niftiFileName = originalFilename;
        deleteNifti = false;
    }

    QString outputFilePath = outputFile.fileName();
    totalSegmentatorProcess = new QProcess(); // Reset
    totalSegmentatorProcess->setProgram(executablePath);
    // ⚠ setting working directory is mandatory when using a python executable
    // without it, some python package display cryptic errors about missing libraries or files.
    totalSegmentatorProcess->setWorkingDirectory(QFileInfo(executablePath).dir().absolutePath());

    // Arguments for TotalSegmentator
    QStringList arguments;
    arguments << "-i" << niftiFileName
              << "-o" << outputFilePath
              << "--ml"; // Output one single image with multiple labels (cf https://github.com/wasserth/TotalSegmentator?tab=readme-ov-file#class-details)

    // Which task to perform ?
    int taskId = getParameterValue("Task").value<int>(); // enum index
    QString taskName = getProperty("Task")->getAttribute("enumNames").toStringList()[taskId];
    arguments << "--task" << taskName;

    // GPU or CPU
    if (getParameterValue("Gpu").toBool()) {
        arguments << "-d" << "gpu";
    }
    else {
        arguments << "-d" << "cpu";
    }

    // Speed
    switch (getParameterValue("Fast").value<int>()) { // enum as an int
        case 1: // Fast, 3mm resolution
            arguments << "--fast";
            break;

        case 2: //Fastest, 6mm resolution
            arguments << "--fastest";
            break;

        default: // Normal speed, full resolution
            break;
    }

    totalSegmentatorProcess->setArguments(arguments);

    // Debug command
    CAMITK_TRACE("Running: " + totalSegmentatorProcess->program() + " " + totalSegmentatorProcess->arguments().join(" "));
    totalSegmentatorProcess->setProcessChannelMode(QProcess::MergedChannels);

    // Define lambda to be run when the process finishes
    QObject::connect(totalSegmentatorProcess, QOverload<int, QProcess::ExitStatus>::of(&QProcess::finished), this,
    [this, outputFilePath, processedImage, niftiFileName, deleteNifti, taskName](int exitCode, QProcess::ExitStatus exitStatus) {

        if (exitStatus == QProcess::ExitStatus::CrashExit) {
            // Manage error
            CAMITK_ERROR_ALT(totalSegmentatorProcess->errorString());
        }
        else {
            // Display output in log
            CAMITK_INFO_ALT(totalSegmentatorProcess->readAll());
            CAMITK_INFO_ALT(QString("Saved in %1").arg(outputFilePath));
            // Open result file
            Component* labelImage = Application::open(outputFilePath);
            if (labelImage != nullptr) {
                // Set its frames
                labelImage->setFrameFrom(processedImage);
                // Remove its path and set it to modified, set a default name
                labelImage->setFileName("");
                labelImage->setModified(true);
                labelImage->setName(Application::getUniqueComponentName(processedImage->getName() + " " + taskName));

                // Generate meshes
                if (getParameterValue("Generate Meshes").toBool()) {
                    Action* reconstruction = Application::getAction("Meshes From Image");
                    if (reconstruction != nullptr) {
                        reconstruction->setInputComponent(labelImage);
                        reconstruction->setParameterValue("Selection Of Label Values", "");
                        reconstruction->apply();
                    }
                }
            }
        }
        // Remove the temporary nifti file and output file
        if (deleteNifti) {
            QFile niftiFile(niftiFileName);
            niftiFile.remove();
        }

        QFile outputFile(outputFilePath);
        outputFile.remove();

        // Remove waiting cursor
        Application::restoreOverrideCursor();

        // Refresh to show the new data
        refreshApplication();

        // TODO enable Apply Button again
    }); // end of QProcess::finished lambda

    // TODO Disable Apply Button
    // TODO: kill button if the process takes too much time ?

    // Set waiting cursor
    QApplication::setOverrideCursor(QCursor(Qt::WaitCursor));
    totalSegmentatorProcess->start();

    return SUCCESS;
}

// -------------------- targetDefined --------------------
void TotalSegmentator::targetDefined() {
    // When an image is selected, try to get the parameters' values from the settings
    QSettings& settings = Application::getSettings();
    settings.beginGroup("TotalSegmentator");
    if (settings.contains("Task")) {
        setParameterValue("Task", settings.value("Task"));
    }
    if (settings.contains("Fast")) {
        setParameterValue("Fast", settings.value("Fast"));
    }
    if (settings.contains("Gpu")) {
        setParameterValue("Gpu", settings.value("Gpu"));
    }
    if (settings.contains("Generate Meshes")) {
        setParameterValue("Generate Meshes", settings.value("Generate Meshes").toBool());
    }
    settings.endGroup();
}

// -------------------- parameterChanged --------------------
void TotalSegmentator::parameterChanged(QString parameterName) {
    // Save parameter value in settings
    QSettings& settings = Application::getSettings();
    settings.beginGroup("TotalSegmentator");
    settings.setValue(parameterName, getParameterValue(parameterName));
    settings.endGroup();
}

