#include "imageprep.h"

#include <QBuffer>
#include <QDir>
#include <QFile>
#include <QImage>
#include <QImageReader>
#include <QImageWriter>

namespace {
// nginx is configured with client_max_body_size 20M; leave headroom.
constexpr qint64 kMaxUploadBytes = 18LL * 1024 * 1024;
constexpr int kJpegQuality = 92;
const int kQualityLadder[] = {92, 85, 75};

// Encodes to JPEG in memory at the given quality.
QByteArray encodeJpeg(const QImage &image, int quality)
{
    QByteArray out;
    QBuffer buffer(&out);
    buffer.open(QIODevice::WriteOnly);
    QImageWriter writer(&buffer, "jpeg");
    writer.setQuality(quality);
    writer.write(image);
    return out;
}

// Decodes with EXIF orientation baked into the pixels. Works for HEIC too, via
// the qtimageformats plugin.
QImage decodeOriented(const QString &path)
{
    QImageReader reader(path);
    reader.setAutoTransform(true);
    return reader.read();
}
} // namespace

ImagePrep::Prepared ImagePrep::prepare(const QString &sourcePath, const QString &cacheDir,
                                       const QString &baseName)
{
    Prepared result;

    QFile in(sourcePath);
    if (!in.open(QIODevice::ReadOnly)) {
        result.error = QStringLiteral("Could not read the selected image");
        return result;
    }
    const QByteArray source = in.readAll();
    in.close();

    QStringList notes;

    // A rotation flag cannot survive metadata stripping, so bake it into the
    // pixels instead, which costs a re-encode.
    QImageReader orientationReader(sourcePath);
    const bool mustRotate =
        orientationReader.transformation() != QImageIOHandler::TransformationNone;

    QByteArray bytes;
    QString extension;

    if (!mustRotate && isJpeg(source)) {
        bytes = stripJpegSegments(source);
        extension = QStringLiteral("jpg");
    } else if (!mustRotate && isPng(source)) {
        bytes = stripPngChunks(source);
        extension = QStringLiteral("png");
    }

    if (bytes.isEmpty()) {
        // Re-encode path: rotated originals, HEIC, or a JPEG/PNG whose marker
        // structure did not parse.
        const QImage image = decodeOriented(sourcePath);
        if (image.isNull()) {
            result.error = QStringLiteral("Could not decode the selected image");
            return result;
        }
        if (mustRotate)
            notes << QStringLiteral("rotated to match its EXIF orientation");
        else if (!isJpeg(source) && !isPng(source))
            notes << QStringLiteral("converted to JPEG");
        bytes = encodeJpeg(image, kJpegQuality);
        extension = QStringLiteral("jpg");
    }

    // Size valve. Only trips on genuinely huge originals; quality is stepped
    // before touching resolution.
    if (bytes.size() > kMaxUploadBytes) {
        QImage image = decodeOriented(sourcePath);
        if (image.isNull()) {
            result.error = QStringLiteral("Could not decode the selected image");
            return result;
        }
        QByteArray shrunk;
        for (int quality : kQualityLadder) {
            const QByteArray attempt = encodeJpeg(image, quality);
            if (attempt.size() <= kMaxUploadBytes) {
                shrunk = attempt;
                notes << QStringLiteral("re-encoded at quality %1 to fit the 20M upload limit").arg(quality);
                break;
            }
        }
        if (shrunk.isEmpty()) {
            double factor = 1.0;
            while (factor > 0.2) {
                factor *= 0.75;
                const QImage scaled = image.scaledToWidth(
                    qMax(1, int(image.width() * factor)), Qt::SmoothTransformation);
                const QByteArray attempt = encodeJpeg(scaled, kQualityLadder[2]);
                if (attempt.size() <= kMaxUploadBytes) {
                    shrunk = attempt;
                    notes << QStringLiteral("scaled to %1×%2 to fit the 20M upload limit")
                                 .arg(scaled.width()).arg(scaled.height());
                    break;
                }
            }
        }
        if (shrunk.isEmpty()) {
            result.error = QStringLiteral("Image is too large to upload even after downscaling");
            return result;
        }
        bytes = shrunk;
        extension = QStringLiteral("jpg");
    }

    const QString outPath = QDir(cacheDir).filePath(baseName + QLatin1Char('.') + extension);
    QFile out(outPath);
    if (!out.open(QIODevice::WriteOnly) || out.write(bytes) != bytes.size()) {
        result.error = QStringLiteral("Could not write the prepared image");
        return result;
    }
    out.close();

    result.filePath = outPath;
    result.note = notes.join(QStringLiteral("; "));
    return result;
}

