#include "ScreenshotEditor.h"
#include "UiHelpers.h"

#include <QApplication>
#include <QGuiApplication>
#include <QHBoxLayout>
#include <QKeyEvent>
#include <QKeySequence>
#include <QMouseEvent>
#include <QPainter>
#include <QPainterPath>
#include <QPaintEvent>
#include <QPushButton>
#include <QScreen>
#include <QToolButton>

namespace
{
    QPixmap grabDesktop()
    {
        const QList<QScreen *> screens = QGuiApplication::screens();
        if (screens.isEmpty())
            return {};
        QRect virt;
        for (QScreen *s : screens)
            virt = virt.united(s->geometry());
        if (!virt.isValid())
            return {};

        qreal maxDpr = 1.0;
        for (QScreen *s : screens)
            maxDpr = qMax(maxDpr, s->devicePixelRatio());

        QPixmap canvas(qRound(virt.width() * maxDpr), qRound(virt.height() * maxDpr));
        canvas.fill(Qt::black);
        canvas.setDevicePixelRatio(maxDpr);

        QPainter p(&canvas);
        for (QScreen *s : screens)
        {
            if (!s)
                continue;
            const QPixmap g = s->grabWindow(0);
            if (g.isNull())
                continue;
            const QPoint topLeft = s->geometry().topLeft() - virt.topLeft();
            p.drawPixmap(QRect(topLeft, s->geometry().size()), g);
        }
        return canvas;
    }

    QRect virtGeometry()
    {
        QRect virt;
        for (QScreen *s : QGuiApplication::screens())
            virt = virt.united(s->geometry());
        return virt;
    }
}

QPixmap ScreenshotEditor::capture(QWidget *host)
{
    Q_UNUSED(host);
    const QPixmap shot = grabDesktop();
    if (shot.isNull())
        return {};
    ScreenshotEditor dlg(shot, nullptr);
    dlg.setWindowModality(Qt::ApplicationModal);
    if (dlg.exec() != QDialog::Accepted)
        return {};
    return dlg.resultPixmap();
}

ScreenshotEditor::ScreenshotEditor(const QPixmap &shot, QWidget *parent)
    : QDialog(parent, Qt::FramelessWindowHint | Qt::WindowStaysOnTopHint | Qt::Tool)
    , m_shot(shot)
{
    setWindowModality(Qt::ApplicationModal);
    setMouseTracking(true);
    setCursor(Qt::CrossCursor);
    setFocusPolicy(Qt::StrongFocus);
    const QRect virt = virtGeometry();
    if (virt.isValid())
        setGeometry(virt);
    else
        showFullScreen();
    buildToolbar();
}

void ScreenshotEditor::buildToolbar()
{
    m_toolbar = new QWidget(this);
    m_toolbar->setObjectName(QStringLiteral("ShotToolbar"));
    m_toolbar->setStyleSheet(QStringLiteral(
        "QWidget#ShotToolbar { background:#FFFFFF; border:1px solid #E5E5E5; border-radius:6px; }"));
    auto *lay = new QHBoxLayout(m_toolbar);
    lay->setContentsMargins(8, 6, 8, 6);
    lay->setSpacing(6);

    const QColor colors[] = {
        QColor(QStringLiteral("#FA5151")),
        QColor(QStringLiteral("#FFC300")),
        QColor(QStringLiteral("#07C160")),
        QColor(QStringLiteral("#10AEFF")),
        QColor(QStringLiteral("#FFFFFF")),
        QColor(QStringLiteral("#191919")),
    };
    for (const QColor &c : colors)
    {
        auto *btn = new QToolButton(m_toolbar);
        btn->setFixedSize(22, 22);
        btn->setCursor(Qt::PointingHandCursor);
        btn->setAutoRaise(true);
        const QString border = (c == QColor(QStringLiteral("#FFFFFF")))
            ? QStringLiteral("1px solid #CCCCCC")
            : QStringLiteral("1px solid transparent");
        btn->setStyleSheet(QStringLiteral(
            "QToolButton { background:%1; border:%2; border-radius:11px; }"
            "QToolButton:hover { border:2px solid #07C160; }")
                               .arg(c.name(), border));
        connect(btn, &QToolButton::clicked, this, [this, c]() { m_color = c; });
        lay->addWidget(btn);
    }

    auto *undo = UiHelpers::ghostButton(QStringLiteral("撤销"), m_toolbar);
    undo->setFixedSize(56, 28);
    connect(undo, &QPushButton::clicked, this, [this]() {
        if (!m_strokes.isEmpty())
        {
            m_strokes.removeLast();
            update();
        }
    });
    auto *cancel = UiHelpers::ghostButton(QStringLiteral("取消"), m_toolbar);
    cancel->setFixedSize(56, 28);
    connect(cancel, &QPushButton::clicked, this, &QDialog::reject);
    auto *ok = UiHelpers::primaryButton(QStringLiteral("完成"), m_toolbar);
    ok->setFixedSize(56, 28);
    connect(ok, &QPushButton::clicked, this, [this]() { finish(); });

    lay->addWidget(undo);
    lay->addWidget(cancel);
    lay->addWidget(ok);
    m_toolbar->adjustSize();
    m_toolbar->hide();
}

