#pragma once

#include "api_server.h"
#ifdef USE_API
#ifdef USE_API_HOMEASSISTANT_SERVICES
#include <functional>
#include <utility>
#include <vector>
#include "api_pb2.h"
#ifdef USE_API_HOMEASSISTANT_ACTION_RESPONSES_JSON
#include "esphome/components/json/json_util.h"
#endif
#include "esphome/core/automation.h"
#include "esphome/core/helpers.h"
#include "esphome/core/string_ref.h"

namespace esphome::api {

template<typename... X> class TemplatableStringValue : public TemplatableValue<std::string, X...> {
  // Verify that const char* uses the base class STATIC_STRING optimization (no heap allocation)
  // rather than being wrapped in a lambda. The base class constructor for const char* is more
  // specialized than the templated constructor here, so it should be selected.
  static_assert(std::is_constructible_v<TemplatableValue<std::string, X...>, const char *>,
                "Base class must have const char* constructor for STATIC_STRING optimization");

 private:
  // Helper to convert value to string - handles the case where value is already a string
  template<typename T> static std::string value_to_string(T &&val) {
    return to_string(std::forward<T>(val));  // NOLINT
  }

  // Overloads for string types - needed because std::to_string doesn't support them
  static std::string value_to_string(char *val) {
    return val ? std::string(val) : std::string();
  }  // For lambdas returning char* (e.g., itoa)
  static std::string value_to_string(const char *val) { return std::string(val); }  // For lambdas returning .c_str()
  static std::string value_to_string(const std::string &val) { return val; }
  static std::string value_to_string(std::string &&val) { return std::move(val); }
  static std::string value_to_string(const StringRef &val) { return val.str(); }
  static std::string value_to_string(StringRef &&val) { return val.str(); }

 public:
  TemplatableStringValue() : TemplatableValue<std::string, X...>() {}

  template<typename F, enable_if_t<!is_invocable<F, X...>::value, int> = 0>
  TemplatableStringValue(F value) : TemplatableValue<std::string, X...>(value) {}

  template<typename F, enable_if_t<is_invocable<F, X...>::value, int> = 0>
  TemplatableStringValue(F f)
      : TemplatableValue<std::string, X...>([f](X... x) -> std::string { return value_to_string(f(x...)); }) {}
};

template<typename... Ts> class TemplatableKeyValuePair {
 public:
  // Default constructor needed for FixedVector::emplace_back()
  TemplatableKeyValuePair() = default;

  // Keys are always string literals from YAML dictionary keys (e.g., "code", "event")
  // and never templatable values or lambdas. Only the value parameter can be a lambda/template.
  // Using const char* avoids std::string heap allocation - keys remain in flash.
  template<typename T> TemplatableKeyValuePair(const char *key, T value) : key(key), value(value) {}

  const char *key{nullptr};
  TemplatableStringValue<Ts...> value;
};

#ifdef USE_API_HOMEASSISTANT_ACTION_RESPONSES
// Represents the response data from a Home Assistant action
// Note: This class holds a StringRef to the error_message from the protobuf message.
// The protobuf message must outlive the ActionResponse (which is guaranteed since
// the callback is invoked synchronously while the message is on the stack).
class ActionResponse {
 public:
  ActionResponse(bool success, StringRef error_message) : success_(success), error_message_(error_message) {}

#ifdef USE_API_HOMEASSISTANT_ACTION_RESPONSES_JSON
  ActionResponse(bool success, StringRef error_message, const uint8_t *data, size_t data_len)
      : success_(success), error_message_(error_message) {
    if (data == nullptr || data_len == 0)
      return;
    JsonDocument tmp = json::parse_json(data, data_len);
    swap(this->json_document_, tmp);
  }
#endif

  bool is_success() const { return this->success_; }
  // Returns reference to error message - can be implicitly converted to std::string if needed
  const StringRef &get_error_message() const { return this->error_message_; }

#ifdef USE_API_HOMEASSISTANT_ACTION_RESPONSES_JSON
  // Get data as parsed JSON object (const version returns read-only view)
  JsonObjectConst get_json() const { return this->json_document_.as<JsonObjectConst>(); }
#endif

