/*============================================================================

The Medical Imaging Interaction Toolkit (MITK)

Copyright (c) German Cancer Research Center (DKFZ)
All rights reserved.

Use of this source code is governed by a 3-clause BSD license that can be
found in the LICENSE file.

============================================================================*/

#include <QmitkPipInstaller.h>

#include <mitkLog.h>
#include <mitkPythonHelper.h>

#include <nlohmann/json.hpp>

#include <QFile>
#include <QRegularExpression>

namespace
{
  // Quote arguments containing whitespace so the logged command line is
  // copy-pasteable into a shell for reproduction.
  QString FormatCommand(const QString& program, const QStringList& args)
  {
    const auto quote = [](const QString& token) {
      return token.contains(QChar::Space) ? '"' + token + '"' : token;
    };

    QString result = quote(program);
    for (const auto& arg : args)
    {
      result += ' ';
      result += quote(arg);
    }
    // Redact credentials from URLs: "://user:password@" -> "://****:****@"
    static QRegularExpression credentialPattern("://[^/@]+:[^/@]+@");
    result.replace(credentialPattern, "://****:****@");

    return result;
  }

  // Canonicalize a package name per PEP 503: lowercase, and treat runs of
  // '-', '_', '.' as a single '-'. Used to match resolved package names
  // against the original requirement specifiers.
  std::string CanonicalizePackageName(const std::string& name)
  {
    std::string result;
    result.reserve(name.size());

    for (char c : name)
    {
      if (c == '_' || c == '.' || c == '-')
      {
        if (!result.empty() && result.back() != '-')
          result.push_back('-');
      }
      else
      {
        result.push_back(static_cast<char>(std::tolower(static_cast<unsigned char>(c))));
      }
    }

    while (!result.empty() && result.back() == '-')
      result.pop_back();

    return result;
  }

  // Extract the package-name prefix of a requirement specifier.
  // Stops at the first character that is not part of a PEP 508 distribution name.
  std::string ExtractPackageName(const std::string& requirement)
  {
    std::size_t end = 0;
    while (end < requirement.size())
    {
      char c = requirement[end];
      if (std::isalnum(static_cast<unsigned char>(c)) || c == '_' || c == '.' || c == '-')
        ++end;
      else
        break;
    }
    return requirement.substr(0, end);
  }

  // Check whether a requirement is a direct reference (VCS URL or PEP 508
  // "name @ url" form) rather than a regular PyPI specifier.
  bool IsDirectReference(const std::string& req)
  {
    // Bare VCS URL: git+https://..., hg+https://..., etc.
    if (req.size() > 4 &&
        (req.compare(0, 4, "git+") == 0 || req.compare(0, 3, "hg+") == 0 ||
         req.compare(0, 4, "svn+") == 0 || req.compare(0, 4, "bzr+") == 0))
      return true;

    // PEP 508 direct reference: name and URL separated by an '@' with arbitrary
    // whitespace on either side (e.g. "name @ https://...", "name@url",
    // "name\t@ git+https://..."). Require at least one alphanumeric character
    // before the '@' and something after it to avoid false positives on stray
    // '@' characters.
    static const QRegularExpression directRef(R"(^\s*[A-Za-z0-9][\w\-.]*\s*@\s*\S)");
    return directRef.match(QString::fromStdString(req)).hasMatch();
  }
}

QmitkPipInstaller::QmitkPipInstaller(QObject* parent)
  : QObject(parent),
    m_Process(new QProcess(this))
{
  connect(m_Process, &QProcess::readyReadStandardOutput, this, &QmitkPipInstaller::OnStandardOutputReady);
  connect(m_Process, &QProcess::readyReadStandardError, this, &QmitkPipInstaller::OnStandardErrorReady);
  connect(m_Process, &QProcess::finished, this, &QmitkPipInstaller::OnProcessFinished);
  // QProcess::FailedToStart is the one error that does not also produce a
  // finished() signal, so we must drive the state machine from here for that
  // case. Every other error code (Crashed, Timedout, ...) is handled by
  // OnProcessFinished via the subsequent finished() emission.
  connect(m_Process, &QProcess::errorOccurred, this, [this](QProcess::ProcessError error) {
    if (error != QProcess::FailedToStart)
      return;

    m_State = State::Failed;
    emit ErrorOccurred("Failed to start pip subprocess.");
    emit InstallFinished(false);
  });
}

QmitkPipInstaller::~QmitkPipInstaller()
{
  if (m_Process->state() != QProcess::NotRunning)
    m_Process->kill();
}