void ScreenshotEditor::updateToolbar()
{
    if (!m_toolbar)
        return;
    if (!m_hasSel || m_sel.width() < 8 || m_sel.height() < 8)
    {
        m_toolbar->hide();
        return;
    }
    m_toolbar->adjustSize();
    const int tw = m_toolbar->width();
    const int th = m_toolbar->height();
    int x = m_sel.right() - tw;
    int y = m_sel.bottom() + 8;
    if (x < 8)
        x = 8;
    if (x + tw > width() - 8)
        x = qMax(8, width() - 8 - tw);
    if (y + th > height() - 8)
        y = m_sel.top() - 8 - th;
    if (y < 8)
        y = 8;
    m_toolbar->move(x, y);
    m_toolbar->show();
    m_toolbar->raise();
}

QPoint ScreenshotEditor::toLocal(const QPoint &pos) const
{
    return pos - m_sel.topLeft();
}

bool ScreenshotEditor::inSelection(const QPoint &pos) const
{
    return m_hasSel && m_sel.adjusted(0, 0, -1, -1).contains(pos);
}

void ScreenshotEditor::finish()
{
    if (!m_hasSel || m_sel.width() < 8 || m_sel.height() < 8)
        return;
    accept();
}

QPixmap ScreenshotEditor::resultPixmap() const
{
    if (!m_hasSel || m_sel.width() < 2 || m_sel.height() < 2)
        return {};
    const qreal dpr = m_shot.devicePixelRatio() > 0 ? m_shot.devicePixelRatio() : 1.0;
    const QRect src(qRound(m_sel.x() * dpr), qRound(m_sel.y() * dpr),
                    qRound(m_sel.width() * dpr), qRound(m_sel.height() * dpr));
    QPixmap out = m_shot.copy(src.intersected(m_shot.rect()));
    if (out.isNull())
        return {};
    out.setDevicePixelRatio(dpr);
    if (!m_strokes.isEmpty())
    {
        QPainter p(&out);
        p.setRenderHint(QPainter::Antialiasing, true);
        p.scale(dpr, dpr);
        for (const Stroke &s : m_strokes)
        {
            QPen pen(s.color, s.width, Qt::SolidLine, Qt::RoundCap, Qt::RoundJoin);
            p.setPen(pen);
            p.setBrush(Qt::NoBrush);
            p.drawPath(s.path);
        }
    }
    return out;
}

