package_fetcher.cpp

// Copyright (c) 2023, QuantStack and Mamba Contributors
//
// Distributed under the terms of the BSD 3-Clause License.
//
// The full license is in the file LICENSE, distributed with this software.

#include <algorithm>

#include "mamba/core/invoke.hpp"
#include "mamba/core/output.hpp"
#include "mamba/core/package_fetcher.hpp"
#include "mamba/core/util.hpp"
#include "mamba/specs/archive.hpp"
#include "mamba/util/string.hpp"
#include "mamba/validation/tools.hpp"

namespace mamba
{

    /**
     * Components passed to the download layer to build a package fetch request.
     *
     * The download system uses mirror_name to look up the appropriate mirror
     * (PassThroughMirror, HTTPMirror, or OCIMirror) and url_path as the resource
     * to fetch. The format of these fields depends on the package source type.
     */
    struct DownloadRequestComponents
    {
        std::string mirror_name;  ///< Mirror lookup key: "" for PassThrough, channel/OCI URL
                                  ///< otherwise
        std::string url_path;     ///< Full URL for PassThrough, or "platform/filename" for
                                  ///< channel-based mirrors
    };

    /**
     * Compute the download request components for a package.
     *
     * The download layer expects different URL formats depending on the source:
     *
     * - Plain HTTP (no auth): PassThroughMirror uses the full package_url directly.
     *   Returns mirror_name="" and url_path=package_url.
     *
     * - OCI registries: OCIMirror builds URLs from channel + path. The package_url
     *   format is not suitable. Returns mirror_name=channel (e.g. oci://.../conda-forge)
     *   and url_path=platform/filename.
     *
     * - Authenticated URLs: HTTPMirror must not receive credentials in the URL
     *   (they are set via libcurl CURLUPart). Returns mirror_name=channel and
     *   url_path=platform/filename so the mirror can build a clean URL.
     */
    auto get_download_request_components(const specs::PackageInfo& pkg) -> DownloadRequestComponents
    {
        constexpr std::string_view oci_scheme = "oci://";
        const bool use_oci = util::starts_with(pkg.package_url, oci_scheme);
        const bool use_auth = std::regex_search(pkg.package_url, http_basicauth_regex())
                              || std::regex_search(pkg.package_url, token_regex());

        if (use_oci || use_auth)
        {
            return {
                .mirror_name = pkg.channel,
                .url_path = util::concat(pkg.platform, '/', pkg.filename),
            };
        }
        return {
            .mirror_name = "",
            .url_path = pkg.package_url,
        };
    }

    /**********************
     * PackageExtractTask *
     **********************/

    PackageExtractTask::PackageExtractTask(PackageFetcher* fetcher, ExtractOptions options)
        : p_fetcher(fetcher)
        , m_options(std::move(options))
    {
    }

    const std::string& PackageExtractTask::name() const
    {
        return p_fetcher->name();
    }

    bool PackageExtractTask::needs_download() const
    {
        return p_fetcher->needs_download();
    }

    void PackageExtractTask::set_progress_callback(progress_callback_t cb)
    {
        m_progress_callback = std::move(cb);
    }

    auto PackageExtractTask::run() -> Result
    {
        bool is_valid = true;
        bool is_extracted = p_fetcher->extract(m_options);
        return { is_valid, is_extracted };
    }

    auto PackageExtractTask::run(std::size_t downloaded_size) -> Result
    {
        using ValidationResult = PackageFetcher::ValidationResult;
        ValidationResult validation_res = p_fetcher->validate(downloaded_size, get_progress_callback());
        const bool is_valid = validation_res == ValidationResult::VALID;
        bool is_extracted = false;
        if (is_valid)
        {
            is_extracted = p_fetcher->extract(m_options, get_progress_callback());
        }
        return { is_valid, is_extracted };
    }

    auto PackageExtractTask::get_progress_callback() -> progress_callback_t*
    {
        if (m_progress_callback.has_value())
        {
            return &m_progress_callback.value();
        }
        else
        {
            return nullptr;
        }
    }

    /*******************
     * PatckageFetcher *
     *******************/

    struct PackageFetcher::CheckSumParams
    {
        std::string_view expected;
        std::string_view actual;
        std::string_view name;
        ValidationResult error;
    };

