// SPDX-License-Identifier: Apache-2.0
// Copyright (C) 2025 Advanced Micro Devices, Inc. All rights reserved.

// ------ I N C L U D E   F I L E S -------------------------------------------
// Local - Include Files
#include "core/common/query_requests.h"
#include "core/common/time.h"
#include "core/common/module_loader.h"
#include "tools/common/SmiWatchMode.h"
#include "tools/common/Table2D.h"
#include "tools/common/XBUtilities.h"
#include "EventTraceConfig.h"
#include "ReportEventTrace.h"

// 3rd Party Library - Include Files
#include <algorithm>
#include <boost/format.hpp>
#include <filesystem>
#include <iomanip>
#include <map>
#include <sstream>
#include <vector>
#include <cstring>
#include "core/common/json/nlohmann/json.hpp"

using bpt = boost::property_tree::ptree;
namespace XBU = XBUtilities;
namespace smi = xrt_core::tools::xrt_smi;

namespace xrt_core::tools::xrt_smi {

// Event trace data structure (must match device.cpp implementation)
struct trace_event {
  uint64_t timestamp;    // Simulated timestamp
  uint16_t event_id;     // Event ID from trace_events.h
  uint64_t payload;      // Event payload/arguments
};

// Function to generate simple dummy event trace data for testing
static void 
generate_dummy_event_trace_data(xrt_core::query::firmware_debug_buffer& log_buffer) 
{
  static uint64_t counter = 0;  // Ever-increasing counter
  counter++;
  
  // Generate simple trace events using real event IDs from trace_events.json
  const size_t num_events = 5; // NOLINT(cppcoreguidelines-avoid-magic-numbers) - dummy data for pretty printing
  auto events = static_cast<trace_event*>(log_buffer.data);
  
  // Base timestamp that increases with each call
  uint64_t base_timestamp = 1000000000ULL + (counter * 10000); // NOLINT(cppcoreguidelines-avoid-magic-numbers) - dummy data for pretty printing
  
  // Use real event IDs that exist in trace_events.json and can be parsed by config
  const std::vector<uint16_t> event_ids{0, 2, 10, 14, 1}; // NOLINT(cppcoreguidelines-avoid-magic-numbers) - dummy data for pretty printing
  const std::vector<uint64_t> payloads{
    (1ULL << 0) | (6ULL << 4),    // PROCESS_APP_MSG_START: context_id=1, msg_opcode=6 // NOLINT(cppcoreguidelines-avoid-magic-numbers) - dummy data for pretty printing
    (1ULL << 0) | (5ULL << 4),    // CREATE_CONTEXT: context_id=1, priority=5 // NOLINT(cppcoreguidelines-avoid-magic-numbers) - dummy data for pretty printing
    (0ULL << 0) | (4ULL << 8),    // GATE_AIE2_CLKS: start_col=0, num_cols=4 // NOLINT(cppcoreguidelines-avoid-magic-numbers) - dummy data for pretty printing
    0ULL,                         // SUSPEND: no_args
    (1ULL << 0) | (6ULL << 4)     // PROCESS_APP_MSG_DONE: context_id=1, msg_opcode=6 // NOLINT(cppcoreguidelines-avoid-magic-numbers) - dummy data for pretty printing
  };
  
  // Create event sequence with ever-increasing data
  for (size_t i = 0; i < num_events; ++i) {
    events[i].timestamp = base_timestamp + (i * 100); // NOLINT(cppcoreguidelines-avoid-magic-numbers) - dummy data for pretty printing
    events[i].event_id = event_ids[i];
    // Add counter to payload to make it ever-increasing while keeping structure
    events[i].payload = payloads[i] + (counter * 0x100); // NOLINT(cppcoreguidelines-avoid-magic-numbers) - dummy data for pretty printing
  }
  
  log_buffer.size = num_events * sizeof(trace_event);
  log_buffer.abs_offset = counter * log_buffer.size;  // Ever-increasing offset
}

static xrt_core::tools::xrt_smi::event_trace_config 
get_event_trace_config(const xrt_core::device* dev) {
  auto archive = XBU::open_archive(dev);
  auto artifacts_repo = XBU::extract_artifacts_from_archive(archive.get(), {"trace_events.json"});
  
  auto& config_data = artifacts_repo["trace_events.json"];
  std::string config_content(config_data.begin(), config_data.end());
  
  nlohmann::json json_config = nlohmann::json::parse(config_content);
  return xrt_core::tools::xrt_smi::event_trace_config(json_config);
}

static void
validate_version_compatibility(const std::pair<uint16_t, uint16_t>& /*version*/,
                               const xrt_core::device* device) 
{
  if (!device) {
    throw std::runtime_error("Warning: Cannot validate event trace version - no device provided");
  }

  /*
  TODO : Add version logic based on what driver provides
  auto firmware_version = xrt_core::device_query<xrt_core::query::event_trace_version>(device);
  if (version.first != firmware_version.major || version.second != firmware_version.minor) {
    std::stringstream err{};
    err << "Warning: Event trace version mismatch!\n"
        << "  JSON file version: " << version.first << "." << version.second << "\n"
        << "  Firmware version: " << firmware_version.major << "." << firmware_version.minor << "\n"
        << "  Event parsing may be incorrect or incomplete.";
    throw std::runtime_error(err.str());
  }
  */
}

static std::string
generate_raw_logs(const xrt_core::device* dev,
                  bool is_watch,
                  uint64_t& watch_mode_offset)
{
  std::stringstream ss{};
  
  smi_debug_buffer debug_buf(watch_mode_offset, is_watch);

  // Always use real device data for raw logs
  xrt_core::device_query<xrt_core::query::event_trace_data>(dev, debug_buf.get_log_buffer());
  
  watch_mode_offset = debug_buf.get_offset();
  
  ss << boost::format("Event Trace Report (Raw) - %s\n") 
        % xrt_core::timestamp();
  ss << "=======================================================\n";
  ss << "\n";

  if (!debug_buf.get_data() || debug_buf.get_size() == 0) {
    ss << "No event trace data available\n";
    return ss.str();
  }

  // Simply print the raw payload data
  ss.write(static_cast<const char*>(debug_buf.get_data()), static_cast<std::streamsize>(debug_buf.get_size()));
  
  return ss.str();
}

static std::string
generate_event_trace_report(const xrt_core::device* dev,
                            const event_trace_config& config,
                            bool is_watch,
                            bool use_dummy = false)
{
  std::stringstream ss{};
  static uint64_t watch_mode_offset = 0;
  
  try {
    auto version = config.get_file_version();
    validate_version_compatibility(version, dev);

    smi_debug_buffer debug_buf(watch_mode_offset, is_watch);

    if (use_dummy) 
    {
      generate_dummy_event_trace_data(debug_buf.get_log_buffer());
    } 
    else 
    {
      xrt_core::device_query<xrt_core::query::event_trace_data>(dev, debug_buf.get_log_buffer());
    }
    watch_mode_offset = debug_buf.get_offset();
    
    ss << boost::format("Event Trace Report (Buffer: %d bytes) - %s\n") 
          % debug_buf.get_size() % xrt_core::timestamp();
    ss << "=======================================================\n";

    ss << "\n";

    if (!debug_buf.get_data() || debug_buf.get_size() == 0) {
      ss << "No event trace data available\n";
      return ss.str();
    }

    // Parse trace events from buffer
    size_t event_count = debug_buf.get_size() / sizeof(trace_event);
    auto events = static_cast<trace_event*>(debug_buf.get_data());

    ss << boost::format("Total Events: %d\n") % event_count;
    ss << boost::format("Buffer Offset: %d\n\n") % debug_buf.get_offset();

    // Create Table2D with headers (enhanced with parsed args)
    const std::vector<Table2D::HeaderData> table_headers = {
      {"Timestamp",         Table2D::Justification::right},
      {"Event ID",          Table2D::Justification::left},
      {"Event Name",        Table2D::Justification::left},
      {"Category",          Table2D::Justification::left},
      {"Payload",           Table2D::Justification::left},
      {"Parsed Args",       Table2D::Justification::left}
    };
    Table2D event_table(table_headers);
    
    // Add data rows
    for (size_t i = 0; i < event_count; ++i) {
      // Parse event using JSON-based configuration
      xrt_core::tools::xrt_smi::event_record record{events[i].timestamp, events[i].event_id, events[i].payload};
      auto parsed_event = config.parse_event(record);

      // Join categories with pipe separator for backward compatibility
      std::string categories_str{};
      for (size_t j = 0; j < parsed_event.categories.size(); ++j) {
        if (j > 0) categories_str += "|";
        categories_str += parsed_event.categories[j];
      }
      
      // Format parsed arguments
      std::string args_str{};
      for (const auto& arg_pair : parsed_event.args) {
        if (!args_str.empty()) args_str += ", ";
        args_str += arg_pair.first + "=" + arg_pair.second;
      }
      
      const std::vector<std::string> entry_data = {
        std::to_string(parsed_event.timestamp),
        (boost::format("0x%04x") % parsed_event.event_id).str(),
        parsed_event.name,
        categories_str,
        (boost::format("0x%x") % parsed_event.raw_payload).str(),
        args_str
      };
      event_table.addEntry(entry_data);
    }

    ss << "  Event Trace Information:\n";
    ss << event_table.toString("    ");

  } 
  catch (const std::exception& e) {
    ss << "Error parsing event trace data: " << e.what() << "\n";
  }
  return ss.str();
}

} // namespace xrt_core::tools::xrt_smi


