// SPDX-FileCopyrightText: 2008 by Jakub Stachowski <qbast@go2.pl>

// SPDX-License-Identifier: GPL-2.0-or-later

#include "mobipocket.h"

#include "decompressor.h"

#include "pdb_p.h"

#include "qmobipocket_debug.h"

#include <QBuffer>

#include <QIODevice>

#include <QImageReader>

#include <QRegularExpression>

#include <QStringConverter>

#include <QtEndian>

namespace Mobipocket

{

struct DocumentPrivate {

    DocumentPrivate(QIODevice *d)

        : pdb(d)

    {

    }

    PDB pdb;

    std::unique_ptr<Decompressor> dec;

    quint16 ntextrecords = 0;

    quint16 maxRecordSize = 0;

    bool valid = false;

    // number of first record holding image. Usually it is directly after end of text, but not always

    quint16 firstImageRecord = 0;

    QMap<Document::MetaKey, QString> metadata;

    QStringDecoder toUtf16;

    bool drm = false;

    quint32 extraflags = 0;

    // index of Thumbnail image in image list. May be specified in EXTH.

    int thumbnailIndex = -1;

    // index of Cover image in image list. May be specified in EXTH.

    int coverIndex = -1;

    void init();

    void findFirstImage();

    void parseEXTH(QByteArrayView data);

    void parseHtmlHead(const QString &data);

};

void DocumentPrivate::parseHtmlHead(const QString &data)

{

    static const QRegularExpression title(QLatin1String("<dc:title.*>(.*)</dc:title>"),

                                          QRegularExpression::CaseInsensitiveOption | QRegularExpression::InvertedGreedinessOption);

    static const QRegularExpression author(QLatin1String("<dc:creator.*>(.*)</dc:creator>"),

                                           QRegularExpression::CaseInsensitiveOption | QRegularExpression::InvertedGreedinessOption);

    static const QRegularExpression copyright(QLatin1String("<dc:rights.*>(.*)</dc:rights>"),

                                              QRegularExpression::CaseInsensitiveOption | QRegularExpression::InvertedGreedinessOption);

    static const QRegularExpression subject(QLatin1String("<dc:subject.*>(.*)</dc:subject>"),

                                            QRegularExpression::CaseInsensitiveOption | QRegularExpression::InvertedGreedinessOption);

    static const QRegularExpression description(QLatin1String("<dc:description.*>(.*)</dc:description>"),

                                                QRegularExpression::CaseInsensitiveOption | QRegularExpression::InvertedGreedinessOption);

    // title could have been already taken from MOBI record

    if (!metadata.contains(Document::Title)) {

        if (const auto titleMatch = title.match(data); titleMatch.hasMatch())

            metadata[Document::Title] = titleMatch.captured(1);

    }

    if (const auto authorMatch = author.match(data); authorMatch.hasMatch())

        metadata[Document::Author] = authorMatch.captured(1);

    if (const auto copyrightMatch = copyright.match(data); copyrightMatch.hasMatch())

        metadata[Document::Copyright] = copyrightMatch.captured(1);

    if (const auto subjectMatch = subject.match(data); subjectMatch.hasMatch())

        metadata[Document::Subject] = subjectMatch.captured(1);

    if (const auto descriptionMatch = description.match(data); descriptionMatch.hasMatch())

        metadata[Document::Description] = descriptionMatch.captured(1);

}

namespace

{

    const QVector<QByteArray> getHuffRecords(const PDB &pdb)

