/******************************************************************************
 * The MIT License (MIT)
 *
 * Copyright (c) 2019-2021 Baldur Karlsson
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 ******************************************************************************/

#include "ShaderMessageViewer.h"
#include <QAction>
#include <QMenu>
#include <QMouseEvent>
#include <QScrollBar>
#include "Code/QRDUtils.h"
#include "Code/Resources.h"
#include "Widgets/Extended/RDHeaderView.h"
#include "Widgets/Extended/RDLineEdit.h"
#include "toolwindowmanager/ToolWindowManager.h"
#include "ui_ShaderMessageViewer.h"

ButtonDelegate::ButtonDelegate(const QIcon &icon, QWidget *parent)
    : m_Icon(icon), QStyledItemDelegate(parent)
{
}

void ButtonDelegate::paint(QPainter *painter, const QStyleOptionViewItem &option,
                           const QModelIndex &index) const
{
  // draw the background to get selection etc
  QApplication::style()->drawControl(QStyle::CE_ItemViewItem, &option, painter);

  QStyleOptionButton button;

  QSize sz = sizeHint(option, index);
  button.rect = option.rect;
  button.rect.setLeft(button.rect.center().x() - sz.width() / 2);
  button.rect.setTop(button.rect.center().y() - sz.height() / 2);
  button.rect.setSize(sz);
  button.icon = m_Icon;
  button.iconSize = sz;
  button.state = QStyle::State_Enabled;

  if(m_ClickedIndex == index)
    button.state |= QStyle::State_Sunken;

  QApplication::style()->drawControl(QStyle::CE_PushButton, &button, painter);
}

QSize ButtonDelegate::sizeHint(const QStyleOptionViewItem &option, const QModelIndex &index) const
{
  QStyleOptionButton button;
  button.icon = m_Icon;
  button.state = QStyle::State_Enabled;

  return QApplication::style()->sizeFromContents(QStyle::CT_PushButton, &button,
                                                 option.decorationSize);
}

bool ButtonDelegate::editorEvent(QEvent *event, QAbstractItemModel *model,
                                 const QStyleOptionViewItem &option, const QModelIndex &index)
{
  if(event->type() == QEvent::MouseButtonPress)
  {
    m_ClickedIndex = index;
  }
  else if(event->type() == QEvent::MouseMove)
  {
    QMouseEvent *e = (QMouseEvent *)event;

    if(m_ClickedIndex != index || (e->buttons() & Qt::LeftButton) == 0)
    {
      m_ClickedIndex = QModelIndex();
    }
    else
    {
      QPoint p = e->pos();

      QSize sz = option.decorationSize;
      QRect rect = option.rect;
      rect.setLeft(rect.center().x() - sz.width() / 2);
      rect.setTop(rect.center().y() - sz.height() / 2);
      rect.setSize(sz);

      if(!rect.contains(p))
      {
        m_ClickedIndex = QModelIndex();
      }
    }
  }
  else if(event->type() == QEvent::MouseButtonRelease)
  {
    if(m_ClickedIndex == index)
    {
      m_ClickedIndex = QModelIndex();

      QMouseEvent *e = (QMouseEvent *)event;

      QPoint p = e->pos();

      QSize sz = option.decorationSize;
      QRect rect = option.rect;
      rect.setLeft(rect.center().x() - sz.width() / 2);
      rect.setTop(rect.center().y() - sz.height() / 2);
      rect.setSize(sz);

      if(rect.contains(p))
        emit messageClicked(index);
    }
  }

  return true;
}

