#include <QHeaderView>
#include <QJsonDocument>
#include <QJsonObject>
#include <QTableView>
#include <QVBoxLayout>
#include <QMessageBox>
#include <QPushButton>
#include <QLabel>

#include "auth_flow.h"

#include "file_functions.h"

void AuthFlow::initialize(KeyringModel *model)
{
  m_model = model;
  m_saved_auth.load();
}

// This function is supposed to open the authentication dialog with various
// fields like method selection and fields to enter password or select a key,
// but for now, it'll simply submit a public key auth request.
void AuthFlow::start(QString user, QString host)
{
  if (m_phase == AuthPhase::requested)
  {
    reset();
  }

  m_username = user;
  m_hostname = host;
  m_key_dlg = new KeySelectDialog(m_model, m_username, m_hostname);
  m_key_dlg->setWindowModality(Qt::ApplicationModal);
  connect(m_key_dlg, &KeySelectDialog::key_selected, this, &AuthFlow::on_key_selected);
  connect(m_key_dlg, &QDialog::rejected, this, &AuthFlow::reset);

  auto saved_key_id = m_saved_auth.lookup(host.toUtf8(), user.toUtf8());
  if (!saved_key_id.isEmpty())
  {
    m_mode = FlowMode::Saved;
    on_key_selected(acquire_keyring_key(saved_key_id));
  }
  else
  {
    m_mode = FlowMode::Default;
    m_key_dlg->show();
  }
}

KeySelectDialog::KeySelectDialog(KeyringModel *model, QStringView username, QStringView hostname, QWidget *parent)
  : QDialog(parent)
    , m_model(model)
{
  this->setWindowTitle(QString("Select key for %1@%2").arg(username, hostname));
  auto view = new QTableView(this);
  view->setModel(m_model);
  view->horizontalHeader()->setSectionResizeMode(QHeaderView::Stretch);
  view->setSelectionBehavior(QAbstractItemView::SelectRows);
  view->setSelectionMode(QAbstractItemView::SingleSelection);
  QDialogButtonBox *buttons = new QDialogButtonBox(QDialogButtonBox::Ok | QDialogButtonBox::Cancel, this);
  buttons->button(QDialogButtonBox::Ok)->setEnabled(false);
  QVBoxLayout *layout = new QVBoxLayout(this);
  layout->addWidget(view);
  layout->addWidget(buttons);
  connect(buttons, &QDialogButtonBox::rejected, this, &QDialog::reject);
  connect(buttons, &QDialogButtonBox::accepted, this, [this, view] {
    auto rows = view->selectionModel()->selectedRows();
    if (rows.isEmpty())
    {
      return;
    }

    QByteArray selected_key_id = m_model->data(rows.first(), KeyringModel::KeyIDRole).toByteArray();

    emit key_selected(acquire_keyring_key(selected_key_id));
  });
  connect(view->selectionModel(), &QItemSelectionModel::selectionChanged, this, [view, buttons](const QItemSelection &, const QItemSelection &) {
    bool selected = !view->selectionModel()->selectedRows().isEmpty();
    buttons->button(QDialogButtonBox::Ok)->setEnabled(selected);
  });
}

KeyActivationDialog::KeyActivationDialog(Mode mode, QStringView key_name, QWidget *parent)
  : QDialog(parent)
{
  this->setWindowTitle(QString("Unlock key %1").arg(key_name));
  QVBoxLayout *pw_layout = new QVBoxLayout(this);

  QString mode_prompt;
  switch (mode) {
    case Mode::passphrase:
      mode_prompt = QStringLiteral("Passphrase:");
      break;
    case Mode::pin:
      mode_prompt = QStringLiteral("PIN:");
      break;
    default:
      break;
  }

  pw_layout->addWidget(new QLabel(mode_prompt, this));
  m_line = new QLineEdit(this);
  m_line->setEchoMode(QLineEdit::Password);
  pw_layout->addWidget(m_line);
  QDialogButtonBox *pw_buttons = new QDialogButtonBox(QDialogButtonBox::Ok | QDialogButtonBox::Cancel, this);
  pw_layout->addWidget(pw_buttons);
  m_err_lbl = new QLabel(this);
  m_err_lbl->setVisible(false);
  pw_layout->addWidget(m_err_lbl);
  connect(pw_buttons, &QDialogButtonBox::rejected, this, &QDialog::reject);
  connect(pw_buttons, &QDialogButtonBox::accepted, this, [this] {
    emit input_submitted(m_line->text().toUtf8());
    m_line->clear();
  });
}