void QmitkPipInstaller::SetInstallSpec(const mitk::PipInstallSpec& spec)
{
  if (this->IsRunning())
  {
    MITK_WARN << "Cannot change install spec while an operation is running.";
    return;
  }

  m_Spec = spec;
}

void QmitkPipInstaller::StartInstall()
{
  if (this->IsRunning())
  {
    MITK_WARN << "QmitkPipInstaller is already running.";
    return;
  }

  m_ResolvedPackages.clear();
  m_CurrentGroup = 0;
  m_CurrentPackage = 0;
  m_GroupStartIndex = 0;
  m_CurrentPostInstallStep = 0;
  m_AnyFailed = false;

  // Working copy of the install groups. A post-install step (PipInstallSpec::
  // postInstallSteps) must satisfy its own imports from packages installed by
  // these groups; the installer no longer injects any synthetic group.
  m_Groups = m_Spec.groups;

  // m_CreatedVirtualEnv is intentionally not reset here. If a prior attempt
  // created the venv, a retry reuses it (BeginVirtualEnvPhase sees it exists
  // and skips creation), so cleanup responsibility must survive the retry.
  // The flag is cleared only in RemoveCreatedVirtualEnv after the venv has
  // actually been removed.

  this->BeginVirtualEnvPhase();
}

void QmitkPipInstaller::AbandonInstall()
{
  // Only act after a failed install. After a successful install the venv
  // belongs to the consumer; while an operation is in progress the consumer
  // should call Cancel() instead.
  if (m_State != State::Failed)
    return;

  this->RemoveCreatedVirtualEnv();
}

void QmitkPipInstaller::Cancel()
{
  if (m_State == State::Idle || m_State == State::Done ||
      m_State == State::Failed || m_State == State::Cancelling)
    return;

  m_State = State::Cancelling;

  if (m_Process->state() != QProcess::NotRunning)
  {
    // Async path: kill the process and let OnProcessFinished call FinalizeCancel.
    m_Process->kill();
  }
  else
  {
    // No process running (we were in a synchronous transition between phases).
    // Finalize immediately.
    this->FinalizeCancel();
  }
}

bool QmitkPipInstaller::IsRunning() const
{
  return m_State != State::Idle && m_State != State::Done && m_State != State::Failed;
}

const std::vector<mitk::PipPackageInfo>& QmitkPipInstaller::GetResolvedPackages() const
{
  return m_ResolvedPackages;
}

void QmitkPipInstaller::OnStandardOutputReady()
{
  const auto output = QString::fromLocal8Bit(m_Process->readAllStandardOutput());
  emit OutputReceived(output, false);
}

void QmitkPipInstaller::OnStandardErrorReady()
{
  const auto output = QString::fromLocal8Bit(m_Process->readAllStandardError());
  emit OutputReceived(output, true);
}