    {

        const QByteArray header = pdb.getRecord(0);

        if (header.size() < 0x78 || header[1] != 'H') {

            return {};

        }

        const quint32 huff_ofs = qFromBigEndian<quint32>(header.constData() + 0x70);

        const quint32 huff_num = qFromBigEndian<quint32>(header.constData() + 0x74);

        // Check for overflow and out-of-bounds access

        if (((huff_ofs + huff_num) < huff_num) || ((huff_ofs + huff_num) > pdb.recordCount())) {

            return {};

        }

        QVector<QByteArray> records(huff_num);

        for (quint32 i = 0; i < huff_num; i++) {

            if (auto r = pdb.getRecord(huff_ofs + i); r.isNull()) {

                return {};

            } else {

                records[i] = r;

            }

        }

        return records;

    };

}

void DocumentPrivate::init()

{

    quint32 encoding = 0;

    if (!pdb.isValid())

        return;

    const QByteArray mhead = pdb.getRecord(0);

    if (mhead.isNull() || mhead.size() < 14)

        return;

    dec = Decompressor::create(mhead[1], getHuffRecords(pdb));

    if ((int)mhead[12] != 0 || (int)mhead[13] != 0)

        drm = true;

    if (!dec || !dec->isValid())

        return;

    ntextrecords = qFromBigEndian<quint16>(mhead.constData() + 8);

    maxRecordSize = qFromBigEndian<quint16>(mhead.constData() + 10);

    if (mhead.size() > 31)

        encoding = qFromBigEndian<quint32>(mhead.constData() + 28);

    if (encoding == 65001) {

        toUtf16 = QStringDecoder(QStringDecoder::Utf8);

    } else {

        toUtf16 = QStringDecoder("windows-1252");

        if (!toUtf16.isValid()) {

            qCWarning(QMOBIPOCKET_LOG) << "Text codec \"windows-1252\" not supported by Qt library, falling back to Latin1";

            toUtf16 = QStringDecoder(QStringConverter::Latin1);

        }

    }

    parseEXTH(mhead);

    if (mhead.size() >= 244) {

        quint32 exthoffs = qFromBigEndian<quint32>(mhead.constData() + 20);

        if ((exthoffs + 16) > 244) {

            extraflags = qFromBigEndian<quint32>(mhead.constData() + 240);

        }

    }

    // try getting metadata from HTML if nothing or only title was recovered from MOBI and EXTH records

    if (metadata.size() < 2 && !drm)

        parseHtmlHead(toUtf16(dec->decompress(pdb.getRecord(1))));

    valid = true;

}

void DocumentPrivate::findFirstImage()

{

    firstImageRecord = ntextrecords + 1;

    while (firstImageRecord < pdb.recordCount()) {

        QByteArray rec = pdb.getRecord(firstImageRecord);

        if (rec.isNull())

            return;

        QBuffer buf(&rec);

        buf.open(QIODevice::ReadOnly);

        QImageReader r(&buf);

        if (r.canRead())

            return;

        firstImageRecord++;

    }

}

void DocumentPrivate::parseEXTH(QByteArrayView data)

{

    if (data.size() < 92)

        return;

    // try to get name

    const qint32 nameoffset = qFromBigEndian<quint32>(data.constData() + 84);

    const qint32 namelen = qFromBigEndian<quint32>(data.constData() + 88);

    const qint32 ssize = qint32(data.size());

    if (nameoffset >= 0 && namelen >= 0 && nameoffset <= ssize && namelen <= ssize - nameoffset) {

        metadata[Document::Title] = toUtf16(data.mid(nameoffset, namelen));

    }

    const quint32 size = quint32(data.size());

    const quint32 exthoffs = qFromBigEndian<quint32>(data.constData() + 20);

    if (exthoffs > size - 28) {

        return;

    }

    if (data.mid(exthoffs + 16, 4) != "EXTH")

        return;

    quint32 records = qFromBigEndian<quint32>(data.constData() + exthoffs + 24);

    quint32 offset = exthoffs + 28;

    for (unsigned int i = 0; i < records; i++) {

        if (offset > size - 8)

            break;

        quint32 type = qFromBigEndian<quint32>(data.constData() + offset);

        quint32 len = qFromBigEndian<quint32>(data.constData() + offset + 4);

        if (len < 8 || len > size - offset)

            break;

        switch (type) {

        case 100:

            metadata[Document::Author] = toUtf16(data.mid(offset + 8, len - 8));

            break;

        case 103:

            metadata[Document::Description] = toUtf16(data.mid(offset + 8, len - 8));

            break;

        case 105:

            metadata[Document::Subject] = toUtf16(data.mid(offset + 8, len - 8));

            break;

        case 109:

            metadata[Document::Copyright] = toUtf16(data.mid(offset + 8, len - 8));

            break;

        case 201:

            if (len >= 12)

                coverIndex = qFromBigEndian<quint32>(data.constData() + offset + 8);

            break;

        case 202:

            if (len >= 12)

                thumbnailIndex = qFromBigEndian<quint32>(data.constData() + offset + 8);

            break;

        default:

            // ignore

            break;

        }

        offset += len;

    }

}

Document::Document(QIODevice *dev)