QString auth_error_to_string(ResponseResult error)
{
  switch (error)
  {
  case ResponseResult::unsupported_version:
    return QStringLiteral("Unsupported key version.");
  case ResponseResult::decryption_failed:
    return QStringLiteral("Wrong passphrase (or corrupted key).");
  case ResponseResult::hardware_fault:
    return QStringLiteral("Hardware rejected this operation. It might be due to corrupted key, invalid server challenge, or a TPM fault.");
  case ResponseResult::invalid_pin:
    return QStringLiteral("Invalid PIN.");
  case ResponseResult::hardware_lockout:
    return QStringLiteral("Your hardware is locked out after too many failed PIN attempts. Try again later.");
  case ResponseResult::incompatible_arguments:
    return QStringLiteral("Server authentication incompatible with your key.");
  default:
    return QString("Error unlocking key (code %1). Catastrophic failure.").arg((int)error);
  }
}

void KeyActivationDialog::display_error(ResponseResult error) const
{
  m_err_lbl->setText(auth_error_to_string(error));
  m_err_lbl->setVisible(true);
}

void AuthFlow::on_key_selected(KeyringKey key)
{
  m_key = key;
  // Choosing the key for the first request.
  if (m_phase == AuthPhase::none)
  {
    AuthRequest req = {
      .username = m_username,
      .method = AuthMethod::certificate,
      .variant = m_key.type == KeyringStorageType::hardware ? CertVariant::hardware : CertVariant::software,
    };

    m_phase = AuthPhase::requested;

    emit request_ready(req);
  }
  // Already authenticating; choosing a different key for the same request.
  // This can be erroneous if a server challenge is issued for a different type
  // of key (software X25519 vs hardware P-256), creating a mismatch.
  else if (m_phase == AuthPhase::requested)
  {
    respond(m_challenge);
  }
}

void AuthFlow::authenticate(QByteArrayView input)
{
  AuthResponse response;
  ResponseResult result = unlock_and_auth(m_key, input, m_challenge.challenge, m_username.toUtf8(), response);
  if (result == ResponseResult::success)
  {
    emit response_ready(response);
    if (m_mode == FlowMode::Default)
    {
      m_saved_auth.insert(m_hostname.toUtf8(), m_username.toUtf8(), m_key.id);
    }
    reset();
  }
  // The error behavior depends on whether we need the application-level key
  // unlock dialog (for TSS2) or not (for Windows CNG, which uses its own).
  else if (m_input_dlg)
  {
    m_input_dlg->display_error(result);
  }
  else
  {
    QMessageBox::warning(nullptr, "Error", auth_error_to_string(result));
    reject_and_select_new_key();
  }
}

void AuthFlow::reject_and_select_new_key()
{
  if (m_mode == FlowMode::Saved)
  {
    m_saved_auth.remove(m_hostname.toUtf8(), m_username.toUtf8());
    m_mode = FlowMode::Default;
  }

  if (m_input_dlg)
  {
    m_input_dlg->deleteLater();
  }

  // The key selection dialog remains in the memory while the authentication
  // attempt is performed. It's only cleared on finalization of the entire flow.
  m_key_dlg->show();
}

void AuthFlow::call_input_dialog(KeyActivationDialog::Mode mode)
{
  m_input_dlg = new KeyActivationDialog(mode, m_key.name);
  m_input_dlg->setWindowModality(Qt::ApplicationModal);

  connect(m_input_dlg, &KeyActivationDialog::input_submitted, this, &AuthFlow::authenticate);
  connect(m_input_dlg, &QDialog::rejected, this, &AuthFlow::reject_and_select_new_key);
  m_key_dlg->hide();
  m_input_dlg->show();
}

bool AuthFlow::awaiting_challenge() const
{
  return m_phase == AuthPhase::requested;
}

void AuthFlow::respond(const AuthChallenge &challenge)
{
  if (m_phase != AuthPhase::requested)
  {
    return;
  }

  m_challenge = challenge;

  switch (m_key.type) {
  case KeyringStorageType::software:
    call_input_dialog(KeyActivationDialog::Mode::passphrase);
    break;
  case KeyringStorageType::hardware:
#ifdef SOF_AO_HARDWARE_KEY_REQUIRES_PIN_UI
    call_input_dialog(KeyActivationDialog::Mode::pin);
#else
    m_key_dlg->hide();
    authenticate(QByteArray());
#endif
    break;
  default:
    QMessageBox::warning(nullptr, "Error", QStringLiteral("Unable to use this key."));
    break;
  }
}

void AuthFlow::reset()
{
  if (m_key_dlg)
  {
    m_key_dlg->deleteLater();
    m_key_dlg = nullptr;
  }

  if (m_input_dlg)
  {
    m_input_dlg->deleteLater();
    m_input_dlg = nullptr;
  }

  m_phase = AuthPhase::none;
}