void QmitkPipInstaller::OnProcessFinished(int exitCode, QProcess::ExitStatus exitStatus)
{
  // Cancellation: the process was killed by Cancel(). Run cleanup and emit the
  // terminal signal here, off the original Cancel() call stack, so the UI
  // thread is not blocked while we tear down the venv.
  if (m_State == State::Cancelling)
  {
    this->FinalizeCancel();
    return;
  }

  bool success = exitStatus == QProcess::NormalExit && exitCode == 0;
  auto numPackages = static_cast<int>(m_ResolvedPackages.size());

  switch (m_State)
  {
  case State::CreatingVirtualEnv:
  {
    if (!success)
    {
      m_State = State::Failed;
      emit ErrorOccurred("Failed to create virtual environment.");
      emit InstallFinished(false);
      return;
    }

    if (!mitk::PythonHelper::ActivateVirtualEnv(m_Spec.venvName))
    {
      m_State = State::Failed;
      emit ErrorOccurred("Failed to activate virtual environment.");
      emit InstallFinished(false);
      return;
    }

    if (m_Spec.upgradePipFirst)
    {
      this->StartPipUpgrade();
    }
    else
    {
      emit ResolveStarted();
      this->StartResolveGroup();
    }
    break;
  }

  case State::UpgradingPip:
  {
    // pip upgrade failure is non-fatal - proceed to resolve first group.
    m_State = State::Resolving;
    emit ResolveStarted();
    this->StartResolveGroup();
    break;
  }

  case State::Resolving:
  {
    if (!success)
    {
      m_State = State::Failed;
      emit ErrorOccurred("Could not resolve all dependencies.");
      emit InstallFinished(false);
      return;
    }

    if (!this->ParseResolveReport(m_ReportFile.fileName(), m_CurrentGroup))
    {
      // ParseResolveReport has already emitted ErrorOccurred with a specific message.
      m_State = State::Failed;
      emit InstallFinished(false);
      return;
    }

    numPackages = static_cast<int>(m_ResolvedPackages.size());

    // Install this group's packages before resolving the next group.
    // This ensures that when the next group is resolved, pip sees
    // the current group's packages as already installed (e.g. torch
    // from a CUDA index won't be re-resolved from PyPI).
    if (m_GroupStartIndex < numPackages)
    {
      m_State = State::Installing;
      m_CurrentPackage = m_GroupStartIndex;
      emit ProgressChanged(m_GroupStartIndex, numPackages);
      this->StartInstallPackage();
    }
    else
    {
      // No new packages in this group. Advance to next group.
      this->AdvanceToNextGroup();
    }
    break;
  }

  case State::Installing:
  {
    // Defensive: in the documented state machine m_CurrentPackage always
    // indexes into m_ResolvedPackages here, but guard against a stray slot
    // fire so we don't read out of range.
    if (m_CurrentPackage < 0 || m_CurrentPackage >= numPackages)
    {
      MITK_WARN << "OnProcessFinished(Installing) with out-of-range m_CurrentPackage; ignoring.";
      return;
    }

    const auto& pkg = m_ResolvedPackages[m_CurrentPackage];
    const auto name = QString::fromStdString(pkg.name);
    const auto status = success ? mitk::PackageStatus::Installed : mitk::PackageStatus::Failed;
    emit PackageStatusChanged(m_CurrentPackage, name, status);

    if (!success)
    {
      // Report the first package failure so the UI's one-shot terminal
      // status has something specific to display. Subsequent failures are
      // still logged via PackageStatusChanged / OutputReceived, but the
      // first one wins the headline.
      if (!m_AnyFailed)
        emit ErrorOccurred(QString("Failed to install %1.").arg(name));
      m_AnyFailed = true;
    }

    m_CurrentPackage++;
    emit ProgressChanged(m_CurrentPackage, numPackages);

    if (m_CurrentPackage < numPackages)
    {
      this->StartInstallPackage();
    }
    else
    {
      // All packages in the current group are installed.
      this->AdvanceToNextGroup();
    }
    break;
  }

  case State::RunningPostInstall:
  {
    const auto& step = m_Spec.postInstallSteps[m_CurrentPostInstallStep];

    if (!success && !step.optional)
    {
      const auto displayName = QString::fromStdString(
        step.displayName.empty() ? std::string("post-install step") : step.displayName);
      if (!m_AnyFailed)
        emit ErrorOccurred(QString("Failed to run %1.").arg(displayName));
      m_AnyFailed = true;
    }

    m_CurrentPostInstallStep++;

    if (m_CurrentPostInstallStep < static_cast<int>(m_Spec.postInstallSteps.size()))
    {
      this->StartPostInstallStep();
    }
    else
    {
      m_State = m_AnyFailed ? State::Failed : State::Done;
      emit InstallFinished(!m_AnyFailed);
    }
    break;
  }

  default:
    break;
  }
}

// Either kicks off a `python -m venv <path>` subprocess (asynchronous) or, if
// the venv is unnecessary or already exists, activates it and dispatches
// synchronously to the next phase.
void QmitkPipInstaller::BeginVirtualEnvPhase()
{
  if (m_Spec.venvName.empty() || mitk::PythonHelper::VirtualEnvExists(m_Spec.venvName))
  {
    // No venv needed or already exists - just activate and move on.
    if (!m_Spec.venvName.empty() && !mitk::PythonHelper::ActivateVirtualEnv(m_Spec.venvName))
    {
      m_State = State::Failed;
      emit ErrorOccurred("Failed to activate virtual environment.");
      emit InstallFinished(false);
      return;
    }

    if (m_Spec.upgradePipFirst)
    {
      this->StartPipUpgrade();
    }
    else
    {
      emit ResolveStarted();
      this->StartResolveGroup();
    }
    return;
  }

  const auto python = this->RequirePythonExecutable();

  if (python.isEmpty())
    return;

  m_State = State::CreatingVirtualEnv;
  m_CreatedVirtualEnv = true;
  emit VirtualEnvCreationStarted();

  const auto venvPath = mitk::PythonHelper::GetVirtualEnvPath(m_Spec.venvName);
  QStringList args = { "-m", "venv", QString::fromStdString(venvPath.string()) };
  MITK_INFO << FormatCommand(python, args).toStdString();
  m_Process->start(python, args);
}