ShaderMessageViewer::ShaderMessageViewer(ICaptureContext &ctx, ShaderStageMask stages, QWidget *parent)
    : QFrame(parent), ui(new Ui::ShaderMessageViewer), m_Ctx(ctx)
{
  ui->setupUi(this);

  ui->messages->setFont(Formatter::PreferredFont());

  ui->messages->setMouseTracking(true);

  m_API = m_Ctx.APIProps().pipelineType;

  QObject::connect(ui->vertex, &QToolButton::toggled, [this](bool) { refreshMessages(); });
  QObject::connect(ui->hull, &QToolButton::toggled, [this](bool) { refreshMessages(); });
  QObject::connect(ui->domain, &QToolButton::toggled, [this](bool) { refreshMessages(); });
  QObject::connect(ui->geometry, &QToolButton::toggled, [this](bool) { refreshMessages(); });
  QObject::connect(ui->pixel, &QToolButton::toggled, [this](bool) { refreshMessages(); });
  QObject::connect(ui->filterButton, &QToolButton::clicked, [this]() { refreshMessages(); });
  QObject::connect(ui->filter, &RDLineEdit::returnPressed, [this]() { refreshMessages(); });

  ui->vertex->setText(ToQStr(ShaderStage::Vertex, m_API));
  ui->hull->setText(ToQStr(ShaderStage::Hull, m_API));
  ui->domain->setText(ToQStr(ShaderStage::Domain, m_API));
  ui->geometry->setText(ToQStr(ShaderStage::Geometry, m_API));
  ui->pixel->setText(ToQStr(ShaderStage::Pixel, m_API));

  m_EID = m_Ctx.CurEvent();
  m_Drawcall = m_Ctx.GetDrawcall(m_EID);

  const PipeState &pipe = m_Ctx.CurPipelineState();

  // check if we have multiview enabled
  m_Multiview = pipe.MultiviewBroadcastCount() > 1;

  // only display sample information if one of the targets is multisampled
  m_Multisampled = false;
  rdcarray<BoundResource> outs = pipe.GetOutputTargets();
  outs.push_back(pipe.GetDepthTarget());
  for(const BoundResource &o : outs)
  {
    if(o.resourceId == ResourceId())
      continue;

    const TextureDescription *tex = m_Ctx.GetTexture(o.resourceId);

    if(tex->msSamp > 1)
    {
      m_Multisampled = true;
      break;
    }
  }

  RDHeaderView *header = new RDHeaderView(Qt::Horizontal, this);
  ui->messages->setHeader(header);
  header->setStretchLastSection(true);
  header->setMinimumSectionSize(40);

  int sortColumn = 0;

  m_debugDelegate = new ButtonDelegate(Icons::wrench(), this);

  if(m_Drawcall && (m_Drawcall->flags & DrawFlags::Dispatch))
  {
    ui->stageFilters->hide();

    ui->messages->setColumns({lit("Debug"), tr("Workgroup"), lit("Thread"), lit("Message")});
    sortColumn = 1;

    ui->messages->setItemDelegateForColumn(0, m_debugDelegate);

    m_OrigShaders[5] = pipe.GetShader(ShaderStage::Compute);
  }
  else
  {
    ui->messages->setColumns({lit("Debug"), lit("Go to"), tr("Location"), lit("Message")});
    sortColumn = 2;

    m_gotoDelegate = new ButtonDelegate(Icons::find(), this);

    ui->messages->setItemDelegateForColumn(0, m_debugDelegate);
    ui->messages->setItemDelegateForColumn(1, m_gotoDelegate);

    QCheckBox *boxes[] = {
        ui->vertex, ui->hull, ui->domain, ui->geometry, ui->pixel,
    };

    for(ShaderStage s : values<ShaderStage>())
    {
      if(s == ShaderStage::Compute)
        continue;

      uint32_t idx = (uint32_t)s;

      m_OrigShaders[idx] = pipe.GetShader(s);

      boxes[idx]->setChecked(bool(stages & MaskForStage(s)));

      // if there's no shader bound, we currently don't support adding stages at runtime so just
      // hide this box as no messages can come from the unbound stage
      if(m_OrigShaders[idx] == ResourceId())
        boxes[idx]->hide();
    }
  }

  ui->messages->setContextMenuPolicy(Qt::CustomContextMenu);
  QObject::connect(ui->messages, &RDTreeWidget::customContextMenuRequested, [this](const QPoint &pos) {
    QModelIndex idx = ui->messages->indexAt(pos);
    RDTreeWidgetItem *item = ui->messages->itemForIndex(idx);

    QMenu contextMenu(this);

    QAction copy(tr("&Copy"), this);

    contextMenu.addAction(&copy);

    copy.setIcon(Icons::copy());

    QObject::connect(&copy, &QAction::triggered,
                     [this, pos, item]() { ui->messages->copyItem(pos, item); });

    QAction debugAction(tr("&Debug"), this);
    debugAction.setIcon(Icons::wrench());
    QAction gotoAction(tr("&Go to"), this);
    gotoAction.setIcon(Icons::find());

    QObject::connect(&debugAction, &QAction::triggered,
                     [this, idx]() { m_debugDelegate->messageClicked(idx); });

    QObject::connect(&gotoAction, &QAction::triggered,
                     [this, idx]() { m_gotoDelegate->messageClicked(idx); });

    contextMenu.addAction(&debugAction);
    if(m_gotoDelegate)
      contextMenu.addAction(&gotoAction);

    RDDialog::show(&contextMenu, ui->messages->viewport()->mapToGlobal(pos));
  });

  QObject::connect(m_debugDelegate, &ButtonDelegate::messageClicked, [this](const QModelIndex &idx) {
    RDTreeWidgetItem *item = ui->messages->itemForIndex(idx);

    int msgIdx = 0;
    if(item)
      msgIdx = item->tag().toInt();
    else
      return;

    ShaderMessage msg = m_Messages[msgIdx];

    const ShaderReflection *refl = m_Ctx.CurPipelineState().GetShaderReflection(msg.stage);

    if(refl->debugInfo.debuggable)
    {
      bool done = false;
      ShaderDebugTrace *trace = NULL;

      m_Ctx.Replay().AsyncInvoke([&trace, &done, msg](IReplayController *r) {
        if(msg.stage == ShaderStage::Compute)
          trace = r->DebugThread(msg.location.compute.workgroup, msg.location.compute.thread);
        else if(msg.stage == ShaderStage::Vertex)
          trace = r->DebugVertex(msg.location.vertex.vertexIndex, msg.location.vertex.instance,
                                 msg.location.vertex.vertexIndex, msg.location.vertex.view);
        else if(msg.stage == ShaderStage::Pixel)
          trace = r->DebugPixel(msg.location.pixel.x, msg.location.pixel.y,
                                msg.location.pixel.sample, msg.location.pixel.primitive);

        if(trace->debugger == NULL)
        {
          r->FreeTrace(trace);
          trace = NULL;
        }

        done = true;
      });

      QString debugContext;

      if(msg.stage == ShaderStage::Compute)
        debugContext = lit("Group [%1,%2,%3] Thread [%4,%5,%6]")
                           .arg(msg.location.compute.workgroup[0])
                           .arg(msg.location.compute.workgroup[1])
                           .arg(msg.location.compute.workgroup[2])
                           .arg(msg.location.compute.thread[0])
                           .arg(msg.location.compute.thread[1])
                           .arg(msg.location.compute.thread[2]);
      else if(msg.stage == ShaderStage::Vertex)
        debugContext = tr("Vertex %1").arg(msg.location.vertex.vertexIndex);
      else if(msg.stage == ShaderStage::Pixel)
        debugContext = tr("Pixel %1,%2").arg(msg.location.pixel.x).arg(msg.location.pixel.y);

      // wait a short while before displaying the progress dialog (which won't show if we're already
      // done by the time we reach it)
      for(int i = 0; !done && i < 100; i++)
        QThread::msleep(5);

      ShowProgressDialog(this, tr("Debugging %1").arg(debugContext), [&done]() { return done; });

      if(!trace)
      {
        RDDialog::critical(this, tr("Debug Error"), tr("Error debugging pixel."));
        return;
      }

      const ShaderBindpointMapping &bindMapping =
          m_Ctx.CurPipelineState().GetBindpointMapping(msg.stage);
      ResourceId pipeline = msg.stage == ShaderStage::Compute
                                ? m_Ctx.CurPipelineState().GetComputePipelineObject()
                                : m_Ctx.CurPipelineState().GetGraphicsPipelineObject();

      // viewer takes ownership of the trace
      IShaderViewer *s = m_Ctx.DebugShader(&bindMapping, refl, pipeline, trace, debugContext);

      if(msg.disassemblyLine >= 0)
      {
        s->ToggleBreakpointOnDisassemblyLine(msg.disassemblyLine);
        s->RunForward();
      }

      m_Ctx.AddDockWindow(s->Widget(), DockReference::AddTo, this);
    }
    else
    {
      RDDialog::critical(
          this, tr("Shader can't be debugged"),
          tr("The shader does not support debugging: %1").arg(refl->debugInfo.debugStatus));
    }
  });

  if(m_gotoDelegate)
  {
    QObject::connect(m_gotoDelegate, &ButtonDelegate::messageClicked, [this](const QModelIndex &idx) {
      RDTreeWidgetItem *item = ui->messages->itemForIndex(idx);

      int msgIdx = 0;
      if(item)
        msgIdx = item->tag().toInt();
      else
        return;

      ShaderMessage msg = m_Messages[msgIdx];

      m_Ctx.SetEventID({}, m_EID, m_EID);

      if(msg.stage == ShaderStage::Vertex)
      {
        m_Ctx.ShowMeshPreview();
        m_Ctx.GetMeshPreview()->SetCurrentInstance(msg.location.vertex.instance);
        m_Ctx.GetMeshPreview()->SetCurrentView(msg.location.vertex.view);
        // TODO, not accurate for indices
        m_Ctx.GetMeshPreview()->ScrollToRow(msg.location.vertex.vertexIndex, MeshDataStage::VSIn);
      }
      else if(msg.stage == ShaderStage::Pixel)
      {
        m_Ctx.ShowTextureViewer();
        Subresource sub = m_Ctx.GetTextureViewer()->GetSelectedSubresource();
        sub.sample = msg.location.pixel.sample;
        m_Ctx.GetTextureViewer()->SetSelectedSubresource(sub);

        // select an actual output. Prefer the first colour output, but if there's no colour output
        // pick depth.
        rdcarray<BoundResource> cols = m_Ctx.CurPipelineState().GetOutputTargets();
        bool hascol = false;
        for(size_t i = 0; i < cols.size(); i++)
          hascol |= cols[i].resourceId != ResourceId();

        if(hascol)
          m_Ctx.GetTextureViewer()->ViewFollowedResource(FollowType::OutputColor,
                                                         ShaderStage::Pixel, 0, 0);
        else
          m_Ctx.GetTextureViewer()->ViewFollowedResource(FollowType::OutputDepth,
                                                         ShaderStage::Pixel, 0, 0);
        m_Ctx.GetTextureViewer()->GotoLocation(msg.location.pixel.x, msg.location.pixel.y);
      }
      else
      {
        qCritical() << "Can't go to a compute thread";
      }
    });
  }

  // deliberately copy m_OrigShaders to m_ReplacedShaders. This is impossible because we should
  // either see a ResourceId() for unedited, or a new resource for edited. This means when we first
  // get OnEventChanged() called we will definitely detect the situation as 'stale' and refresh the
  // messages.
  memcpy(m_ReplacedShaders, m_OrigShaders, sizeof(m_ReplacedShaders));

  header->setDefaultAlignment(Qt::AlignLeft | Qt::AlignVCenter);

  ui->staleStatus->hide();

  ui->label->setText(tr("Shader messages from @%1 - %2")
                         .arg(m_EID)
                         .arg(m_Drawcall ? m_Drawcall->name : rdcstr("Unknown draw")));

  setWindowTitle(tr("Shader messages at @%1").arg(m_EID));

  m_Ctx.AddCaptureViewer(this);

  OnEventChanged(m_Ctx.CurEvent());

  ui->messages->sortByColumn(sortColumn, Qt::SortOrder::AscendingOrder);

  for(int i = 0; i < 4; i++)
  {
    header->setSectionResizeMode(i, QHeaderView::Interactive);
    ui->messages->resizeColumnToContents(i);
  }
}