    : d(new DocumentPrivate(dev))

{

    Q_ASSERT(dev->openMode() & QIODevice::ReadOnly);

    Q_ASSERT(!dev->isSequential());

    d->init();

}

Document::~Document()

{

    delete d;

}

namespace

{

constexpr qsizetype preTrailingDataLength(QByteArrayView data, quint32 flags)

{

    if (flags == 0) {

        return data.size();

    }

    for (int i = 31; i > 0; i--) {

        if ((flags & (1u << i)) == 0) {

            continue;

        }

        qsizetype chopN = 0;

        for (int j = 0; j < 4; j++) {

            if (j + 1 > data.size()) {

                return 0;

            }

            quint8 l = data.at(data.size() - (j + 1));

            chopN |= (l & 0x7f) << (7 * j);

            if (l & 0x80) {

                break;

            }

        }

        data.chop(std::min<qsizetype>(chopN, data.size()));

    }

    if ((flags & 0x1) && !data.isEmpty()) {

        quint8 l = data.back() & 0x3;

        data.chop(std::min<qsizetype>(l + 1, data.size()));

    }

    return data.size();

}

static_assert(preTrailingDataLength({"0\x00", 2}, 0x0) == 2);

static_assert(preTrailingDataLength({"0\x00", 2}, 0x1) == 1);

static_assert(preTrailingDataLength({"0\x01", 2}, 0x1) == 0);

static_assert(preTrailingDataLength({"0\x02", 2}, 0x1) == 0);

static_assert(preTrailingDataLength({"abcd\x03", 5}, 0x1) == 1);

static_assert(preTrailingDataLength({"abcd\x81", 5}, 0x2) == 4);

static_assert(preTrailingDataLength({"\x02\x01", 2}, 0x2) == 0);

static_assert(preTrailingDataLength({"\x80\x02", 2}, 0x2) == 0);

static_assert(preTrailingDataLength({"abcd\x85", 5}, 0x2) == 0);

static_assert(preTrailingDataLength({"abc\x01\x7f\x82", 6}, 0x2) == 4);

static_assert(preTrailingDataLength({"abc\x01\x80\x02", 6}, 0x2) == 4);

static_assert(preTrailingDataLength({"abc\x01\x7f\x82", 6}, 0x3) == 2);

static_assert(preTrailingDataLength({"abc\x81\x80\x02", 6}, 0x6) == 3);

static_assert(preTrailingDataLength({"abc\x00\x81\x81", 6}, 0x7) == 3);

} // namespace

QString Document::text(int size) const

{

    QByteArray whole;

    for (int i = 1; i < d->ntextrecords + 1; i++) {

        auto record = d->pdb.getRecord(i);

        record.resize(preTrailingDataLength(record, d->extraflags));

        QByteArray decompressedRecord = d->dec->decompress(record);

        whole += decompressedRecord;

        if (!d->dec->isValid()) {

            d->valid = false;

            return QString();

        }

        if (size != -1 && whole.size() > size)

            break;

    }

    return d->toUtf16(whole);

}

int Document::imageCount() const

{

    // FIXME: don't count FLIS and FCIS records

    return d->pdb.recordCount() - d->ntextrecords;

}

bool Document::isValid() const

{

    return d->valid;

}

QImage Document::getImage(int i) const

{

    if (!d->firstImageRecord)

        d->findFirstImage();

    if ((i < 0) || (i > std::numeric_limits<quint16>::max()) //

        || (d->firstImageRecord + i) >= d->pdb.recordCount()) {

        return {};

    }

    QByteArray rec = d->pdb.getRecord(d->firstImageRecord + i);

    return (rec.isNull()) ? QImage() : QImage::fromData(rec);

}

QMap<Document::MetaKey, QString> Document::metadata() const

{

    return d->metadata;

}

bool Document::hasDRM() const

{

    return d->drm;

}

QImage Document::thumbnail() const

{

    if (QImage img = getImage(d->thumbnailIndex); !img.isNull()) {

        return img;

    }

    // Fall back to cover image, or return an empty image

    return getImage(d->coverIndex);

}

}