///////////////////////////////////////////////////////////////////////

// File:        tessdatamanager.cpp

// Description: Functions to handle loading/combining tesseract data files.

// Author:      Daria Antonova

//

// (C) Copyright 2009, Google Inc.

// Licensed under the Apache License, Version 2.0 (the "License");

// you may not use this file except in compliance with the License.

// You may obtain a copy of the License at

// http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software

// distributed under the License is distributed on an "AS IS" BASIS,

// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

// See the License for the specific language governing permissions and

// limitations under the License.

//

///////////////////////////////////////////////////////////////////////

#ifdef HAVE_CONFIG_H

#  include "config_auto.h"

#endif

#include "tessdatamanager.h"

#include <cstdio>

#include <string>

#if defined(HAVE_LIBARCHIVE)

#  include <archive.h>

#  include <archive_entry.h>

#endif

#include <tesseract/version.h>

#include "errcode.h"

#include "helpers.h"

#include "params.h"

#include "serialis.h"

#include "tprintf.h"

namespace tesseract {

TessdataManager::TessdataManager() : reader_(nullptr), is_loaded_(false), swap_(false) {

  SetVersionString(TESSERACT_VERSION_STR);

}

TessdataManager::TessdataManager(FileReader reader)

    : reader_(reader), is_loaded_(false), swap_(false) {

  SetVersionString(TESSERACT_VERSION_STR);

}

// Lazily loads from the given filename. Won't actually read the file

// until it needs it.

void TessdataManager::LoadFileLater(const char *data_file_name) {

  Clear();

  data_file_name_ = data_file_name;

}

#if defined(HAVE_LIBARCHIVE)

bool TessdataManager::LoadArchiveFile(const char *filename) {

  bool result = false;

  archive *a = archive_read_new();

  if (a != nullptr) {

    archive_read_support_filter_all(a);

    archive_read_support_format_all(a);

    if (archive_read_open_filename(a, filename, 8192) == ARCHIVE_OK) {

      archive_entry *ae;

      while (archive_read_next_header(a, &ae) == ARCHIVE_OK) {

        const char *component = archive_entry_pathname(ae);

        if (component != nullptr) {

          TessdataType type;

          if (TessdataTypeFromFileName(component, &type)) {

            int64_t size = archive_entry_size(ae);

            if (size > 0) {

              entries_[type].resize(size);

              if (archive_read_data(a, &entries_[type][0], size) == size) {

                is_loaded_ = true;

              }

            }

          }

        }

      }

      result = is_loaded_;

    }

    archive_read_free(a);

  }

  return result;

}

#endif

bool TessdataManager::Init(const char *data_file_name) {

  std::vector<char> data;

  if (reader_ == nullptr) {

#if defined(HAVE_LIBARCHIVE)

    if (LoadArchiveFile(data_file_name)) {

      return true;

    }

#endif

    if (!LoadDataFromFile(data_file_name, &data)) {

      return false;

    }

  } else {

    if (!(*reader_)(data_file_name, &data)) {

      return false;

    }

  }

  return LoadMemBuffer(data_file_name, &data[0], data.size());

}

// Loads from the given memory buffer as if a file.

bool TessdataManager::LoadMemBuffer(const char *name, const char *data, int size) {

  // TODO: This method supports only the proprietary file format.

  if (size < 0) {

    return false;

  }

  Clear();

  data_file_name_ = name;

  TFile fp;

  fp.Open(data, size);

  uint32_t num_entries;

  if (!fp.DeSerialize(&num_entries)) {

    return false;

  }

  swap_ = num_entries > kMaxNumTessdataEntries;

  fp.set_swap(swap_);

  if (swap_) {

    ReverseN(&num_entries, sizeof(num_entries));

  }

  if (num_entries > kMaxNumTessdataEntries) {

    return false;

  }

  // TODO: optimize (no init required).

  std::vector<int64_t> offset_table(num_entries);

  if (!fp.DeSerialize(&offset_table[0], num_entries)) {

    return false;

  }

  for (unsigned i = 0; i < num_entries && i < TESSDATA_NUM_ENTRIES; ++i) {

    if (offset_table[i] >= 0) {

      if (offset_table[i] > size) {

        return false;

      }

      int64_t entry_size = size - offset_table[i];

      unsigned j = i + 1;

      while (j < num_entries && offset_table[j] == -1) {

        ++j;

      }

      if (j < num_entries) {

        if (offset_table[j] < 0 || offset_table[j] > size) {

          return false;

        }

        entry_size = offset_table[j] - offset_table[i];

      }

      if (entry_size < 0) {

        return false;

      }

      entries_[i].resize(entry_size);

      if (entry_size > 0 && !fp.DeSerialize(&entries_[i][0], entry_size)) {

        return false;

      }

    }

  }

  if (entries_[TESSDATA_VERSION].empty()) {

    SetVersionString("Pre-4.0.0");

  }

  is_loaded_ = true;

  return true;

}

// Overwrites a single entry of the given type.

void TessdataManager::OverwriteEntry(TessdataType type, const char *data, int size) {

  is_loaded_ = true;

  entries_[type].resize(size);

  memcpy(&entries_[type][0], data, size);

}

// Saves to the given filename.

bool TessdataManager::SaveFile(const char *filename, FileWriter writer) const {

  // TODO: This method supports only the proprietary file format.

  ASSERT_HOST(is_loaded_);

  std::vector<char> data;

  Serialize(&data);

  if (writer == nullptr) {

    return SaveDataToFile(data, filename);

  } else {

    return (*writer)(data, filename);

  }

}

// Serializes to the given vector.

void TessdataManager::Serialize(std::vector<char> *data) const {

  // TODO: This method supports only the proprietary file format.

  ASSERT_HOST(is_loaded_);

  // Compute the offset_table and total size.

  int64_t offset_table[TESSDATA_NUM_ENTRIES];

  int64_t offset = sizeof(int32_t) + sizeof(offset_table);

  for (unsigned i = 0; i < TESSDATA_NUM_ENTRIES; ++i) {

    if (entries_[i].empty()) {

      offset_table[i] = -1;

    } else {

      offset_table[i] = offset;

      offset += entries_[i].size();

    }

  }

  data->resize(offset, 0);

  int32_t num_entries = TESSDATA_NUM_ENTRIES;

  TFile fp;

  fp.OpenWrite(data);

  fp.Serialize(&num_entries);

  fp.Serialize(&offset_table[0], countof(offset_table));

  for (const auto &entry : entries_) {

    if (!entry.empty()) {

      fp.Serialize(&entry[0], entry.size());

    }

  }

}

// Resets to the initial state, keeping the reader.

void TessdataManager::Clear() {

  for (auto &entry : entries_) {

    entry.clear();

  }

  is_loaded_ = false;

}

// Prints a directory of contents.

void TessdataManager::Directory() const {

  printf("Version:%s\n", VersionString().c_str());

  auto offset = TESSDATA_NUM_ENTRIES * sizeof(int64_t);

  for (unsigned i = 0; i < TESSDATA_NUM_ENTRIES; ++i) {

    if (!entries_[i].empty()) {

      printf("%u:%s:size=%zu, offset=%zu\n", i, kTessdataFileSuffixes[i], entries_[i].size(),

              offset);

      offset += entries_[i].size();

    }

  }

}

// Opens the given TFile pointer to the given component type.

// Returns false in case of failure.

bool TessdataManager::GetComponent(TessdataType type, TFile *fp) {

  if (!is_loaded_ && !Init(data_file_name_.c_str())) {

    return false;

  }

  const TessdataManager *const_this = this;

  return const_this->GetComponent(type, fp);

}

// As non-const version except it can't load the component if not already

// loaded.

bool TessdataManager::GetComponent(TessdataType type, TFile *fp) const {

  ASSERT_HOST(is_loaded_);

  if (entries_[type].empty()) {

    return false;

  }

  fp->Open(&entries_[type][0], entries_[type].size());

  fp->set_swap(swap_);

  return true;

}

// Returns the current version string.

std::string TessdataManager::VersionString() const {

  return std::string(&entries_[TESSDATA_VERSION][0], entries_[TESSDATA_VERSION].size());

}

// Sets the version string to the given v_str.

void TessdataManager::SetVersionString(const std::string &v_str) {

  entries_[TESSDATA_VERSION].resize(v_str.size());

  memcpy(&entries_[TESSDATA_VERSION][0], v_str.data(), v_str.size());

}

bool TessdataManager::CombineDataFiles(const char *language_data_path_prefix,

                                       const char *output_filename) {

  // Load individual tessdata components from files.

  for (auto filesuffix : kTessdataFileSuffixes) {

    TessdataType type;

    ASSERT_HOST(TessdataTypeFromFileSuffix(filesuffix, &type));

    std::string filename = language_data_path_prefix;

    filename += filesuffix;

    FILE *fp = fopen(filename.c_str(), "rb");

    if (fp != nullptr) {

      fclose(fp);

      if (!LoadDataFromFile(filename.c_str(), &entries_[type])) {

        tprintf("Load of file %s failed!\n", filename.c_str());

        return false;

      }

    }

  }

  is_loaded_ = true;

  // Make sure that the required components are present.

  if (!IsBaseAvailable() && !IsLSTMAvailable()) {

    tprintf(

        "Error: traineddata file must contain at least (a unicharset file"

        " and inttemp) OR an lstm file.\n");

    return false;

  }

  // Write updated data to the output traineddata file.

  return SaveFile(output_filename, nullptr);

}

bool TessdataManager::OverwriteComponents(const char *new_traineddata_filename,

                                          char **component_filenames, int num_new_components) {

  // Open the files with the new components.

  // TODO: This method supports only the proprietary file format.

  for (int i = 0; i < num_new_components; ++i) {

    TessdataType type;

    if (TessdataTypeFromFileName(component_filenames[i], &type)) {

      if (!LoadDataFromFile(component_filenames[i], &entries_[type])) {

        tprintf("Failed to read component file:%s\n", component_filenames[i]);

        return false;

      }

    }

  }

  // Write updated data to the output traineddata file.

  return SaveFile(new_traineddata_filename, nullptr);

}

bool TessdataManager::ExtractToFile(const char *filename) {

  TessdataType type = TESSDATA_NUM_ENTRIES;

  ASSERT_HOST(tesseract::TessdataManager::TessdataTypeFromFileName(filename, &type));

  if (entries_[type].empty()) {

    return false;

  }

  return SaveDataToFile(entries_[type], filename);

}

bool TessdataManager::TessdataTypeFromFileSuffix(const char *suffix, TessdataType *type) {

  for (unsigned i = 0; i < TESSDATA_NUM_ENTRIES; ++i) {

    if (strcmp(kTessdataFileSuffixes[i], suffix) == 0) {

      *type = static_cast<TessdataType>(i);

      return true;

    }

  }

#if !defined(NDEBUG)

  tprintf(

      "TessdataManager can't determine which tessdata"

      " component is represented by %s\n",

      suffix);

#endif

  return false;

}

bool TessdataManager::TessdataTypeFromFileName(const char *filename, TessdataType *type) {

  // Get the file suffix (extension)

  const char *suffix = strrchr(filename, '.');

  if (suffix == nullptr || *(++suffix) == '\0') {

    return false;

  }

  return TessdataTypeFromFileSuffix(suffix, type);

}

} // namespace tesseract