/*
 * Copyright 2011, Oliver Tappe <zooey@hirschkaefer.de>
 * Distributed under the terms of the MIT License.
 */


#include <package/hpkg/RepositoryWriterImpl.h>

#include <algorithm>
#include <new>

#include <ByteOrder.h>
#include <Message.h>
#include <Path.h>

#include <AutoDeleter.h>
#include <HashSet.h>

#include <package/hpkg/BlockBufferPoolNoLock.h>
#include <package/hpkg/HPKGDefsPrivate.h>
#include <package/hpkg/PackageDataReader.h>
#include <package/hpkg/PackageEntry.h>
#include <package/hpkg/PackageFileHeapWriter.h>
#include <package/hpkg/PackageInfoAttributeValue.h>
#include <package/hpkg/PackageReader.h>
#include <package/ChecksumAccessors.h>
#include <package/PackageInfoContentHandler.h>
#include <package/RepositoryInfo.h>

#include "HashableString.h"


namespace BPackageKit {

namespace BHPKG {

namespace BPrivate {


using BPackageKit::BPrivate::GeneralFileChecksumAccessor;
using BPackageKit::BPrivate::HashableString;


namespace {


// #pragma mark - PackageEntryDataFetcher


struct PackageEntryDataFetcher {
        PackageEntryDataFetcher(BErrorOutput* errorOutput,
                BPackageData& packageData)
                :
                fErrorOutput(errorOutput),
                fPackageData(packageData)
        {
        }

        status_t ReadIntoString(BAbstractBufferedDataReader* heapReader,
                BString& _contents)
        {
                // create a PackageDataReader
                BAbstractBufferedDataReader* reader;
                status_t result = BPackageDataReaderFactory()
                        .CreatePackageDataReader(heapReader, fPackageData, reader);
                if (result != B_OK)
                        return result;
                ObjectDeleter<BAbstractBufferedDataReader> readerDeleter(reader);

                // copy data into the given string
                int32 bufferSize = fPackageData.Size();
                char* buffer = _contents.LockBuffer(bufferSize);
                if (buffer == NULL)
                        return B_NO_MEMORY;

                result = reader->ReadData(0, buffer, bufferSize);
                if (result != B_OK) {
                        fErrorOutput->PrintError("Error: Failed to read data: %s\n",
                                strerror(result));
                        _contents.UnlockBuffer(0);
                } else
                        _contents.UnlockBuffer(bufferSize);

                return result;
        }

private:
        BErrorOutput*                   fErrorOutput;
        BPackageData&                   fPackageData;
};


// #pragma mark - PackageContentHandler


struct PackageContentHandler : public BPackageInfoContentHandler {
        PackageContentHandler(BErrorOutput* errorOutput, BPackageInfo* packageInfo,
                BAbstractBufferedDataReader* heapReader,
                BRepositoryInfo* repositoryInfo)
                :
                BPackageInfoContentHandler(*packageInfo, errorOutput),
                fHeapReader(heapReader),
                fRepositoryInfo(repositoryInfo)
        {
        }

        virtual status_t HandleEntry(BPackageEntry* entry)
        {
                // if license must be approved, read any license files from package such
                // that those can be stored in the repository later
                if ((fPackageInfo.Flags() & B_PACKAGE_FLAG_APPROVE_LICENSE) == 0
                        || entry == NULL)
                        return B_OK;

                // return if not in ./data/licenses folder
                const BPackageEntry* parent = entry->Parent();
                BString licenseFolderName("licenses");
                if (parent == NULL || licenseFolderName != parent->Name())
                        return B_OK;

                parent = parent->Parent();
                BString dataFolderName("data");
                if (parent == NULL || dataFolderName != parent->Name())
                        return B_OK;

                if (parent->Parent() != NULL)
                        return B_OK;

                // check if license already is in repository
                const BStringList& licenseNames = fRepositoryInfo->LicenseNames();
                for (int i = 0; i < licenseNames.CountStrings(); ++i) {
                        if (licenseNames.StringAt(i).ICompare(entry->Name()) == 0) {
                                // license already exists
                                return B_OK;
                        }
                }

                // fetch contents of license file
                BPackageData& packageData = entry->Data();
                PackageEntryDataFetcher dataFetcher(fErrorOutput, packageData);

                BString licenseText;
                status_t result = dataFetcher.ReadIntoString(fHeapReader, licenseText);
                if (result != B_OK)
                        return result;

                // add license to repository
                return fRepositoryInfo->AddLicense(entry->Name(), licenseText);
        }