void
ReportEventTrace::
getPropertyTreeInternal(const xrt_core::device* dev, bpt& pt) const
{
  // Defer to the 20202 format.  If we ever need to update JSON data,
  // Then update this method to do so.
  getPropertyTree20202(dev, pt);
}

void
ReportEventTrace::
getPropertyTree20202(const xrt_core::device* dev, bpt& pt) const
{
  bpt event_trace_pt{};

  try {
    // Get the event trace configuration
    auto config = smi::get_event_trace_config(dev);
    
    smi_debug_buffer debug_buf(0, false);
    
    // Query event trace data from device using specific query struct (like telemetry)
    xrt_core::device_query<xrt_core::query::event_trace_data>(dev, debug_buf.get_log_buffer());

    // Parse trace events from buffer
    if (debug_buf.get_data() && debug_buf.get_size() > 0) {
      size_t event_count = debug_buf.get_size() / sizeof(smi::trace_event);
      auto events = static_cast<smi::trace_event*>(debug_buf.get_data());

      bpt events_array{};
      for (size_t i = 0; i < event_count; ++i) {
        // Parse event using json based configuration
        xrt_core::tools::xrt_smi::event_record record{events[i].timestamp, 
                            events[i].event_id, 
                            events[i].payload};

        auto parsed_event = config.parse_event(record);

        bpt event_pt;
        event_pt.put("timestamp", parsed_event.timestamp);
        event_pt.put("event_id", parsed_event.event_id);
        event_pt.put("event_name", parsed_event.name);
        
        // Join categories with pipe separator for backward compatibility
        std::string categories_str{};
        for (size_t j = 0; j < parsed_event.categories.size(); ++j) {
          if (j > 0) categories_str += "|";
          categories_str += parsed_event.categories[j];
        }
        event_pt.put("category", categories_str);
        
        event_pt.put("payload", parsed_event.raw_payload);
        
        // Add parsed arguments
        bpt args_pt;
        for (const auto& arg_pair : parsed_event.args) {
          args_pt.put(arg_pair.first, arg_pair.second);
        }
        if (!parsed_event.args.empty()) {
          event_pt.add_child("args", args_pt);
        }
        events_array.push_back(std::make_pair("", event_pt));
      }
      event_trace_pt.add_child("events", events_array);
      event_trace_pt.put("event_count", event_count);
      event_trace_pt.put("buffer_offset", debug_buf.get_offset());
      event_trace_pt.put("buffer_size", debug_buf.get_size());
    } else {
      event_trace_pt.put("event_count", 0);
      event_trace_pt.put("buffer_offset", 0);
      event_trace_pt.put("buffer_size", 0);
    }
  } 
  catch (const std::exception& e) {
    event_trace_pt.put("event_count", 0);
    event_trace_pt.put("error", e.what());
  }

  // There can only be 1 root node
  pt.add_child("event_trace", event_trace_pt);
}