 protected:
  bool success_;
  StringRef error_message_;
#ifdef USE_API_HOMEASSISTANT_ACTION_RESPONSES_JSON
  JsonDocument json_document_;
#endif
};

// Callback type for action responses
template<typename... Ts> using ActionResponseCallback = std::function<void(const ActionResponse &, Ts...)>;
#endif

template<typename... Ts> class HomeAssistantServiceCallAction : public Action<Ts...> {
 public:
  explicit HomeAssistantServiceCallAction(APIServer *parent, bool is_event) : parent_(parent) {
    this->flags_.is_event = is_event;
  }

  template<typename T> void set_service(T service) { this->service_ = service; }

  // Initialize FixedVector members - called from Python codegen with compile-time known sizes.
  // Must be called before any add_* methods; capacity must match the number of subsequent add_* calls.
  void init_data(size_t count) { this->data_.init(count); }
  void init_data_template(size_t count) { this->data_template_.init(count); }
  void init_variables(size_t count) { this->variables_.init(count); }

  // Keys are always string literals from the Python code generation (e.g., cg.add(var.add_data("tag_id", templ))).
  // The value parameter can be a lambda/template, but keys are never templatable.
  // Using const char* for keys avoids std::string heap allocation - keys remain in flash.
  template<typename V> void add_data(const char *key, V &&value) {
    this->add_kv_(this->data_, key, std::forward<V>(value));
  }
  template<typename V> void add_data_template(const char *key, V &&value) {
    this->add_kv_(this->data_template_, key, std::forward<V>(value));
  }
  template<typename V> void add_variable(const char *key, V &&value) {
    this->add_kv_(this->variables_, key, std::forward<V>(value));
  }

#ifdef USE_ESP8266
  // On ESP8266, ESPHOME_F() returns __FlashStringHelper* (PROGMEM pointer).
  // Store as const char* — populate_service_map copies from PROGMEM at play() time.
  template<typename V> void add_data(const __FlashStringHelper *key, V &&value) {
    this->add_kv_(this->data_, reinterpret_cast<const char *>(key), std::forward<V>(value));
  }
  template<typename V> void add_data_template(const __FlashStringHelper *key, V &&value) {
    this->add_kv_(this->data_template_, reinterpret_cast<const char *>(key), std::forward<V>(value));
  }
  template<typename V> void add_variable(const __FlashStringHelper *key, V &&value) {
    this->add_kv_(this->variables_, reinterpret_cast<const char *>(key), std::forward<V>(value));
  }
#endif

#ifdef USE_API_HOMEASSISTANT_ACTION_RESPONSES
  template<typename T> void set_response_template(T response_template) {
    this->response_template_ = response_template;
    this->flags_.has_response_template = true;
  }

  void set_wants_status() { this->flags_.wants_status = true; }
  void set_wants_response() { this->flags_.wants_response = true; }

#ifdef USE_API_HOMEASSISTANT_ACTION_RESPONSES_JSON
  Trigger<JsonObjectConst, Ts...> *get_success_trigger_with_response() { return &this->success_trigger_with_response_; }
#endif
  Trigger<Ts...> *get_success_trigger() { return &this->success_trigger_; }
  Trigger<std::string, Ts...> *get_error_trigger() { return &this->error_trigger_; }
#endif  // USE_API_HOMEASSISTANT_ACTION_RESPONSES

  void play(const Ts &...x) override {
    HomeassistantActionRequest resp;
    std::string service_value = this->service_.value(x...);
    resp.service = StringRef(service_value);
    resp.is_event = this->flags_.is_event;

    // Local storage for lambda-evaluated strings - lives until after send
    FixedVector<std::string> data_storage;
    FixedVector<std::string> data_template_storage;
    FixedVector<std::string> variables_storage;

    this->populate_service_map(resp.data, this->data_, data_storage, x...);
    this->populate_service_map(resp.data_template, this->data_template_, data_template_storage, x...);
    this->populate_service_map(resp.variables, this->variables_, variables_storage, x...);

#ifdef USE_API_HOMEASSISTANT_ACTION_RESPONSES
#ifdef USE_API_HOMEASSISTANT_ACTION_RESPONSES_JSON
    // IMPORTANT: Declare at outer scope so it lives until send_homeassistant_action returns.
    std::string response_template_value;
#endif
    if (this->flags_.wants_status) {
      // Generate a unique call ID for this service call
      static uint32_t call_id_counter = 1;
      uint32_t call_id = call_id_counter++;
      resp.call_id = call_id;
#ifdef USE_API_HOMEASSISTANT_ACTION_RESPONSES_JSON
      if (this->flags_.wants_response) {
        resp.wants_response = true;
        // Set response template if provided
        if (this->flags_.has_response_template) {
          response_template_value = this->response_template_.value(x...);
          resp.response_template = StringRef(response_template_value);
        }
      }
#endif

      auto captured_args = std::make_tuple(x...);
      this->parent_->register_action_response_callback(call_id, [this, captured_args](const ActionResponse &response) {
        std::apply(
            [this, &response](auto &&...args) {
              if (response.is_success()) {
#ifdef USE_API_HOMEASSISTANT_ACTION_RESPONSES_JSON
                if (this->flags_.wants_response) {
                  this->success_trigger_with_response_.trigger(response.get_json(), args...);
                } else
#endif
                {
                  this->success_trigger_.trigger(args...);
                }
              } else {
                this->error_trigger_.trigger(response.get_error_message(), args...);
              }
            },
            captured_args);
      });
    }
#endif

    this->parent_->send_homeassistant_action(resp);
  }

 protected:
  // Helper to add key-value pairs to FixedVectors
  // Keys are always string literals (const char*), values can be lambdas/templates
  template<typename V> void add_kv_(FixedVector<TemplatableKeyValuePair<Ts...>> &vec, const char *key, V &&value) {
    auto &kv = vec.emplace_back();
    kv.key = key;
    kv.value = std::forward<V>(value);
  }

  template<typename VectorType, typename SourceType>
  static void populate_service_map(VectorType &dest, SourceType &source, FixedVector<std::string> &value_storage,
                                   Ts... x) {
    dest.init(source.size());

#ifdef USE_ESP8266
    // On ESP8266, all static strings from codegen are FLASH_STRING (PROGMEM),
    // so is_static_string() is always false — the zero-copy STATIC_STRING fast
    // path from the non-ESP8266 branch cannot trigger. We copy all keys and
    // values unconditionally: keys via _P functions (may be in PROGMEM), values
    // via value() which handles FLASH_STRING internally.
    value_storage.init(source.size() * 2);

    for (auto &it : source) {
      auto &kv = dest.emplace_back();

      // Key: copy from possible PROGMEM
      {
        size_t key_len = strlen_P(it.key);
        value_storage.push_back(std::string(key_len, '\0'));
        memcpy_P(value_storage.back().data(), it.key, key_len);
        kv.key = StringRef(value_storage.back());
      }

      // Value: value() handles FLASH_STRING via _P functions internally
      value_storage.push_back(it.value.value(x...));
      kv.value = StringRef(value_storage.back());
    }
#else
    // On non-ESP8266, strings are directly readable from flash-mapped memory.
    // Count non-static strings to allocate exact storage needed.
    size_t lambda_count = 0;
    for (const auto &it : source) {
      if (!it.value.is_static_string()) {
        lambda_count++;
      }
    }
    value_storage.init(lambda_count);

    for (auto &it : source) {
      auto &kv = dest.emplace_back();
      kv.key = StringRef(it.key);

      if (it.value.is_static_string()) {
        // Static string — pointer directly readable, zero allocation
        kv.value = StringRef(it.value.get_static_string());
      } else {
        // Lambda — evaluate and store result
        value_storage.push_back(it.value.value(x...));
        kv.value = StringRef(value_storage.back());
      }
    }
#endif
  }

  APIServer *parent_;
  TemplatableStringValue<Ts...> service_{};
  FixedVector<TemplatableKeyValuePair<Ts...>> data_;
  FixedVector<TemplatableKeyValuePair<Ts...>> data_template_;
  FixedVector<TemplatableKeyValuePair<Ts...>> variables_;
#ifdef USE_API_HOMEASSISTANT_ACTION_RESPONSES
#ifdef USE_API_HOMEASSISTANT_ACTION_RESPONSES_JSON
  TemplatableStringValue<Ts...> response_template_{""};
  Trigger<JsonObjectConst, Ts...> success_trigger_with_response_;
#endif  // USE_API_HOMEASSISTANT_ACTION_RESPONSES_JSON
  Trigger<Ts...> success_trigger_;
  Trigger<std::string, Ts...> error_trigger_;
#endif  // USE_API_HOMEASSISTANT_ACTION_RESPONSES

  struct Flags {
    uint8_t is_event : 1;
    uint8_t wants_status : 1;
    uint8_t wants_response : 1;
    uint8_t has_response_template : 1;
    uint8_t reserved : 5;
  } flags_{0};
};

}  // namespace esphome::api

#endif
#endif