void QmitkPipInstaller::StartPipUpgrade()
{
  const auto python = this->RequirePythonExecutable();

  if (python.isEmpty())
    return;

  m_State = State::UpgradingPip;
  emit PipUpgradeStarted();

  QStringList args = { "-m", "pip", "install", "--upgrade", "pip" };
  MITK_INFO << FormatCommand(python, args).toStdString();
  m_Process->start(python, args);
}

void QmitkPipInstaller::StartResolveGroup()
{
  const auto python = this->RequirePythonExecutable();

  if (python.isEmpty())
    return;

  if (m_CurrentGroup >= static_cast<int>(m_Groups.size()))
  {
    this->BeginPostInstallPhase();
    return;
  }

  m_GroupStartIndex = static_cast<int>(m_ResolvedPackages.size());
  const auto& group = m_Groups[m_CurrentGroup];

  // m_ReportFile is reused across groups: the underlying temp file path is
  // stable for the lifetime of this QmitkPipInstaller, but we close the Qt
  // handle each time so pip can write to the path freely. The file is deleted
  // when the QTemporaryFile object is destroyed.
  m_ReportFile.close();

  if (!m_ReportFile.open())
  {
    m_State = State::Failed;
    emit ErrorOccurred("Could not create temporary file for pip report.");
    emit InstallFinished(false);
    return;
  }

  const auto reportPath = m_ReportFile.fileName();
  m_ReportFile.close();

  QStringList args = { "-m", "pip", "install", "--dry-run", "--report", reportPath };
  args = this->BuildPipArgs(args, group);

  for (const auto& req : group.requirements)
    args.append(QString::fromStdString(req));

  m_State = State::Resolving;
  MITK_INFO << FormatCommand(python, args).toStdString();
  m_Process->start(python, args);
}

void QmitkPipInstaller::StartInstallPackage()
{
  const auto python = this->RequirePythonExecutable();

  if (python.isEmpty())
    return;

  if (m_CurrentPackage >= static_cast<int>(m_ResolvedPackages.size()))
  {
    this->BeginPostInstallPhase();
    return;
  }

  const auto& pkg = m_ResolvedPackages[m_CurrentPackage];
  emit PackageStatusChanged(m_CurrentPackage, QString::fromStdString(pkg.name), mitk::PackageStatus::Installing);

  const auto& group = m_Groups[pkg.group];

  // Direct references (VCS URLs, PEP 508 "name @ url") must be passed as-is
  // so pip fetches from the URL instead of searching PyPI for name==version.
  QString installArg;
  if (!pkg.specifier.empty() && IsDirectReference(pkg.specifier))
    installArg = QString::fromStdString(pkg.specifier);
  else
    installArg = QString::fromStdString(pkg.name) + "==" + QString::fromStdString(pkg.version);

  QStringList args = { "-m", "pip", "install", "--no-deps", installArg };
  args = this->BuildPipArgs(args, group);

  MITK_INFO << FormatCommand(python, args).toStdString();
  m_Process->start(python, args);
}

QStringList QmitkPipInstaller::BuildPipArgs(const QStringList& baseArgs, const mitk::PipInstallGroup& group) const
{
  QStringList args = baseArgs;

  if (!group.indexUrl.empty())
    args.append({ "--index-url", QString::fromStdString(group.indexUrl) });

  for (const auto& arg : group.extraPipArgs)
    args.append(QString::fromStdString(arg));

  return args;
}