ShaderMessageViewer::~ShaderMessageViewer()
{
  m_Ctx.RemoveCaptureViewer(this);
  delete ui;
}

bool ShaderMessageViewer::IsOutOfDate()
{
  return ui->staleStatus->isVisible();
}

void ShaderMessageViewer::OnCaptureLoaded()
{
}

void ShaderMessageViewer::OnCaptureClosed()
{
  ToolWindowManager::closeToolWindow(this);
}

void ShaderMessageViewer::OnEventChanged(uint32_t eventId)
{
  ResourceId shaders[6];
  bool editsChanged = false;
  QString staleReason;

  for(ShaderStage s : values<ShaderStage>())
  {
    uint32_t idx = (uint32_t)s;
    shaders[idx] = m_Ctx.GetResourceReplacement(m_OrigShaders[idx]);

    // either an edit has been applied, updated, or removed if these don't match
    if(shaders[idx] != m_ReplacedShaders[idx])
    {
      editsChanged = true;
      staleReason += QFormatStr(", %1").arg(ToQStr(s, m_API));
    }
  }

  // if the edits haven't changed, just skip
  if(!editsChanged)
    return;

  // if it's the current event we can update with the latest
  if(m_EID == eventId)
  {
    m_Messages = m_Ctx.CurPipelineState().GetShaderMessages();

    // not stale anymore
    ui->staleStatus->hide();

    // update current set of replaced shaders
    memcpy(m_ReplacedShaders, shaders, sizeof(m_ReplacedShaders));

    refreshMessages();
  }
  else
  {
    staleReason.remove(0, 2);

    // otherwise we can't - just update the stale status
    ui->staleStatus->show();
    ui->staleStatus->setText(
        tr("Messages are stale because edits to %1 shaders have changed since they were fetched.\n"
           "Select the event @%2 to refresh.")
            .arg(staleReason)
            .arg(m_EID));

    ui->messages->beginUpdate();

    for(int i = 0; i < ui->messages->topLevelItemCount(); i++)
      ui->messages->topLevelItem(i)->setItalic(true);

    ui->messages->endUpdate();
  }
}

