#include "erf_disk.h"

#include <saas/library/mapping/mapping.h>
#include <saas/protos/factors.pb.h>
#include <library/cpp/logger/global/global.h>

#include <google/protobuf/messagext.h>

#include <util/stream/file.h>

void TRTYErfDiskHeader::Remap() {
    VERIFY_WITH_LOG(!!FilePath, "Incorrect header remap usage (1)");
    VERIFY_WITH_LOG(RemapTable.size(), "Incorrect header remap usage (2)");
    NRTYServer::ReleaseFileMapping(FilePath + "/" + NameOfFile);
    {
        TRTYErfDiskManager::TCreationContext sourceContext(TPathName{FilePath}, NameOfFile, FHeaderDescription.Get(), true, true);
        TRTYErfDiskManager edm(sourceContext);
        edm.Open();

        TRTYErfDiskManager::TCreationContext targetContext(TPathName{FilePath}, "~" + NameOfFile, FDescription);
        targetContext.BlockCount = edm.Size();
        TRTYErfDiskManager edmWriter(targetContext);
        edmWriter.Open();

        // transform the ERF block for each document
        TBasicFactorStorage factorsAreaNew(FDescription->GetStaticFactorsCount());
        TBasicFactorStorage factorsAreaRemap(FHeaderDescription->GetStaticFactorsCount());
        for (size_t i = 0; i < edm.Size(); i++) {
            edm.ReadRaw(factorsAreaRemap, i);
            factorsAreaNew.Clear();
            for (int f = 0; f < Factors.ysize(); f++)
                if (RemapTable[f] != -1)
                    factorsAreaNew[RemapTable[f]] = factorsAreaRemap[f];
            edmWriter.Write(factorsAreaNew, i);
        }
    }

    TFsPath(FilePath + "/" + NameOfFile).DeleteIfExists();
    TFsPath(FilePath + "/" + NameOfFile + ".hdr").DeleteIfExists();

    TFsPath(FilePath + "/~" + NameOfFile).RenameTo(FilePath + "/" + NameOfFile);
    TFsPath(FilePath + "/~" + NameOfFile + ".hdr").RenameTo(FilePath + "/" + NameOfFile + ".hdr");

    TVector<int>().swap(RemapTable);
    NeedSaveFactors = false;
}

bool TRTYErfDiskHeader::NeedInRemap() const {
    VERIFY_WITH_LOG(!!FilePath, "Incorrect header usage");
    RemapTable.resize(Factors.size(), -1);
    bool keyRemap = false;

    if (FDescription->GetStaticFactorsCount() != FHeaderDescription->GetStaticFactorsCount())
        keyRemap = true;

    for (size_t i = 0; i < Factors.size(); i++) {
        if ((i >= FDescription->GetStaticFactorsCount()) ||
            (Factors[i].Name != FDescription->GetFactor(i).Name) ||
            (Factors[i].Width != FDescription->GetFactor(i).Width) ||
            ((ui32)Factors[i].Type != FDescription->GetFactor(i).Type)
            ) {
            keyRemap = true;
            for (size_t f = 0; f < FDescription->GetStaticFactorsCount(); f++) {
                if (FDescription->GetFactor(f).Name == Factors[i].Name) {
                    RemapTable[i] = f;
                    break;
                }
            }
        } else
            RemapTable[i] = i;
    }

    if (!keyRemap) {
        TVector<int>().swap(RemapTable);
    }

    for (ui32 i = 0; i < RemapTable.size(); ++i) {
        INFO_LOG << "RemapTable[" << i << "] = " << RemapTable[i] << Endl;
    }

    return keyRemap;
}

void TRTYErfDiskHeader::SaveHeader(const TString& filePath) const {
    if ((FilePath == filePath) && !NeedSaveFactors && !NeedSaveDocsCount) {
        return;
    }
    if ((FilePath == filePath) && !NeedSaveFactors && !RemapTable.empty()) {
        // the only thing that changed was DocsCount
        // avoid writing the header now as it should be re-mapped
        return;
    }
    VERIFY_WITH_LOG(RemapTable.empty(), "Erf Header %s needs Remap", (filePath + "/" + NameOfFile + ".hdr").data());
    TString hdrPath = filePath + "/" + NameOfFile + ".hdr";
    TFsPath(hdrPath).DeleteIfExists();
    try {
        TString tempHdrPath = filePath + "/~" + NameOfFile + ".hdr";
        {
            TFixedBufferFileOutput data(tempHdrPath);
            SaveHeaderToStream(data);
        }
        TFsPath(tempHdrPath).RenameTo(hdrPath);
    } catch (...) {
        FAIL_LOG("Can't write header file %s", hdrPath.data());
    }
    if (FilePath == filePath) {
        NeedSaveFactors = false;
        NeedSaveDocsCount = false;
    }
}