void
ReportEventTrace::
writeReport(const xrt_core::device* device,
            const bpt& /*pt*/,
            const std::vector<std::string>& elements_filter,
            std::ostream& output) const
{
  // Check for dummy option
  bool use_dummy = std::find(elements_filter.begin(), elements_filter.end(), "dummy") != elements_filter.end();
  
  try {
    // Get the event trace configuration
    auto config = smi::get_event_trace_config(device);
    
    // Check for watch mode
    if (smi_watch_mode::parse_watch_mode_options(elements_filter)) {
      // Create report generator lambda for watch mode
      auto report_generator = [config, use_dummy](const xrt_core::device* dev) -> std::string {
        return smi::generate_event_trace_report(dev, config, true, use_dummy);
      };

      smi_watch_mode::run_watch_mode(device, output, report_generator, "Event Trace");
      return;
    }
    output << "Event Trace Report\n";
    output << "==================\n\n";
    output << smi::generate_event_trace_report(device, config, false, use_dummy);
    output << std::endl;
  }
  catch (const std::exception& e) {
    output << "Error loading event trace config: " << e.what() << "\n";
    output << "Falling back to raw event trace data:\n\n";
    
    static uint64_t watch_mode_offset = 0;
    if (smi_watch_mode::parse_watch_mode_options(elements_filter)) {
      // Create report generator lambda for watch mode with raw logs
      auto report_generator = [](const xrt_core::device* dev) -> std::string {
        return smi::generate_raw_logs(dev, true, watch_mode_offset);
      };

      smi_watch_mode::run_watch_mode(device, output, report_generator, "Event Trace (Raw)");
      return;
    }
    output << smi::generate_raw_logs(device, false, watch_mode_offset);
    output << std::endl;
  }
}
