#include "engine/models/pocket_tts/loader.h"

#include "engine/framework/model_spec/package.h"
#include "engine/framework/runtime/options.h"
#include "engine/models/pocket_tts/assets.h"
#include "engine/models/pocket_tts/session.h"

#include <filesystem>
#include <optional>
#include <stdexcept>
#include <string>
#include <string_view>
#include <utility>

namespace engine::models::pocket_tts {
namespace {

std::string requested_language(const runtime::ModelLoadRequest & request) {
    return runtime::find_option(request.options, {"language"}).value_or("english");
}

runtime::CapabilitySet capabilities(const PocketTTSAssets & assets) {
    runtime::CapabilitySet out;
    out.supported_tasks = {
        {runtime::VoiceTaskKind::Tts, {runtime::RunMode::Offline}},
    };
    out.languages = {assets.language};
    out.supports_speaker_reference = true;
    out.supports_style_condition = false;
    return out;
}

runtime::ModelMetadata metadata(const PocketTTSAssets & assets) {
    runtime::ModelMetadata out;
    out.family = "pocket_tts";
    out.variant = assets.language;
    out.description = "PocketTTS loaded from local assets.";
    return out;
}

runtime::ModelCliInterface cli(const PocketTTSAssets &) {
    runtime::ModelCliInterface out;
    out.session_options = {
        {
            "pocket_tts.voice_state_cache_slots",
            "n",
            "Prepared voice-state cache slots; default 4, set 0 to disable.",
        },
    };
    return out;
}

std::shared_ptr<const PocketTTSAssets> load_request_assets(const runtime::ModelLoadRequest & request) {
    return load_pocket_tts_assets(request.model_path, requested_language(request));
}

class PocketTTSLoader final : public runtime::IVoiceModelLoader {
public:
    std::string family() const override {
        return "pocket_tts";
    }

    runtime::CapabilitySet advertised_capabilities() const override {
        runtime::CapabilitySet out;
        out.supported_tasks = {
            {runtime::VoiceTaskKind::Tts, {runtime::RunMode::Offline}},
        };
        out.supports_speaker_reference = true;
        out.supports_style_condition = true;
        return out;
    }

    bool can_load(const runtime::ModelLoadRequest & request) const override {
        if (request.family_hint.has_value() && *request.family_hint != family()) {
            return false;
        }
        try {
            (void) engine::model_spec::load_resource_bundle(
                request.model_path,
                engine::model_spec::default_spec_path(family()));
            return true;
        } catch (...) {
            return false;
        }
    }

    runtime::ModelInspection inspect(const runtime::ModelLoadRequest & request) const override {
        const auto assets = load_request_assets(request);
        runtime::ModelInspection inspection;
        inspection.metadata = metadata(*assets);
        inspection.capabilities = capabilities(*assets);
        inspection.cli = cli(*assets);
        inspection.model_root = assets->resources.model_root();
        const auto spec_path = engine::model_spec::default_spec_path(family());
        inspection.discovered_configs = runtime::discover_named_assets_from_package_spec(
            request.model_path,
            spec_path,
            engine::model_spec::ResourceKind::Files);
        inspection.discovered_weights = runtime::discover_named_assets_from_package_spec(
            request.model_path,
            spec_path,
            engine::model_spec::ResourceKind::Tensors);
        return inspection;
    }

    std::unique_ptr<runtime::ILoadedVoiceModel> load(const runtime::ModelLoadRequest & request) const override {
        return load_pocket_tts_model(request);
    }
};

}  // namespace

PocketTTSModel::PocketTTSModel(
    std::filesystem::path model_dir,
    std::shared_ptr<const PocketTTSAssets> manifest)
    : model_dir_(std::move(model_dir)),
      manifest_(std::move(manifest)) {}

PocketTTSModel PocketTTSModel::load(const ModelConfig & config) {
    if (config.model_dir.empty()) {
        throw std::runtime_error("PocketTTS model_dir is required");
    }
    runtime::ModelLoadRequest request;
    request.model_path = config.model_dir;
    if (!config.language.empty()) {
        request.options.emplace("language", config.language);
    }
    auto manifest = load_request_assets(request);
    auto voice_asset_root = manifest->voice_asset_root;
    return PocketTTSModel(std::move(voice_asset_root), std::move(manifest));
}

std::unique_ptr<runtime::IVoiceTaskSession> PocketTTSModel::create_task_session(
    const runtime::TaskSpec & task,
    const runtime::SessionOptions & options) const {
    if (!manifest_) {
        throw std::runtime_error("PocketTTS model is not loaded");
    }
    return std::make_unique<PocketTTSSession>(task, options, manifest_, model_dir_);
}

const std::filesystem::path & PocketTTSModel::model_dir() const noexcept {
    return model_dir_;
}

PocketTTSLoadedModel::PocketTTSLoadedModel(
    runtime::ModelMetadata metadata,
    runtime::CapabilitySet capabilities,
    PocketTTSModel model)
    : metadata_(std::move(metadata)),
      capabilities_(std::move(capabilities)),
      model_(std::move(model)) {}

const runtime::ModelMetadata & PocketTTSLoadedModel::metadata() const noexcept {
    return metadata_;
}

const runtime::CapabilitySet & PocketTTSLoadedModel::capabilities() const noexcept {
    return capabilities_;
}

std::unique_ptr<runtime::IVoiceTaskSession> PocketTTSLoadedModel::create_task_session(
    const runtime::TaskSpec & task,
    const runtime::SessionOptions & options) const {
    return model_.create_task_session(task, options);
}

std::unique_ptr<PocketTTSLoadedModel> load_pocket_tts_model(const runtime::ModelLoadRequest & request) {
    auto assets = load_request_assets(request);
    PocketTTSModel model(assets->voice_asset_root, assets);

    return std::make_unique<PocketTTSLoadedModel>(
        metadata(*assets),
        capabilities(*assets),
        std::move(model));
}

std::shared_ptr<runtime::IVoiceModelLoader> make_pocket_tts_loader() {
    return std::make_shared<PocketTTSLoader>();
}

}  // namespace engine::models::pocket_tts