bool ImagePrep::isJpeg(const QByteArray &b)
{
    return b.size() > 3 && quint8(b[0]) == 0xFF && quint8(b[1]) == 0xD8 && quint8(b[2]) == 0xFF;
}

bool ImagePrep::isPng(const QByteArray &b)
{
    return b.size() > 8 && quint8(b[0]) == 0x89 && quint8(b[1]) == 0x50 && quint8(b[2]) == 0x4E
        && quint8(b[3]) == 0x47;
}

// Drops APP1 (EXIF and XMP) and COM segments, keeping the compressed scan data
// byte-for-byte. APP0/JFIF and APP2/ICC stay: removing the colour profile would
// shift the rendered colours. Returns an empty array if the marker structure
// does not parse, so the caller falls back to a re-encode.
QByteArray ImagePrep::stripJpegSegments(const QByteArray &src)
{
    const auto u8 = [&](int i) { return quint8(src[i]); };
    QByteArray out;
    out.append(char(0xFF));
    out.append(char(0xD8));

    int i = 2;
    while (i + 1 < src.size()) {
        if (u8(i) != 0xFF)
            return {};
        const int marker = u8(i + 1);

        // Standalone markers: no length field, no payload.
        if (marker == 0x01 || (marker >= 0xD0 && marker <= 0xD7)) {
            out.append(char(0xFF));
            out.append(char(marker));
            i += 2;
            continue;
        }

        // Start of scan: entropy-coded data runs to the end of the file.
        if (marker == 0xDA) {
            out.append(src.mid(i));
            return out;
        }

        if (i + 3 >= src.size())
            return {};
        const int length = (u8(i + 2) << 8) | u8(i + 3);
        if (length < 2 || i + 2 + length > src.size())
            return {};

        const bool isExifOrXmp = marker == 0xE1;
        const bool isComment = marker == 0xFE;
        if (!isExifOrXmp && !isComment)
            out.append(src.mid(i, 2 + length));

        i += 2 + length;
    }
    return {};
}

// Drops the metadata chunks, keeps everything else including the critical
// IHDR/PLTE/IDAT/IEND run and colour chunks (iCCP, sRGB). Whole chunks go, so no
// CRC needs recomputing.
QByteArray ImagePrep::stripPngChunks(const QByteArray &src)
{
    static const QByteArray drop[] = {"eXIf", "tEXt", "zTXt", "iTXt", "tIME"};
    const auto u8 = [&](int i) { return quint8(src[i]); };
    if (src.size() < 8)
        return {};

    QByteArray out;
    out.append(src.mid(0, 8));

    int i = 8;
    while (i + 8 <= src.size()) {
        const qint64 length =
            (qint64(u8(i)) << 24) | (qint64(u8(i + 1)) << 16) | (qint64(u8(i + 2)) << 8) | qint64(u8(i + 3));
        if (length < 0)
            return {};

        const QByteArray type = src.mid(i + 4, 4);
        const qint64 total = 12 + length;
        if (i + total > src.size())
            return {};

        bool dropped = false;
        for (const QByteArray &d : drop) {
            if (type == d) {
                dropped = true;
                break;
            }
        }
        if (!dropped)
            out.append(src.mid(i, int(total)));
        i += int(total);

        if (type == "IEND")
            return out;
    }
    return {};
}