// Author: Derek Barnett

#include "PbbamInternalConfig.h"

#include "DataSetIO.h"

#include <algorithm>
#include <cassert>
#include <cstddef>
#include <exception>
#include <fstream>
#include <iostream>

#include <boost/algorithm/string.hpp>

#include "FileUtils.h"
#include "FofnReader.h"
#include "XmlReader.h"
#include "XmlWriter.h"
#include "pbbam/MakeUnique.h"
#include "pbbam/StringUtilities.h"

namespace PacBio {
namespace BAM {
namespace {

std::unique_ptr<DataSetBase> DataSetFromXml(const std::string& xmlFn)
{
    std::ifstream in(xmlFn);
    if (!in) throw std::runtime_error{"DataSet: could not open XML file for reading: " + xmlFn};
    return XmlReader::FromStream(in);
}

std::unique_ptr<DataSetBase> DataSetFromBam(const std::string& bamFn)
{
    // peek at sort order to determine if file should be an AlignmentSet or else SubreadSet
    const auto bamFile = BamFile{bamFn};
    const auto& header = bamFile.Header();
    const auto aligned = header.SortOrder() == "coordinate";

    std::unique_ptr<DataSetBase> dataset;
    if (aligned)
        dataset = std::make_unique<AlignmentSet>();
    else
        dataset = std::make_unique<SubreadSet>();

    auto& resources = dataset->ExternalResources();
    resources.Add(ExternalResource(BamFile(bamFn)));
    return dataset;
}

std::unique_ptr<DataSetBase> DataSetFromFasta(const std::string& fasta)
{
    // make FASTA data set
    auto dataset = std::make_unique<ReferenceSet>();
    auto& resources = dataset->ExternalResources();
    resources.Add(ExternalResource("PacBio.ReferenceFile.ReferenceFastaFile", fasta));
    return std::move(dataset);
}

std::unique_ptr<DataSetBase> DataSetFromFofn(const std::string& fofn)
{
    const auto fofnDir = FileUtils::DirectoryName(fofn);
    std::ifstream in(fofn);
    if (!in) throw std::runtime_error{"DataSet: could not open FOFN for reading: " + fofn};

    auto filenames = FofnReader::Files(in);
    std::transform(
        filenames.begin(), filenames.end(), filenames.begin(),
        [&fofnDir](const std::string fn) { return FileUtils::ResolvedFilePath(fn, fofnDir); });
    return DataSetIO::FromUris(filenames);
}

std::unique_ptr<DataSetBase> DataSetFromUri(const std::string& uri)
{
    // NOTE: this says URI, but we're not quite handling filenames as true URIs
    //       basically just treating as a regular filename for now

    // handle on extension
    if (boost::algorithm::iends_with(uri, ".xml"))
        return DataSetFromXml(uri);
    else if (boost::algorithm::iends_with(uri, ".bam"))
        return DataSetFromBam(uri);
    else if (boost::algorithm::iends_with(uri, ".fofn"))
        return DataSetFromFofn(uri);
    else if (boost::algorithm::iends_with(uri, ".fasta") ||
             boost::algorithm::iends_with(uri, ".fa")) {
        return DataSetFromFasta(uri);
    }

    // unknown filename extension
    throw std::runtime_error{"DataSet: unsupported extension on input file: " + uri};
}

}  // anonymous

std::unique_ptr<DataSetBase> DataSetIO::FromUri(const std::string& uri)
{
    return FromUris(std::vector<std::string>(1, uri));
}

std::unique_ptr<DataSetBase> DataSetIO::FromUris(const std::vector<std::string>& uris)
{
    if (uris.empty()) throw std::runtime_error{"DataSet: empty input URI list"};

    // create dataset(s) from URI(s)
    std::vector<std::unique_ptr<DataSetBase> > datasets;
    datasets.reserve(uris.size());
    for (const auto& uri : uris)
        datasets.emplace_back(DataSetFromUri(uri));
    assert(!datasets.empty());

    // if only 1, just return
    if (datasets.size() == 1) return std::unique_ptr<DataSetBase>(datasets.front().release());

    // else merge
    else {
        auto& result = datasets.at(0);
        for (size_t i = 1; i < datasets.size(); ++i) {
            const auto& next = datasets.at(i);
            *result += *next;
        }
        return std::move(result);
    }
}

std::unique_ptr<DataSetBase> DataSetIO::FromXmlString(const std::string& xml)
{
    if (xml.empty()) throw std::runtime_error{"DataSet: cannot load from empty XML string"};
    std::istringstream s{xml};
    return XmlReader::FromStream(s);
}

void DataSetIO::ToFile(const std::unique_ptr<DataSetBase>& dataset, const std::string& fn)
{
    DataSetIO::ToFile(*dataset, fn);
}

void DataSetIO::ToStream(const std::unique_ptr<DataSetBase>& dataset, std::ostream& out)
{
    DataSetIO::ToStream(*dataset, out);
}

void DataSetIO::ToFile(const DataSetBase& dataset, const std::string& fn)
{
    std::ofstream out(fn);
    if (!out) throw std::runtime_error{"DataSet: could not open XML file for writing: " + fn};
    XmlWriter::ToStream(dataset, out);
}

void DataSetIO::ToStream(const DataSetBase& dataset, std::ostream& out)
{
    XmlWriter::ToStream(dataset, out);
}

}  // namespace BAM
}  // namespace PacBio
