// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
// Mobius Forensic Toolkit
// Copyright (C) 2008,2009,2010,2011,2012,2013,2014,2015,2016,2017,2018,2019,2020,2021,2022 Eduardo Aguiar
//
// This program is free software; you can redistribute it and/or modify it
// under the terms of the GNU General Public License as published by the
// Free Software Foundation; either version 2, or (at your option) any later
// version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
// Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
#include "imagefile.h"
#include "imagefile_impl_null.h"
#include <mobius/exception.inc>
#include <mobius/io/file.h>
#include <mobius/string_functions.h>
#include <algorithm>
#include <stdexcept>
#include <unordered_map>

namespace mobius
{
namespace vfs
{
namespace
{
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Imagefile implementation struct
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
struct implementation
{
  std::string description;			// descriptions
  std::string extensions;		// file extensions, separated by '|'
  bool is_writeable;				// has write support
  imagefile_is_instance_type f_is_instance;	// is instance functor
  imagefile_builder_type f_builder;		// builder functor
};

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Imagefile implementations
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
static std::unordered_map <std::string, implementation> implementations_;

} // namespace

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Construct object
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
imagefile::imagefile ()
{
  impl_ = std::make_shared <imagefile_impl_null> ();
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Constructor from implementation pointer
//! \param impl implementation pointer
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
imagefile::imagefile (const std::shared_ptr <imagefile_impl_base>& impl)
  : impl_ (impl)
{
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Constructor from file object
//! \param f File object
//! \param type Imagefile type (default = "autodetect")
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
imagefile::imagefile (const mobius::io::file& f, const std::string& type)
  : impl_ (build_imagefile_implementation (f, type))
{
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Get metadata
//! \return imagefile metadata
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
mobius::metadata
imagefile::get_metadata () const
{
  // imagefile metadata
  mobius::metadata metadata = 
  {
    {
      "type",
      "Type",
      "std::string",
      get_type ()
    },
    {
      "size",
      "Size",
      "size_type",
      std::to_string (get_size ()) + " bytes"
    },
    {
      "sectors",
      "Number of sectors",
      "size_type",
      std::to_string (get_sectors ())
    },
    {
      "sector_size",
      "Sector size",
      "size_type",
      std::to_string (get_sector_size ()) + " bytes"
    }
  };

  // implementation specific metadata
  for (const auto& p : get_attributes ())
    {
      auto description = mobius::string::capitalize (mobius::string::replace (p.first, "_", " "));
      metadata.add (p.first, description, "std::string", p.second.to_string ());
    }

  return metadata;
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Add imagefile implementation class
//! \param id Implementation ID
//! \param description Description
//! \param extensions File extensions, separated by '|'
//! \param is_writeable True/False if imagefile has writing support
//! \param f_is_instance Is instance function
//! \param f_builder Builder function
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
void
add_imagefile_implementation (
  const std::string& id,
  const std::string& description,
  const std::string& extensions,
  bool is_writeable,
  imagefile_is_instance_type f_is_instance,
  imagefile_builder_type f_builder)
{
  implementations_[id] = implementation {
          description,
          extensions,
          is_writeable,
          f_is_instance,
          f_builder
  };
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Remove implementation class
//! \param id Implementation ID
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
void
remove_imagefile_implementation (const std::string& id)
{
  implementations_.erase (id);
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief List implementations available
//! \return Implementations available
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
std::vector <imagefile_implementation_info>
get_imagefile_implementations ()
{
  std::vector <imagefile_implementation_info> implementations;
  
  for (const auto& p : implementations_)
    implementations.push_back ({
            p.first,
            p.second.description,
            p.second.extensions,
            p.second.is_writeable}
    );
  
  std::sort (implementations.begin (),
             implementations.end (),
             [] (const auto& a, const auto& b){ return a.type < b.type; }
            );
  
  return implementations;
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Build implementation, according to ID
//! \param f File object
//! \param id Implementation ID
//! \return shared_ptr to implementation object
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
std::shared_ptr <imagefile_impl_base>
build_imagefile_implementation (const mobius::io::file& f, const std::string& id)
{
  // If type == "autodetect", use f_is_instance function to check if file
  // is an instance of this imagefile type.
  if (id == "autodetect")
    {
      for (const auto& p : implementations_)
        {
          if (p.second.f_is_instance (f))
            {
              // Use f_builder function to create imagefile implementation from
              // file <i>f</i>
              return p.second.f_builder (f);
            }
        }
      
      // fallback: raw imagefile
      return build_imagefile_implementation (f, "raw");
    }

  // Otherwise, if type is given, create imagefile using type implementation
  else
    {
      auto p = implementations_.find (id);

      if (p == implementations_.end ())
        throw std::runtime_error (MOBIUS_EXCEPTION_MSG ("implementation class '" + id + "' not found"));

      return p->second.f_builder (f);
    }
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Create new imagefile object by URL
//! \param url Imagefile URL
//! \param type Imagefile type (default = "autodetect")
//! \return Imagefile object
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
imagefile
new_imagefile_by_url (const std::string& url, const std::string& type)
{
  auto f = mobius::io::new_file_by_url (url);
  return imagefile (f, type);
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Create new imagefile object by path
//! \param path Imagefile path
//! \param type Imagefile type (default = "autodetect")
//! \return Imagefile object
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
imagefile
new_imagefile_by_path (const std::string& path, const std::string& type)
{
  auto f = mobius::io::new_file_by_path (path);
  return imagefile (f, type);
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Create new imagefile object from file
//! \param f File object
//! \param type Imagefile type (default = "autodetect")
//! \return Imagefile object
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
imagefile
new_imagefile_from_file (const mobius::io::file& f, const std::string& type)
{
  return imagefile (f, type);
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Normalize drive info
//! \param drive_vendor drive vendor
//! \param drive_model drive model
//! \param drive_serial_number drive serial number
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
void
normalize_drive_info (
  std::string& vendor,
  std::string& model,
  std::string& serial_number)
{
  // get drive vendor, if any
  size_t pos = model.find (' ');

  if (pos != std::string::npos)
    {
      vendor = mobius::string::capitalize (model.substr (0, pos));
      ++pos;
    }

  else
    {
      pos = 0;

      if (mobius::string::startswith (model, "ST"))
        vendor = "Seagate";
    }

  // get drive model
  model = model.substr (pos);
  pos = model.find ('-');

  if (pos != std::string::npos)
    model.erase (pos);

  // Western Digital serial numbers
  if (mobius::string::startswith (serial_number, "WD-"))
    {
      serial_number = serial_number.substr (3);
      vendor = "Western Digital";
    }
}

} // namespace vfs
} // namespace mobius