bool TRTYErfDiskHeader::LoadHeader(const TString& filePath) {
    FilePath = filePath;
    TString hdrPath = FilePath + "/" + NameOfFile + ".hdr";
    if (!NFs::Exists(hdrPath)) {
        LoadStaticFactors();
        NeedSaveFactors = true;
        return true;
    }
    TUnbufferedFileInput data(hdrPath.data());
    try {
        LoadHeaderFromStream(data);
    } catch (...) {
        FAIL_LOG("Incorrect erf header %s format", hdrPath.data());
    }
    return true;
}

void TRTYErfDiskHeader::LoadStaticFactors() {
    Factors.clear();
    for (ui32 i = 0; i < FDescription->GetStaticFactorsCount(); ++i) {
        Factors.push_back(FDescription->GetFactor(i));
    }
    FHeaderDescription.Reset(new TBaseStaticFactorsDescription(Factors));
    HeaderSumSizeBytes = FHeaderDescription->GetSumWidthBytes();
}

void TRTYErfDiskHeader::LoadFactorsFromFactorsHeader(const NRTYServer::TFactorsHeader &fh) {
    Factors.clear();
    for (ui32 i = 0; i < fh.FactorsSize(); i++) {
        const NRTYServer::TFactorSimple& fs = fh.GetFactors(i);
        ui32 width = fs.HasWidth() ? fs.GetWidth() : 32;
        TBitsReader::TFieldType type = fs.HasType() ? (TBitsReader::TFieldType)fs.GetType() : TBitsReader::ftInt;
        NRTYFactors::TSimpleFactorDescription factorDescription(fs.GetName(), fs.GetIndexGlobal(), width, type);
        Factors.push_back(NRTYFactors::TFactor(factorDescription, NRTYFactors::ftStatic));
    }
    FHeaderDescription.Reset(new TBaseStaticFactorsDescription(Factors));
    HeaderSumSizeBytes = FHeaderDescription->GetSumWidthBytes();
}

void TRTYErfDiskHeader::SaveFactorsToFactorsHeader(NRTYServer::TFactorsHeader &fh) const {
    for (int i = 0; i < Factors.ysize(); ++i) {
        NRTYServer::TFactorSimple* fs = fh.AddFactors();
        fs->SetName(Factors[i].Name);
        fs->SetIndexGlobal(Factors[i].IndexGlobal);
        fs->SetWidth(Factors[i].Width);
        fs->SetType((ui32)Factors[i].Type);
    }
}

bool TRTYErfDiskHeader::LoadHeaderFromStream(IInputStream& is) {
    NRTYServer::TFactorsHeader fh;
    ::google::protobuf::io::TProtoSerializer::Load(&is, fh);
    LoadFactorsFromFactorsHeader(fh);
    DocsCount = fh.HasDocsCount() ? fh.GetDocsCount() : 0;
    return true;
}

void TRTYErfDiskHeader::SaveHeaderToStream(IOutputStream& os) const {
    NRTYServer::TFactorsHeader fh;
    SaveFactorsToFactorsHeader(fh);
    fh.SetDocsCount(DocsCount);
    ::google::protobuf::io::TProtoSerializer::Save(&os, fh);
}


bool TRTYErfDiskManager::Write(const TBasicFactorStorage& erfBlock, ui32 position) const {
    if (Y_LIKELY(!!Data && !IsReadOnly)) {
        if (BlocksCount < position) {
            return false;
        }
        Remapper.PackSimple(erfBlock, Data + (ui64)position * BlockSizeBytes());
        return true;
    } else {
        FAIL_LOG("File %s not opened for write", GetFileName().c_str());
        return false;
    }
}

bool TRTYErfDiskManager::ReadRaw(TBasicFactorStorage& erfBlock, ui32 position) const {
    if (Y_LIKELY(!!Data)) {
        if (BlocksCount < position) {
            return false;
        }

        // read and unpack as is
        Remapper.UnpackSimple(erfBlock, Data + (ui64)position * BlockSizeBytes());
        return true;
    } else {
        FAIL_LOG("File %s not opened for read", GetFileName().c_str());
        return false;
    }
}

bool TRTYErfDiskManager::Read(TFactorView& factors, ui32 position) const {
    if (Y_LIKELY(!!Data)) {
        if (BlocksCount < position) {
            return false;
        }

        // read, unpack and remap
        Remapper.UnpackFormula(factors, Data + (ui64)position * BlockSizeBytes());
        return true;
    } else {
        FAIL_LOG("File %s not opened for read", GetFileName().c_str());
        return false;
    }
}

TRTYErfDiskManager::TRTYErfDiskManager(const TCreationContext& context, const TString& component)
    : IRTYErfManager(component)
    , Directory(context.Directory.PathName())
    , FileName(context.FileName)
    , IsReadOnly(context.ReadOnly)
    , UseGlobalMapping(context.UseGlobalMapping)
    , Data(nullptr)
    , BlocksCount(context.BlockCount)
    , Header(context.Description, FileName)
    , Remapper(*context.Description)
{
    VERIFY_WITH_LOG(BlocksCount < (1 << DOC_LEVEL_Bits), "Blocks count is incorrect: %lu", BlocksCount);
}