    PackageFetcher::PackageFetcher(const specs::PackageInfo& pkg_info, MultiPackageCache& caches)
        : m_package_info(pkg_info)
        , m_caches(&caches)
    {
        const fs::u8path extracted_cache = m_caches->get_extracted_dir_path(m_package_info);
        if (extracted_cache.empty())
        {
            const fs::u8path tarball_cache = m_caches->get_tarball_path(m_package_info);
            auto& cache = m_caches->first_writable_cache(true);
            m_cache_path = cache.path() / package_cache_folder_relative_path(m_package_info);
            fs::create_directories(m_cache_path);

            if (!tarball_cache.empty())
            {
                LOG_DEBUG << "Found valid tarball cache at '" << tarball_cache.string() << "'";
                cache.clear_query_cache(m_package_info);
                m_tarball_path = tarball_cache / filename();
                m_needs_extract = true;
                LOG_DEBUG << "Using cached tarball '" << filename() << "'";
            }
            else
            {
                m_caches->clear_query_cache(m_package_info);
                // need to download this file
                const DownloadRequestComponents components = get_download_request_components(
                    m_package_info
                );
                LOG_DEBUG << "Adding '" << name() << "' to download targets from '"
                          << hide_secrets(components.mirror_name) << "/" << components.url_path
                          << "'";
                m_tarball_path = m_cache_path / filename();
                m_needs_extract = true;
                m_needs_download = true;
            }
        }
        else
        {
            LOG_DEBUG << "Using cached '" << name() << "'";
        }
    }

    const std::string& PackageFetcher::name() const
    {
        return m_package_info.name;
    }

    bool PackageFetcher::needs_download() const
    {
        return m_needs_download;
    }

    bool PackageFetcher::needs_extract() const
    {
        return m_needs_extract;
    }

    download::Request
    PackageFetcher::build_download_request(std::optional<post_download_success_t> callback)
    {
        // download::Request request(name(), download::MirrorName(""), url(),
        // m_tarball_path.string());
        const DownloadRequestComponents components = get_download_request_components(m_package_info);
        download::Request request(
            name(),
            download::MirrorName(components.mirror_name),
            components.url_path,
            m_tarball_path.string()
        );
        request.expected_size = expected_size();
        request.sha256 = sha256();

        request.on_success = [this, cb = std::move(callback)](const download::Success& success)
        {
            LOG_INFO << "Download finished, tarball available at '" << m_tarball_path.string() << "'";
            if (cb.has_value())
            {
                cb.value()(success.transfer.downloaded_size);
            }
            m_needs_download = false;
            m_downloaded_url = m_package_info.package_url;
            return expected_t<void>();
        };

        request.on_failure = [](const download::Error& error)
        {
            if (error.transfer.has_value())
            {
                LOG_ERROR << "Failed to download package from "
                          << error.transfer.value().effective_url << " (status "
                          << error.transfer.value().http_status << ")\n"
                          << "If you see this message repeatedly, the state of your installation might be corrupted,\n"
                          << "in which case running `mamba clean --all` might fix it.\n\n"
                          << "If you continue to meet this problem, please search or report an issue\n"
                          << "on  mamba's issue tracker: https://github.com/mamba-org/mamba/issues/";
            }
            else
            {
                LOG_WARNING << error.message;
            }

            if (error.retry_wait_seconds.has_value())
            {
                LOG_WARNING << "Retrying in " << error.retry_wait_seconds.value() << " seconds";
            }
        };

        return request;
    }

    auto PackageFetcher::validate(std::size_t downloaded_size, progress_callback_t* cb) const
        -> ValidationResult
    {
        update_monitor(cb, PackageExtractEvent::validate_update);
        ValidationResult res = validate_size(downloaded_size);
        if (res != ValidationResult::VALID)
        {
            update_monitor(cb, PackageExtractEvent::validate_failure);
            return res;
        }

        interruption_point();

        if (!sha256().empty())
        {
            res = validate_checksum(
                {
                    /* .expected= */ sha256(),
                    /* .actual= */ validation::sha256sum(m_tarball_path),
                    /* .name= */ "SHA256",
                    /* .error= */ ValidationResult::SHA256_ERROR,
                }
            );
        }
        else if (!md5().empty())
        {
            res = validate_checksum(
                {
                    /* .expected= */ md5(),
                    /* .actual= */ validation::md5sum(m_tarball_path),
                    /* .name= */ "MD5",
                    /* .error= */ ValidationResult::MD5SUM_ERROR,
                }
            );
        }

        auto event = res == ValidationResult::VALID ? PackageExtractEvent::validate_success
                                                    : PackageExtractEvent::validate_failure;
        update_monitor(cb, event);
        return res;
    }

    namespace
    {
        fs::u8path get_extract_path(const std::string& filename, const fs::u8path& cache_path)
        {
            std::string fn = filename;
            if (util::ends_with(fn, ".tar.bz2"))
            {
                fn = fn.substr(0, fn.size() - 8);
            }
            else if (util::ends_with(fn, ".conda"))
            {
                fn = fn.substr(0, fn.size() - 6);
            }
            else
            {
                LOG_ERROR << "Unknown package format '" << filename << "'";
                throw std::runtime_error("Unknown package format.");
            }
            return cache_path / fn;
        }