void ShaderMessageViewer::refreshMessages()
{
  ShaderStageMask mask = ShaderStageMask::Compute;

  if(!m_Drawcall || !(m_Drawcall->flags & DrawFlags::Dispatch))
  {
    mask = ShaderStageMask::Unknown;

    if(ui->vertex->isChecked())
      mask |= ShaderStageMask::Vertex;
    if(ui->hull->isChecked())
      mask |= ShaderStageMask::Hull;
    if(ui->domain->isChecked())
      mask |= ShaderStageMask::Domain;
    if(ui->geometry->isChecked())
      mask |= ShaderStageMask::Geometry;
    if(ui->pixel->isChecked())
      mask |= ShaderStageMask::Pixel;
  }

  int vs = ui->messages->verticalScrollBar()->value();
  int curMsg = -1;
  {
    RDTreeWidgetItem *item = ui->messages->currentItem();
    if(item)
      curMsg = item->tag().toInt();
  }
  RDTreeWidgetItem *newCurrentItem = NULL;
  ui->messages->beginUpdate();
  ui->messages->clear();

  QString filter = ui->filter->text().trimmed();

  for(int i = 0; i < m_Messages.count(); i++)
  {
    const ShaderMessage &msg = m_Messages[i];

    // filter by stages
    if(!(MaskForStage(msg.stage) & mask))
      continue;

    QString text(msg.message);

    QString location;
    if(msg.stage == ShaderStage::Vertex)
    {
      // only show the view if the draw has multiview enabled
      if(m_Multiview)
      {
        location += lit("View %1, ").arg(msg.location.vertex.view);
      }

      // only show the instance if the draw is actually instanced
      if(m_Drawcall && (m_Drawcall->flags & DrawFlags::Instanced) && m_Drawcall->numInstances > 1)
      {
        location += lit("Inst %1, ").arg(msg.location.vertex.instance);
      }

      if(m_Drawcall && (m_Drawcall->flags & DrawFlags::Indexed))
      {
        location += lit("Idx %1").arg(msg.location.vertex.vertexIndex);
      }
      else
      {
        location += lit("Vert %1").arg(msg.location.vertex.vertexIndex);
      }
    }
    else if(msg.stage == ShaderStage::Pixel)
    {
      location = QFormatStr("%1 %2,%3")
                     .arg(IsD3D(m_API) ? lit("Pixel") : lit("Frag"))
                     .arg(msg.location.pixel.x)
                     .arg(msg.location.pixel.y);

      if(msg.location.pixel.primitive == ~0U)
        location += lit(", Prim ?");
      else
        location += lit(", Prim %1").arg(msg.location.pixel.primitive);

      if(m_Multisampled)
      {
        if(msg.location.pixel.sample == ~0U)
          location += lit(", Samp ?");
        else
          location += lit(", Samp %1").arg(msg.location.pixel.sample);
      }
    }
    else if(msg.stage == ShaderStage::Compute)
    {
    }
    else
    {
      // no location info for other stages
      location = tr("Unknown shader");
    }

    // filter by text on location and messag
    if(!filter.isEmpty() && !text.contains(filter, Qt::CaseInsensitive) &&
       !location.contains(filter, Qt::CaseInsensitive))
      continue;

    RDTreeWidgetItem *node = NULL;

    if(msg.stage == ShaderStage::Compute)
    {
      node = new RDTreeWidgetItem({
          QString(), QFormatStr("%1, %2, %3")
                         .arg(msg.location.compute.workgroup[0])
                         .arg(msg.location.compute.workgroup[1])
                         .arg(msg.location.compute.workgroup[2]),
          QFormatStr("%1, %2, %3")
              .arg(msg.location.compute.thread[0])
              .arg(msg.location.compute.thread[1])
              .arg(msg.location.compute.thread[2]),
          text,
      });
    }
    else
    {
      node = new RDTreeWidgetItem({QString(), QString(), location, text});
    }

    if(node)
    {
      if(i == curMsg)
        newCurrentItem = node;

      node->setItalic(ui->staleStatus->isVisible());
      node->setTag(i);
      ui->messages->addTopLevelItem(node);
    }
  }

  ui->messages->clearSelection();
  ui->messages->endUpdate();
  ui->messages->verticalScrollBar()->setValue(vs);

  if(newCurrentItem)
  {
    ui->messages->setCurrentItem(newCurrentItem);
    ui->messages->scrollToItem(newCurrentItem);
  }
}