bool TRTYErfDiskManager::DoOpen() {
    INFO_LOG << "OPENING MANAGER: " << (ui64)Data << "/" << GetFileName() << "/" << (!!FileMap ? (i64)FileMap->GetFile().GetHandle() : 0) << Endl;
    if (!Data) {
        VERIFY_WITH_LOG(Header.LoadHeader(Directory), "Incorrect erf header");

        const TFsPath& path = GetFileName();
        const TString& name = path.GetPath();
        if ((IsReadOnly || !BlocksCount) && !NFs::Exists(name)) {
            WARNING_LOG << "File doesn't exist in read-only mode: " << name << Endl;
            return false;
        }

        if (IsReadOnly) {
            INFO_LOG << "File is opening in read-only mode: " << name << Endl;
        }
        TFile file(name, IsReadOnly ? RdOnly : (RdWr | OpenAlways));
        const size_t dataLength = file.GetLength();
        const ui32 blockSize = BlockSizeBytes();

        if (!BlocksCount && blockSize) {
            // we need to find out the BlocksCount from the file size and the header
            if (dataLength % blockSize == 0) {
                // this is the normal path, no warnings
                BlocksCount = dataLength / blockSize;

                if (Header.GetDocsCount() && BlocksCount != Header.GetDocsCount()) {
                    // no action needed, just a warning
                    WARNING_LOG << "Mismatch between erf data size and header for file " << name.Quote()
                                << ": " << BlocksCount << " vs " << Header.GetDocsCount()
                                << ". Using: " << BlocksCount << Endl;
                }
            } else {
                // start guessing
                WARNING_LOG << "Uneven erf file size for " << name.Quote()
                            << ". File size " << dataLength << " is not a multiple of block size " << blockSize << Endl;

                if (Header.GetDocsCount() != 0) {
                    BlocksCount = Header.GetDocsCount();
                    WARNING_LOG << "Taking erf size from header: " << Header.GetDocsCount() << Endl;
                } else {
                    BlocksCount = dataLength / blockSize;
                    WARNING_LOG << "No DocsCount in the erf header, estimating size: " << BlocksCount << Endl;
                }
            }
        }
        Header.SetDocsCount(BlocksCount);
        const size_t needFileSize = BlocksCount * BlockSizeBytes(); // used to have AlignUp here, not anymore

        bool resized = false;
        if (!IsReadOnly && (dataLength != needFileSize)) {
            file.Resize(needFileSize);
            resized = true;
            Header.SaveHeader(Directory);
        }

        if (resized && UseGlobalMapping) {
            NRTYServer::ReleaseFileMapping(name);
        }
        FileMap.Reset(NRTYServer::GetFileMapping(UseGlobalMapping, name, IsReadOnly));
        FileMap->Map(0, needFileSize);
        Data = (ui8*)FileMap->Ptr();

        if (dataLength < needFileSize) {
            memset((char*)Data + dataLength, 0, needFileSize - dataLength);
        }
    }
    DEBUG_LOG << "OPENED MANAGER: " << (ui64)Data << "/" << GetFileName() << "/" << (!!FileMap ? (i64)FileMap->GetFile().GetHandle() : 0) << Endl;
    return true;
}

ui32 TRTYErfDiskManager::Size() const {
    return BlocksCount;
}

bool TRTYErfDiskManager::DoClose() {
    DEBUG_LOG << "CLOSING MANAGER: " << (ui64)Data << "/" << GetFileName() << "/" << (!!FileMap ? (i64)FileMap->GetFile().GetHandle() : 0) << Endl;
    if (UseGlobalMapping)
        NRTYServer::ReleaseFileMapping(GetFileName());
    if (!!Data) {
        Data = nullptr;
        FileMap.Reset(nullptr);
        Header.SaveHeader(Directory);
        BlocksCount = 0;
    }
    DEBUG_LOG << "CLOSED MANAGER: " << (ui64)Data << "/" << GetFileName() << "/" << (!!FileMap ? (i64)FileMap->GetFile().GetHandle() : 0) << Endl;
    return true;
}

void TRTYErfDiskManager::RenameTo(const TString& dir, const TString& fileName) {
    TFsPath(GetFileName()).RenameTo(dir + "/" + fileName);
    Header.RenameTo(dir, fileName);
}

bool TRTYErfDiskManager::Resize() {
    ui32 blocksCountTemp = BlocksCount;
    Close();
    BlocksCount = 2 * blocksCountTemp;
    return Open();
}

bool TRTYErfDiskManager::RemapRequired() const {
    return Header.NeedInRemap();
}