        void clear_extract_path(const fs::u8path& path)
        {
            if (fs::exists(path))
            {
                LOG_DEBUG << "Removing '" << path.string() << "' before extracting it again";
                fs::remove_all(path);
            }
        }

        void extract_impl(
            const fs::u8path& tarball_path,
            const fs::u8path& extract_path,
            const ExtractOptions& options
        )
        {
            // Use non-subproc version if concurrency is disabled to avoid
            // any potential subprocess issues
            if (PackageFetcherSemaphore::get_max() == 1)
            {
                mamba::extract(tarball_path, extract_path, options);
            }
            else
            {
                mamba::extract_subproc(tarball_path, extract_path, options);
            }
        }
    }

    bool PackageFetcher::extract(const ExtractOptions& options, progress_callback_t* cb)
    {
        // Extracting is __not__ yet thread safe it seems...
        interruption_point();

        LOG_DEBUG << "Waiting for decompression " << m_tarball_path;
        update_monitor(cb, PackageExtractEvent::extract_update);

        {
            std::lock_guard<counting_semaphore> lock(PackageFetcherSemaphore::semaphore);
            interruption_point();
            LOG_DEBUG << "Decompressing '" << m_tarball_path.string() << "'";
            try
            {
                const fs::u8path extract_path = get_extract_path(filename(), m_cache_path);
                // Be sure the first writable cache doesn't contain invalid extracted package
                clear_extract_path(extract_path);
                extract_impl(m_tarball_path, extract_path, options);

                interruption_point();
                LOG_DEBUG << "Extracted to '" << extract_path.string() << "'";
                write_repodata_record(extract_path);
                update_urls_txt();
                m_caches->clear_query_cache(m_package_info);
                update_monitor(cb, PackageExtractEvent::extract_success);
            }
            catch (const std::logic_error&)
            {
                // `std::logic_error` indicates a programming bug (e.g., missing
                // `_initialized` sentinel). Re-throw to fail hard.
                throw;
            }
            catch (const std::exception& e)
            {
                Console::instance().print(filename() + " extraction failed");
                LOG_ERROR << "Error when extracting package: " << e.what();
                update_monitor(cb, PackageExtractEvent::extract_failure);
                return false;
            }
        }
        m_needs_extract = false;
        return true;
    }

    PackageExtractTask PackageFetcher::build_extract_task(ExtractOptions options)
    {
        return { this, std::move(options) };
    }

    void PackageFetcher::clear_cache() const
    {
        const auto remove_extracted_at = [&](const fs::u8path& cache_path)
        { fs::remove_all(get_extract_path(filename(), cache_path)); };

        fs::remove_all(m_tarball_path);
        remove_extracted_at(m_cache_path);

        // Tarballs may use the flat layout while extraction uses the hierarchical layout.
        const fs::u8path tarball_parent = m_tarball_path.parent_path();
        if (tarball_parent != m_cache_path)
        {
            remove_extracted_at(tarball_parent);
        }

        if (m_caches)
        {
            m_caches->clear_query_cache(m_package_info);
        }
    }

    /*******************
     * Private methods *
     *******************/

    const std::string& PackageFetcher::filename() const
    {
        return m_package_info.filename;
    }

    const std::string& PackageFetcher::url() const
    {
        return m_downloaded_url;
    }

    const std::string& PackageFetcher::sha256() const
    {
        return m_package_info.sha256;
    }

    const std::string& PackageFetcher::md5() const
    {
        return m_package_info.md5;
    }

    std::size_t PackageFetcher::expected_size() const
    {
        return m_package_info.size;
    }

    auto PackageFetcher::validate_size(std::size_t downloaded_size) const -> ValidationResult
    {
        auto res = ValidationResult::VALID;
        if (expected_size() && expected_size() != downloaded_size)
        {
            res = ValidationResult::SIZE_ERROR;
            LOG_ERROR << "File not valid: file size doesn't match expectation " << m_tarball_path
                      << "\nExpected: " << expected_size() << "\nActual: " << downloaded_size
                      << "\n";
            Console::instance().print(filename() + " tarball has incorrect size");
        }
        return res;
    }

    auto PackageFetcher::validate_checksum(const CheckSumParams& params) const -> ValidationResult
    {
        auto res = ValidationResult::VALID;
        if (params.actual != params.expected)
        {
            res = params.error;
            LOG_ERROR << "File not valid: " << params.name << " doesn't match expectation "
                      << m_tarball_path << "\nExpected: " << params.expected
                      << "\nActual: " << params.actual << "\n";
            Console::instance().print(util::concat(filename(), " tarball has incorrect ", params.name));
            // TODO: terminate monitor
        }
        return res;
    }