void ScreenshotEditor::paintEvent(QPaintEvent *)
{
    QPainter p(this);
    p.drawPixmap(rect(), m_shot);
    const bool liveSel = (m_selecting || m_hasSel) && m_sel.width() > 1 && m_sel.height() > 1;
    QPainterPath dim;
    dim.addRect(rect());
    if (liveSel)
    {
        QPainterPath hole;
        hole.addRect(m_sel);
        dim = dim.subtracted(hole);
        p.fillPath(dim, QColor(0, 0, 0, 120));
        p.setPen(QPen(QColor(QStringLiteral("#07C160")), 1));
        p.setBrush(Qt::NoBrush);
        p.drawRect(m_sel.adjusted(0, 0, -1, -1));

        if (m_hasSel && !m_selecting)
        {
            p.save();
            p.setRenderHint(QPainter::Antialiasing, true);
            p.setClipRect(m_sel);
            p.translate(m_sel.topLeft());
            auto drawStroke = [&p](const Stroke &s) {
                QPen pen(s.color, s.width, Qt::SolidLine, Qt::RoundCap, Qt::RoundJoin);
                p.setPen(pen);
                p.setBrush(Qt::NoBrush);
                p.drawPath(s.path);
            };
            for (const Stroke &s : m_strokes)
                drawStroke(s);
            if (m_drawing)
                drawStroke(m_cur);
            p.restore();
        }

        const QString t = QStringLiteral("%1 × %2").arg(m_sel.width()).arg(m_sel.height());
        const QFontMetrics fm(font());
        const int tw = fm.horizontalAdvance(t) + 10;
        const int th = fm.height() + 6;
        QRect lab(m_sel.left(), m_sel.top() - th - 4, tw, th);
        if (lab.top() < 0)
            lab.moveTop(m_sel.top() + 4);
        p.fillRect(lab, QColor(0, 0, 0, 160));
        p.setPen(Qt::white);
        p.drawText(lab, Qt::AlignCenter, t);
    }
    else
    {
        p.fillPath(dim, QColor(0, 0, 0, 80));
        if (!m_selecting)
        {
            p.setPen(QColor(255, 255, 255, 220));
            p.drawText(rect(), Qt::AlignCenter, QStringLiteral("拖动鼠标选择截图区域，Esc 取消"));
        }
    }
}

void ScreenshotEditor::mousePressEvent(QMouseEvent *event)
{
    if (event->button() != Qt::LeftButton)
        return;
    const QPoint pos = event->pos();
    if (m_hasSel && inSelection(pos))
    {
        m_drawing = true;
        m_cur.color = m_color;
        m_cur.width = 3.2;
        m_cur.path = QPainterPath();
        m_cur.path.moveTo(toLocal(pos));
        update();
        return;
    }
    if (m_hasSel && m_toolbar && m_toolbar->isVisible() && m_toolbar->geometry().contains(pos))
        return;
    m_hasSel = false;
    m_strokes.clear();
    m_selecting = true;
    m_dragOrigin = pos;
    m_sel = QRect(pos, QSize(1, 1));
    if (m_toolbar)
        m_toolbar->hide();
    setCursor(Qt::CrossCursor);
    grabMouse();
    update();
}

void ScreenshotEditor::mouseMoveEvent(QMouseEvent *event)
{
    const QPoint pos = event->pos();
    if (m_selecting)
    {
        m_sel = QRect(m_dragOrigin, pos).normalized();
        repaint();
        return;
    }
    if (m_drawing)
    {
        m_cur.path.lineTo(toLocal(pos));
        update();
        return;
    }
    if (m_hasSel && inSelection(pos))
        setCursor(Qt::CrossCursor);
    else if (m_hasSel)
        setCursor(Qt::ArrowCursor);
}

void ScreenshotEditor::mouseReleaseEvent(QMouseEvent *event)
{
    if (event->button() != Qt::LeftButton)
        return;
    if (mouseGrabber() == this)
        releaseMouse();
    if (m_selecting)
    {
        m_selecting = false;
        m_sel = QRect(m_dragOrigin, event->pos()).normalized();
        if (m_sel.width() < 8 || m_sel.height() < 8)
        {
            m_hasSel = false;
            m_sel = QRect();
            if (m_toolbar)
                m_toolbar->hide();
        }
        else
        {
            m_hasSel = true;
            updateToolbar();
        }
        update();
        return;
    }
    if (m_drawing)
    {
        m_drawing = false;
        if (!m_cur.path.isEmpty())
            m_strokes.push_back(m_cur);
        m_cur = Stroke();
        update();
    }
}

void ScreenshotEditor::keyPressEvent(QKeyEvent *event)
{
    if (event->key() == Qt::Key_Escape)
    {
        if (mouseGrabber() == this)
            releaseMouse();
        reject();
        return;
    }
    if (event->key() == Qt::Key_Return || event->key() == Qt::Key_Enter)
    {
        finish();
        return;
    }
    if (event->matches(QKeySequence::Undo) || (event->modifiers() & Qt::ControlModifier
                                               && event->key() == Qt::Key_Z))
    {
        if (!m_strokes.isEmpty())
        {
            m_strokes.removeLast();
            update();
        }
        return;
    }
    QDialog::keyPressEvent(event);
}