bool QmitkPipInstaller::ParseResolveReport(const QString& reportPath, int groupIndex)
{
  QFile file(reportPath);

  if (!file.open(QIODevice::ReadOnly))
  {
    MITK_ERROR << "Could not open report file: " << reportPath.toStdString();
    emit ErrorOccurred("Could not open pip resolve report.");
    return false;
  }

  const auto data = file.readAll();
  file.close();

  // The report format is specified at
  // https://pip.pypa.io/en/stable/reference/installation-report/. We rely
  // on the outer catch to surface any deviation - every field we access
  // here is required by that spec (plus PEP 426 for name / version).
  // `requested` is the one genuinely optional field.
  try
  {
    const auto report = nlohmann::json::parse(data.constData(), data.constData() + data.size());
    const auto& group = m_Groups[groupIndex];

    for (const auto& entry : report.at("install"))
    {
      mitk::PipPackageInfo info;

      const auto& metadata = entry.at("metadata");
      info.name = metadata.at("name").get<std::string>();
      info.version = metadata.at("version").get<std::string>();
      info.requested = entry.value("requested", false);
      info.group = groupIndex;

      // Match the resolved package against this group's requirements using
      // PEP 503-canonicalized comparison so e.g. "scikit_learn", "scikit-learn"
      // and "Scikit-Learn" all match the same resolved name.
      const auto canonicalName = CanonicalizePackageName(info.name);

      for (const auto& req : group.requirements)
      {
        const auto reqName = CanonicalizePackageName(ExtractPackageName(req));
        if (reqName == canonicalName)
        {
          info.specifier = req;
          break;
        }
      }

      // Bare VCS URLs (git+https://...) have no extractable package name, so
      // the loop above cannot match them. Fall back to matching requested
      // packages against unclaimed direct-reference requirements.
      if (info.specifier.empty() && info.requested)
      {
        for (const auto& req : group.requirements)
        {
          if (!IsDirectReference(req))
            continue;

          bool alreadyClaimed = false;
          for (const auto& existing : m_ResolvedPackages)
          {
            if (existing.specifier == req)
            {
              alreadyClaimed = true;
              break;
            }
          }

          if (!alreadyClaimed)
          {
            info.specifier = req;
            break;
          }
        }
      }

      const auto index = static_cast<int>(m_ResolvedPackages.size());
      m_ResolvedPackages.push_back(info);
      emit PackageStatusChanged(index, QString::fromStdString(info.name), mitk::PackageStatus::Pending);
    }
  }
  catch (const nlohmann::json::exception& e)
  {
    MITK_ERROR << "Failed to parse report JSON: " << e.what();
    emit ErrorOccurred(QString("Failed to parse pip report: %1").arg(e.what()));
    return false;
  }

  return true;
}

void QmitkPipInstaller::AdvanceToNextGroup()
{
  m_CurrentGroup++;
  m_GroupStartIndex = static_cast<int>(m_ResolvedPackages.size());

  if (m_CurrentGroup < static_cast<int>(m_Groups.size()))
  {
    m_State = State::Resolving;
    emit ResolveStarted();
    this->StartResolveGroup();
  }
  else
  {
    this->BeginPostInstallPhase();
  }
}

// Called once all pip groups have finished. If no post-install steps are
// configured, immediately emits the terminal InstallFinished. Otherwise
// enters the RunningPostInstall state and kicks off the first step.
void QmitkPipInstaller::BeginPostInstallPhase()
{
  if (m_Spec.postInstallSteps.empty())
  {
    m_State = m_AnyFailed ? State::Failed : State::Done;
    emit InstallFinished(!m_AnyFailed);
    return;
  }

  m_State = State::RunningPostInstall;
  m_CurrentPostInstallStep = 0;
  this->StartPostInstallStep();
}

void QmitkPipInstaller::StartPostInstallStep()
{
  const auto python = this->RequirePythonExecutable();

  if (python.isEmpty())
    return;

  if (m_CurrentPostInstallStep >= static_cast<int>(m_Spec.postInstallSteps.size()))
  {
    m_State = m_AnyFailed ? State::Failed : State::Done;
    emit InstallFinished(!m_AnyFailed);
    return;
  }

  const auto& step = m_Spec.postInstallSteps[m_CurrentPostInstallStep];
  const auto displayName = QString::fromStdString(
    step.displayName.empty() ? std::string("post-install step") : step.displayName);

  // Run the step's Python source verbatim via `python -c`. QProcess::start with
  // a QStringList handles argument quoting on Windows. Any imports must be
  // satisfied by packages installed in the preceding groups; stdout/stderr
  // (including tqdm progress) flows to the UI details view via
  // OnStandardOutputReady / OnStandardErrorReady -> OutputReceived.
  const auto script = QString::fromStdString(step.pythonCode);

  emit PostInstallStepStarted(displayName);

  QStringList args = { "-c", script };
  MITK_INFO << FormatCommand(python, args).toStdString();
  m_Process->start(python, args);
}

QString QmitkPipInstaller::PythonExecutable() const
{
  const auto path = mitk::PythonHelper::GetExecutablePath();

  if (path.empty())
    return {};

  return QString::fromStdString(path.string());
}

QString QmitkPipInstaller::RequirePythonExecutable()
{
  const auto python = this->PythonExecutable();

  if (python.isEmpty())
  {
    m_State = State::Failed;
    emit ErrorOccurred("Python executable not found.");
    emit InstallFinished(false);
  }

  return python;
}

void QmitkPipInstaller::FinalizeCancel()
{
  m_State = State::Failed;
  this->RemoveCreatedVirtualEnv();
  emit InstallFinished(false);
}

void QmitkPipInstaller::RemoveCreatedVirtualEnv()
{
  if (m_CreatedVirtualEnv && !m_Spec.venvName.empty())
  {
    mitk::PythonHelper::RemoveVirtualEnv(m_Spec.venvName);
    m_CreatedVirtualEnv = false;
  }
}