    void PackageFetcher::write_repodata_record(const fs::u8path& base_path) const
    {
        const fs::u8path repodata_record_path = base_path / "info" / "repodata_record.json";
        const fs::u8path index_path = base_path / "info" / "index.json";

        nlohmann::json index;
        std::ifstream index_file = open_ifstream(index_path);
        index_file >> index;

        nlohmann::json repodata_record = m_package_info;

        // `from_json()` does NOT set `_initialized` because it deserializes
        // already-written cache files for display/query purposes. Those `PackageInfo`
        // objects are never passed to this function — they're used for `mamba list`,
        // dependency computation, etc.
        // See `PackageInfo::defaulted_keys`. Issue #4095.
        const bool contains_initialized = std::ranges::find(
                                              m_package_info.defaulted_keys,
                                              specs::defaulted_key::initialized
                                          )
                                          != m_package_info.defaulted_keys.end();
        if (!contains_initialized)
        {
            throw std::logic_error(
                "`PackageInfo` missing `_initialized` sentinel in `defaulted_keys`. "
                "This indicates a bug in the code path that created this `PackageInfo`. "
                "See GitHub issue #4095."
            );
        }

        // - URL-derived packages: listed fields have stub values (0, "", [])
        //   → erase them so `index.json` provides correct values
        // - Solver-derived packages: only `_initialized` in list
        //   → nothing erased, all fields preserved (including channel patches)
        for (const auto& key : m_package_info.defaulted_keys)
        {
            if (key != specs::defaulted_key::initialized)
            {
                repodata_record.erase(key);
            }
        }

        // `insert()` only adds MISSING keys — solver-derived fields (including
        // channel patches with intentionally empty arrays) are preserved.
        repodata_record.insert(index.cbegin(), index.cend());

        if (repodata_record.find("size") == repodata_record.end() || repodata_record["size"] == 0)
        {
            repodata_record["size"] = fs::file_size(m_tarball_path);
        }

        // Matches conda behavior where `depends` and `constrains` are always present.
        // Some packages (like `nlohmann_json-abi`) don't have `depends` in `index.json`.
        if (!repodata_record.contains("depends"))
        {
            repodata_record["depends"] = nlohmann::json::array();
        }
        if (!repodata_record.contains("constrains"))
        {
            repodata_record["constrains"] = nlohmann::json::array();
        }

        // Matches conda behavior: omit `track_features` when empty to reduce JSON noise.
        if (repodata_record.contains("track_features"))
        {
            const auto& tf = repodata_record["track_features"];
            bool is_empty = tf.is_null() || (tf.is_string() && tf.get<std::string>().empty())
                            || (tf.is_array() && tf.empty());
            if (is_empty)
            {
                repodata_record.erase("track_features");
            }
        }

        // Compute missing checksums from tarball. Issue #4095.
        auto needs_checksum = [&repodata_record](std::string_view key)
        {
            return !repodata_record.contains(key) || !repodata_record[key].is_string()
                   || repodata_record[key].get<std::string>().empty();
        };

        if (needs_checksum("md5"))
        {
            repodata_record["md5"] = validation::md5sum(m_tarball_path);
        }

        if (needs_checksum("sha256"))
        {
            repodata_record["sha256"] = validation::sha256sum(m_tarball_path);
        }

        std::ofstream repodata_record_file(repodata_record_path.std_path());
        repodata_record_file << repodata_record.dump(4);
    }

    namespace
    {
        std::mutex urls_txt_mutex;
    }

    void PackageFetcher::update_urls_txt() const
    {
        // TODO: check if this lock is really required
        std::unique_lock lock{ urls_txt_mutex };
        const auto urls_file_path = m_cache_path / "urls.txt";
        std::ofstream urls_txt(urls_file_path.std_path(), std::ios::app);
        urls_txt << url() << std::endl;
    }

    void PackageFetcher::update_monitor(progress_callback_t* cb, PackageExtractEvent event) const
    {
        if (cb)
        {
            // We dont want to propagate errors coming from user's callbacks
            [[maybe_unused]] auto result = safe_invoke(*cb, event);
        }
    }

    /***************************
     * PackageFetcherSemaphore *
     ***************************/

    counting_semaphore PackageFetcherSemaphore::semaphore(0);

    std::ptrdiff_t PackageFetcherSemaphore::get_max()
    {
        return PackageFetcherSemaphore::semaphore.get_max();
    }

    void PackageFetcherSemaphore::set_max(int value)
    {
        PackageFetcherSemaphore::semaphore.set_max(value);
    }
}