// MIT License
//
// Copyright (c) 2025 Advanced Micro Devices, Inc. All Rights Reserved.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

#include "library/rocprofiler-sdk/counters.hpp"
#include "core/agent_manager.hpp"
#include "core/demangler.hpp"
#include "core/trace_cache/cache_manager.hpp"
#include "core/trace_cache/metadata_registry.hpp"
#include "library/rocprofiler-sdk/fwd.hpp"

#include <memory>
#include <timemory/utility/types.hpp>

#include "logger/debug.hpp"

namespace rocprofsys
{
namespace rocprofiler_sdk
{
namespace
{
void
metadata_initialize_counter_category()
{
    trace_cache::get_metadata_registry().add_string(
        trait::name<category::rocm_counter_collection>::value);
}

void
metadata_initialize_counter_track(const char* name)
{
    trace_cache::get_metadata_registry().add_track({ name, std::nullopt, "{}" });
}

void
metadata_initialize_counters_pmc(size_t dev_id, const std::string& name,
                                 const std::string& metric_description)
{
    const size_t EVENT_CODE       = 0;
    const size_t INSTANCE_ID      = 0;
    const char*  LONG_DESCRIPTION = "";
    const char*  COMPONENT        = "";
    const char*  BLOCK            = "";
    const char*  EXPRESSION       = "";
    const auto*  TARGET_ARCH      = "GPU";

    trace_cache::get_metadata_registry().add_pmc_info(
        { agent_type::GPU, dev_id, TARGET_ARCH, EVENT_CODE, INSTANCE_ID, name.c_str(),
          name.c_str(), metric_description.c_str(), LONG_DESCRIPTION, COMPONENT,
          "Unit Count", rocprofsys::trace_cache::ABSOLUTE, BLOCK, EXPRESSION, 0, 0 });
}
}  // namespace
namespace
{
std::string
get_counter_description(const client_data* tool_data, std::string_view _v)
{
    const auto& _info = tool_data->events_info;
    for(const auto& itr : _info)
    {
        if(itr.symbol().find(_v) == 0 || itr.short_description().find(_v) == 0)
        {
            return itr.long_description();
        }
    }
    return std::string{};
}
}  // namespace

void
counter_event::operator()(const client_data* tool_data, ::perfetto::CounterTrack* _track,
                          const std::string& track_name, timing_interval _timing,
                          scope::config _scope) const
{
    if(!record.dispatch_data) return;

    const auto& _dispatch_info = record.dispatch_data->dispatch_info;
    const auto* _kern_sym_data =
        tool_data->get_kernel_symbol_info(_dispatch_info.kernel_id);

    auto _bundle =
        counter_bundle_t{ rocprofsys::utility::demangle(_kern_sym_data->kernel_name),
                          _scope };

    _bundle.push(_dispatch_info.queue_id.handle)
        .start()
        .store(record.record_counter.counter_value);

    _bundle.stop().pop(_dispatch_info.queue_id.handle);

    if(_track && _timing.start > 0 && _timing.end > _timing.start)
    {
        TRACE_COUNTER(trait::name<category::rocm_counter_collection>::value, *_track,
                      _timing.start, record.record_counter.counter_value);
        TRACE_COUNTER(trait::name<category::rocm_counter_collection>::value, *_track,
                      _timing.end, 0);

        const std::string event_metadata  = "{}";
        const size_t      stack_id        = 0;
        const size_t      parent_stack_id = 0;
        const size_t      correlation_id  = 0;
        const std::string call_stack      = "{}";
        const std::string line_info       = "{}";
        const size_t      agent_handle    = record.record_counter.agent_id.handle;
        const size_t      value           = record.record_counter.counter_value;

        auto agent = get_agent_manager_instance().get_agent_by_handle(agent_handle);

        trace_cache::get_buffer_storage().store(trace_cache::pmc_event_with_sample{
            static_cast<size_t>(
                category_enum_id<category::rocm_counter_collection>::value),
            track_name.c_str(), _timing.start, event_metadata.c_str(), stack_id,
            parent_stack_id, correlation_id, call_stack.c_str(), line_info.c_str(),
            static_cast<uint32_t>(agent.device_id), static_cast<uint8_t>(agent.type),
            track_name.c_str(), static_cast<double>(value) });
    }
}

counter_storage::counter_storage(const client_data* _tool_data, uint64_t _devid,
                                 size_t _idx, std::string_view _name)
: tool_data{ _tool_data }
, device_id{ _devid }
, index{ static_cast<int64_t>(_idx) }
, metric_name{ _name }
, metric_description{ get_counter_description(_tool_data, metric_name) }
{
    auto _metric_name = std::string{ _name };
    _metric_name =
        std::regex_replace(_metric_name, std::regex{ "(.*)\\[([0-9]+)\\]" }, "$1_$2");
    storage_name = fmt::format("rocprof-device-{}-{}", device_id, _metric_name);
    storage = std::make_unique<counter_storage_type>(tim::standalone_storage{}, index,
                                                     storage_name);
    tim::manager::instance()->add_cleanup(
        storage_name + "cleanup", [storage_ptr = storage.get(), metric_name = metric_name,
                                   metric_description = metric_description]() {
            if(storage_ptr)
                counter_storage::write(storage_ptr, metric_name, metric_description);
        });

    {
        constexpr auto _unit = ::perfetto::CounterTrack::Unit::UNIT_COUNT;
        track_name           = fmt::format(" GPU {} [{}]", _metric_name, device_id);
        track                = std::make_unique<counter_track_type>(
            ::perfetto::StaticString(track_name.c_str()));

        metadata_initialize_counter_category();
        metadata_initialize_counters_pmc(device_id, track_name.c_str(),
                                         metric_description);
        metadata_initialize_counter_track(track_name.c_str());
        track->set_is_incremental(false);
        track->set_unit(_unit);
        track->set_unit_multiplier(1);
    }
}

void
counter_storage::operator()(const counter_event& _event, timing_interval _timing,
                            scope::config _scope) const
{
    operation::set_storage<counter_data_tracker>{}(storage.get());
    _event(tool_data, track.get(), track_name, _timing, _scope);
}

void
counter_storage::write(counter_storage_type* storage, const std::string& metric_name,
                       const std::string& metric_description)
{
    if(!trait::runtime_enabled<counter_data_tracker>::get())
    {
        LOG_WARNING("{} counter_data_tracker is disabled. Can't write storage.",
                    metric_name);
        return;
    }

    operation::set_storage<counter_data_tracker>{}(storage);
    counter_data_tracker::label()       = metric_name;
    counter_data_tracker::description() = metric_description;
    storage->write();
}
}  // namespace rocprofiler_sdk
}  // namespace rocprofsys