        virtual status_t HandleEntryAttribute(BPackageEntry* entry,
                BPackageEntryAttribute* attribute)
        {
                return B_OK;
        }

        virtual status_t HandleEntryDone(BPackageEntry* entry)
        {
                return B_OK;
        }

        virtual void HandleErrorOccurred()
        {
        }

private:
        BPackageReader*                                 fPackageReader;
        BAbstractBufferedDataReader*    fHeapReader;
        BRepositoryInfo*                                fRepositoryInfo;
};


}       // anonymous namespace


// #pragma mark - PackageNameSet


struct RepositoryWriterImpl::PackageNameSet
        : public ::BPrivate::HashSet<HashableString> {
};


// #pragma mark - RepositoryWriterImpl


RepositoryWriterImpl::RepositoryWriterImpl(BRepositoryWriterListener* listener,
        BRepositoryInfo* repositoryInfo)
        :
        inherited("repository", listener),
        fListener(listener),
        fRepositoryInfo(repositoryInfo),
        fPackageCount(0),
        fPackageNames(NULL)
{
}


RepositoryWriterImpl::~RepositoryWriterImpl()
{
        delete fPackageNames;
}


status_t
RepositoryWriterImpl::Init(const char* fileName)
{
        try {
                fPackageNames = new PackageNameSet();
                status_t result = fPackageNames->InitCheck();
                if (result != B_OK)
                        return result;
                return _Init(fileName);
        } catch (status_t error) {
                return error;
        } catch (std::bad_alloc&) {
                fListener->PrintError("Out of memory!\n");
                return B_NO_MEMORY;
        }
}


status_t
RepositoryWriterImpl::AddPackage(const BEntry& packageEntry)
{
        try {
                return _AddPackage(packageEntry);
        } catch (status_t error) {
                return error;
        } catch (std::bad_alloc&) {
                fListener->PrintError("Out of memory!\n");
                return B_NO_MEMORY;
        }
}


status_t
RepositoryWriterImpl::AddPackageInfo(const BPackageInfo& packageInfo)
{
        try {
                return _AddPackageInfo(packageInfo);
        } catch (status_t error) {
                return error;
        } catch (std::bad_alloc&) {
                fListener->PrintError("Out of memory!\n");
                return B_NO_MEMORY;
        }
}


status_t
RepositoryWriterImpl::Finish()
{
        try {
                return _Finish();
        } catch (status_t error) {
                return error;
        } catch (std::bad_alloc&) {
                fListener->PrintError("Out of memory!\n");
                return B_NO_MEMORY;
        }
}


status_t
RepositoryWriterImpl::_Init(const char* fileName)
{
        status_t error = inherited::Init(NULL, false, fileName,
                BPackageWriterParameters());
        if (error != B_OK)
                return error;

        return InitHeapReader(sizeof(hpkg_repo_header));
}


status_t
RepositoryWriterImpl::_Finish()
{
        hpkg_repo_header header;

        // write repository info
        uint64 infoLength;
        status_t result = _WriteRepositoryInfo(header, infoLength);
        if (result != B_OK)
                return result;

        // write package attributes
        uint64 packagesLength;
        _WritePackageAttributes(header, packagesLength);

        // flush the heap writer
        result = fHeapWriter->Finish();
        if (result != B_OK)
                return result;
        uint64 compressedHeapSize = fHeapWriter->CompressedHeapSize();
        uint64 totalSize = fHeapWriter->HeapOffset() + compressedHeapSize;

        header.heap_compression = B_HOST_TO_BENDIAN_INT16(
                Parameters().Compression());
        header.heap_chunk_size = B_HOST_TO_BENDIAN_INT32(fHeapWriter->ChunkSize());
        header.heap_size_compressed = B_HOST_TO_BENDIAN_INT64(compressedHeapSize);
        header.heap_size_uncompressed = B_HOST_TO_BENDIAN_INT64(
                fHeapWriter->UncompressedHeapSize());

        fListener->OnRepositoryDone(sizeof(header), infoLength,
                fRepositoryInfo->LicenseNames().CountStrings(), fPackageCount,
                packagesLength, totalSize);

        // update the general header info and write the header
        header.magic = B_HOST_TO_BENDIAN_INT32(B_HPKG_REPO_MAGIC);
        header.header_size = B_HOST_TO_BENDIAN_INT16((uint16)sizeof(header));
        header.version = B_HOST_TO_BENDIAN_INT16(B_HPKG_REPO_VERSION);
        header.total_size = B_HOST_TO_BENDIAN_INT64(totalSize);
        header.minor_version = B_HOST_TO_BENDIAN_INT16(B_HPKG_REPO_MINOR_VERSION);

        RawWriteBuffer(&header, sizeof(header), 0);

        SetFinished(true);
        return B_OK;
}


status_t
RepositoryWriterImpl::_AddPackage(const BEntry& packageEntry)
{
        status_t result = packageEntry.InitCheck();
        if (result != B_OK) {
                fListener->PrintError("entry not initialized!\n");
                return result;
        }

        BPath packagePath;
        if ((result = packageEntry.GetPath(&packagePath)) != B_OK) {
                fListener->PrintError("can't get path for entry '%s'!\n",
                        packageEntry.Name());
                return result;
        }

        BPackageReader packageReader(fListener);
        if ((result = packageReader.Init(packagePath.Path())) != B_OK) {
                fListener->PrintError("can't create package reader for '%s'!\n",
                        packagePath.Path());
                return result;
        }

        fPackageInfo.Clear();

        // parse package
        PackageContentHandler contentHandler(fListener, &fPackageInfo,
                packageReader.HeapReader(), fRepositoryInfo);
        if ((result = packageReader.ParseContent(&contentHandler)) != B_OK)
                return result;

        // determine package's checksum
        GeneralFileChecksumAccessor checksumAccessor(packageEntry);
        BString checksum;
        if ((result = checksumAccessor.GetChecksum(checksum)) != B_OK) {
                fListener->PrintError("can't compute checksum of file '%s'!\n",
                        packagePath.Path());
                return result;
        }
        fPackageInfo.SetChecksum(checksum);

        // register package's attributes
        if ((result = _RegisterCurrentPackageInfo()) != B_OK)
                return result;

        return B_OK;
}


status_t
RepositoryWriterImpl::_AddPackageInfo(const BPackageInfo& packageInfo)
{
        fPackageInfo = packageInfo;

        // register package's attributes
        status_t result = _RegisterCurrentPackageInfo();
        if (result != B_OK)
                return result;

        return B_OK;
}


status_t
RepositoryWriterImpl::_RegisterCurrentPackageInfo()
{
        status_t result = fPackageInfo.InitCheck();
        if (result != B_OK) {
                fListener->PrintError("package %s has incomplete package-info!\n",
                        fPackageInfo.Name().String());
                return result;
        }

        // reject package with a name that we've seen already
        if (fPackageNames->Contains(fPackageInfo.Name())) {
                fListener->PrintError("package %s has already been added!\n",
                        fPackageInfo.Name().String());
                return B_NAME_IN_USE;
        }

        // all packages must have the same vendor as the repository
        const BString& expectedVendor = fRepositoryInfo->Vendor();
        if (fPackageInfo.Vendor().ICompare(expectedVendor) != 0) {
                fListener->PrintError("package '%s' has unexpected vendor '%s' "
                        "(expected '%s')!\n", fPackageInfo.Name().String(),
                        fPackageInfo.Vendor().String(), expectedVendor.String());
                return B_BAD_DATA;
        }

        // all packages must have an architecture that's compatible with the one
        // used by the repository
        BPackageArchitecture expectedArchitecture = fRepositoryInfo->Architecture();
        if (fPackageInfo.Architecture() != expectedArchitecture
                && fPackageInfo.Architecture() != B_PACKAGE_ARCHITECTURE_ANY
                && fPackageInfo.Architecture() != B_PACKAGE_ARCHITECTURE_SOURCE) {
                fListener->PrintError(
                        "package '%s' has non-matching architecture '%s' "
                        "(expected '%s', '%s', or '%s')!\n", fPackageInfo.Name().String(),
                        BPackageInfo::kArchitectureNames[fPackageInfo.Architecture()],
                        BPackageInfo::kArchitectureNames[expectedArchitecture],
                        BPackageInfo::kArchitectureNames[B_PACKAGE_ARCHITECTURE_ANY],
                        BPackageInfo::kArchitectureNames[B_PACKAGE_ARCHITECTURE_SOURCE]);
                return B_BAD_DATA;
        }

        if ((result = fPackageNames->Add(fPackageInfo.Name())) != B_OK)
                return result;

        PackageAttribute* packageAttribute = AddStringAttribute(
                B_HPKG_ATTRIBUTE_ID_PACKAGE, fPackageInfo.Name(), PackageAttributes());
        RegisterPackageInfo(packageAttribute->children, fPackageInfo);
        fPackageCount++;
        fListener->OnPackageAdded(fPackageInfo);

        return B_OK;
}


status_t
RepositoryWriterImpl::_WriteRepositoryInfo(hpkg_repo_header& header,
        uint64& _length)
{
        // archive and flatten the repository info and write it
        BMessage archive;
        status_t result = fRepositoryInfo->Archive(&archive);
        if (result != B_OK) {
                fListener->PrintError("can't archive repository header!\n");
                return result;
        }

        ssize_t flattenedSize = archive.FlattenedSize();
        char buffer[flattenedSize];
        if ((result = archive.Flatten(buffer, flattenedSize)) != B_OK) {
                fListener->PrintError("can't flatten repository header!\n");
                return result;
        }

        WriteBuffer(buffer, flattenedSize);

        // notify listener
        fListener->OnRepositoryInfoSectionDone(flattenedSize);

        // update the header
        header.info_length = B_HOST_TO_BENDIAN_INT32(flattenedSize);

        _length = flattenedSize;
        return B_OK;
}


void
RepositoryWriterImpl::_WritePackageAttributes(hpkg_repo_header& header,
        uint64& _length)
{
        // write the package attributes (zlib writer on top of a file writer)
        uint64 startOffset = fHeapWriter->UncompressedHeapSize();

        uint32 stringsLength;
        uint32 stringsCount = WritePackageAttributes(PackageAttributes(),
                stringsLength);

        uint64 sectionSize = fHeapWriter->UncompressedHeapSize() - startOffset;

        fListener->OnPackageAttributesSectionDone(stringsCount, sectionSize);

        // update the header
        header.packages_length = B_HOST_TO_BENDIAN_INT64(sectionSize);
        header.packages_strings_count = B_HOST_TO_BENDIAN_INT64(stringsCount);
        header.packages_strings_length = B_HOST_TO_BENDIAN_INT64(stringsLength);

        _length = sectionSize;
}


}       // namespace BPrivate

}       // namespace BHPKG

}       // namespace BPackageKit