#include "light_state.h"
#include "esp_color_correction.h"
#include "esphome/core/defines.h"
#include "esphome/core/controller_registry.h"
#include "esphome/core/log.h"
#include "light_output.h"
#include "transformers.h"

namespace esphome::light {

static const char *const TAG = "light";

LightState::LightState(LightOutput *output) : output_(output) {}

LightTraits LightState::get_traits() { return this->output_->get_traits(); }
LightCall LightState::turn_on() { return this->make_call().set_state(true); }
LightCall LightState::turn_off() { return this->make_call().set_state(false); }
LightCall LightState::toggle() { return this->make_call().set_state(!this->remote_values.is_on()); }
LightCall LightState::make_call() { return LightCall(this); }

void LightState::setup() {
  this->output_->setup_state(this);
  for (auto *effect : this->effects_) {
    effect->init_internal(this);
  }

  // Start with loop disabled if idle - respects any effects/transitions set up during initialization
  this->disable_loop_if_idle_();

  // When supported color temperature range is known, initialize color temperature setting within bounds.
  auto traits = this->get_traits();
  float min_mireds = traits.get_min_mireds();
  if (min_mireds > 0) {
    this->remote_values.set_color_temperature(min_mireds);
    this->current_values.set_color_temperature(min_mireds);
  }

  auto call = this->make_call();
  LightStateRTCState recovered{};
  if (this->initial_state_callback_) {
    this->initial_state_callback_(recovered);
    this->initial_state_callback_ = nullptr;  // One-shot — no longer needed
  }
  switch (this->restore_mode_) {
    case LIGHT_RESTORE_DEFAULT_OFF:
    case LIGHT_RESTORE_DEFAULT_ON:
    case LIGHT_RESTORE_INVERTED_DEFAULT_OFF:
    case LIGHT_RESTORE_INVERTED_DEFAULT_ON:
      this->rtc_ = this->make_entity_preference<LightStateRTCState>();
      // Attempt to load from preferences, else fall back to default values
      if (!this->rtc_.load(&recovered)) {
        recovered.state = (this->restore_mode_ == LIGHT_RESTORE_DEFAULT_ON ||
                           this->restore_mode_ == LIGHT_RESTORE_INVERTED_DEFAULT_ON);
      } else if (this->restore_mode_ == LIGHT_RESTORE_INVERTED_DEFAULT_OFF ||
                 this->restore_mode_ == LIGHT_RESTORE_INVERTED_DEFAULT_ON) {
        // Inverted restore state
        recovered.state = !recovered.state;
      }
      break;
    case LIGHT_RESTORE_AND_OFF:
    case LIGHT_RESTORE_AND_ON:
      this->rtc_ = this->make_entity_preference<LightStateRTCState>();
      this->rtc_.load(&recovered);
      recovered.state = (this->restore_mode_ == LIGHT_RESTORE_AND_ON);
      break;
    case LIGHT_ALWAYS_OFF:
      recovered.state = false;
      break;
    case LIGHT_ALWAYS_ON:
      recovered.state = true;
      break;
  }

  call.set_color_mode_if_supported(recovered.color_mode);
  call.set_state(recovered.state);
  call.set_brightness_if_supported(recovered.brightness);
  call.set_color_brightness_if_supported(recovered.color_brightness);
  call.set_red_if_supported(recovered.red);
  call.set_green_if_supported(recovered.green);
  call.set_blue_if_supported(recovered.blue);
  call.set_white_if_supported(recovered.white);
  call.set_color_temperature_if_supported(recovered.color_temp);
  call.set_cold_white_if_supported(recovered.cold_white);
  call.set_warm_white_if_supported(recovered.warm_white);
  if (recovered.effect != 0) {
    call.set_effect(recovered.effect);
  } else {
    call.set_transition_length_if_supported(0);
  }
  call.perform();
}
void LightState::dump_config() {
  ESP_LOGCONFIG(TAG, "Light '%s'", this->get_name().c_str());
  auto traits = this->get_traits();
  if (traits.supports_color_capability(ColorCapability::BRIGHTNESS)) {
    ESP_LOGCONFIG(TAG,
                  "  Default Transition Length: %.1fs\n"
                  "  Gamma Correct: %.2f",
                  this->default_transition_length_ / 1e3f, this->gamma_correct_);
  }
  if (traits.supports_color_capability(ColorCapability::COLOR_TEMPERATURE)) {
    ESP_LOGCONFIG(TAG,
                  "  Min Mireds: %.1f\n"
                  "  Max Mireds: %.1f",
                  traits.get_min_mireds(), traits.get_max_mireds());
  }
}
void LightState::loop() {
  // Apply effect (if any)
  auto *effect = this->get_active_effect_();
  if (effect != nullptr) {
    effect->apply();
  }

  // Apply transformer (if any)
  if (this->transformer_ != nullptr) {
    auto values = this->transformer_->apply();
    this->is_transformer_active_ = true;
    if (values.has_value()) {
      this->current_values = *values;
      this->output_->update_state(this);
      this->next_write_ = true;
    }

    if (this->transformer_->is_finished()) {
      // if the transition has written directly to the output, current_values is outdated, so update it
      this->current_values = this->transformer_->get_target_values();

      this->transformer_->stop();
      this->is_transformer_active_ = false;
      this->transformer_ = nullptr;
      if (this->target_state_reached_listeners_) {
        for (auto *listener : *this->target_state_reached_listeners_) {
          listener->on_light_target_state_reached();
        }
      }

      // Disable loop if idle (no transformer and no effect)
      this->disable_loop_if_idle_();
    }
  }

  // Write state to the light
  if (this->next_write_) {
    this->next_write_ = false;
    this->output_->write_state(this);
    // Disable loop if idle (no transformer and no effect)
    this->disable_loop_if_idle_();
  }
}

float LightState::get_setup_priority() const { return setup_priority::HARDWARE - 1.0f; }

void LightState::publish_state() {
  if (this->remote_values_listeners_) {
    for (auto *listener : *this->remote_values_listeners_) {
      listener->on_light_remote_values_update();
    }
  }
#if defined(USE_LIGHT) && defined(USE_CONTROLLER_REGISTRY)
  ControllerRegistry::notify_light_update(this);
#endif
}

LightOutput *LightState::get_output() const { return this->output_; }

static constexpr auto EFFECT_NONE_REF = StringRef::from_lit("None");

StringRef LightState::get_effect_name() {
  if (this->active_effect_index_ > 0) {
    return this->effects_[this->active_effect_index_ - 1]->get_name();
  }
  return EFFECT_NONE_REF;
}

void LightState::add_remote_values_listener(LightRemoteValuesListener *listener) {
  if (!this->remote_values_listeners_) {
    this->remote_values_listeners_ = make_unique<std::vector<LightRemoteValuesListener *>>();
  }
  this->remote_values_listeners_->push_back(listener);
}
void LightState::add_target_state_reached_listener(LightTargetStateReachedListener *listener) {
  if (!this->target_state_reached_listeners_) {
    this->target_state_reached_listeners_ = make_unique<std::vector<LightTargetStateReachedListener *>>();
  }
  this->target_state_reached_listeners_->push_back(listener);
}

void LightState::set_default_transition_length(uint32_t default_transition_length) {
  this->default_transition_length_ = default_transition_length;
}
uint32_t LightState::get_default_transition_length() const { return this->default_transition_length_; }
void LightState::set_flash_transition_length(uint32_t flash_transition_length) {
  this->flash_transition_length_ = flash_transition_length;
}
uint32_t LightState::get_flash_transition_length() const { return this->flash_transition_length_; }
void LightState::set_gamma_correct(float gamma_correct) { this->gamma_correct_ = gamma_correct; }
void LightState::set_restore_mode(LightRestoreMode restore_mode) { this->restore_mode_ = restore_mode; }
void LightState::set_initial_state(void (*callback)(LightStateRTCState &)) { this->initial_state_callback_ = callback; }
bool LightState::supports_effects() { return !this->effects_.empty(); }
const FixedVector<LightEffect *> &LightState::get_effects() const { return this->effects_; }
void LightState::add_effects(const std::initializer_list<LightEffect *> &effects) {
  // Called once from Python codegen during setup with all effects from YAML config
  this->effects_ = effects;
}

void LightState::current_values_as_binary(bool *binary) { this->current_values.as_binary(binary); }
void LightState::current_values_as_brightness(float *brightness) {
  this->current_values.as_brightness(brightness);
  *brightness = this->gamma_correct_lut(*brightness);
}
void LightState::current_values_as_rgb(float *red, float *green, float *blue) {
  this->current_values.as_rgb(red, green, blue);
  *red = this->gamma_correct_lut(*red);
  *green = this->gamma_correct_lut(*green);
  *blue = this->gamma_correct_lut(*blue);
}
void LightState::current_values_as_rgbw(float *red, float *green, float *blue, float *white) {
  this->current_values.as_rgbw(red, green, blue, white);
  *red = this->gamma_correct_lut(*red);
  *green = this->gamma_correct_lut(*green);
  *blue = this->gamma_correct_lut(*blue);
  *white = this->gamma_correct_lut(*white);
}
void LightState::current_values_as_rgbww(float *red, float *green, float *blue, float *cold_white, float *warm_white,
                                         bool constant_brightness) {
  this->current_values.as_rgb(red, green, blue);
  *red = this->gamma_correct_lut(*red);
  *green = this->gamma_correct_lut(*green);
  *blue = this->gamma_correct_lut(*blue);
  this->current_values_as_cwww(cold_white, warm_white, constant_brightness);
}
void LightState::current_values_as_rgbct(float *red, float *green, float *blue, float *color_temperature,
                                         float *white_brightness) {
  auto traits = this->get_traits();
  this->current_values.as_rgbct(traits.get_min_mireds(), traits.get_max_mireds(), red, green, blue, color_temperature,
                                white_brightness);
  *red = this->gamma_correct_lut(*red);
  *green = this->gamma_correct_lut(*green);
  *blue = this->gamma_correct_lut(*blue);
  *white_brightness = this->gamma_correct_lut(*white_brightness);
}
void LightState::current_values_as_cwww(float *cold_white, float *warm_white, bool constant_brightness) {
  if (!constant_brightness) {
    // Without constant_brightness, gamma commutes with simple multiplication:
    //   gamma(white_level * cw) = gamma(white_level) * gamma(cw)
    // (since gamma(a*b) = (a*b)^g = a^g * b^g = gamma(a) * gamma(b))
    // so applying gamma after is mathematically equivalent and simpler.
    this->current_values.as_cwww(cold_white, warm_white, false);
    *cold_white = this->gamma_correct_lut(*cold_white);
    *warm_white = this->gamma_correct_lut(*warm_white);
    return;
  }

  // For constant_brightness mode, gamma MUST be applied to the individual
  // channel values BEFORE the balancing formula (max/sum ratio), not after.
  //
  // Why: The cold_white_ and warm_white_ values stored in LightColorValues
  // are gamma-uncorrected (see transform_parameters_() which applies
  // gamma_uncorrect to the linear CW/WW fractions derived from color
  // temperature). Applying gamma_correct here recovers the original linear
  // fractions, which the constant_brightness formula then uses to distribute
  // power evenly. The max/sum formula ensures cold+warm PWM output sums to
  // a constant, keeping total power (and perceived brightness) the same
  // across all color temperatures.
  //
  // Applying gamma AFTER the formula would be incorrect because gamma is
  // nonlinear: gamma(a/b) != gamma(a)/gamma(b), so the carefully balanced
  // ratio would be distorted, causing a severe brightness dip at mid-range
  // color temperatures.
  const auto &v = this->current_values;
  if (!(v.get_color_mode() & ColorCapability::COLD_WARM_WHITE)) {
    *cold_white = *warm_white = 0;
    return;
  }

  const float cw_level = this->gamma_correct_lut(v.get_cold_white());
  const float ww_level = this->gamma_correct_lut(v.get_warm_white());
  const float white_level = this->gamma_correct_lut(v.get_state() * v.get_brightness());
  const float sum = cw_level > 0 || ww_level > 0 ? cw_level + ww_level : 1;  // Don't divide by zero.
  *cold_white = white_level * std::max(cw_level, ww_level) * cw_level / sum;
  *warm_white = white_level * std::max(cw_level, ww_level) * ww_level / sum;
}
void LightState::current_values_as_ct(float *color_temperature, float *white_brightness) {
  auto traits = this->get_traits();
  this->current_values.as_ct(traits.get_min_mireds(), traits.get_max_mireds(), color_temperature, white_brightness);
  *white_brightness = this->gamma_correct_lut(*white_brightness);
}

#ifdef USE_LIGHT_GAMMA_LUT
float LightState::gamma_correct_lut(float value) const {
  if (value <= 0.0f)
    return 0.0f;
  if (value >= 1.0f)
    return 1.0f;
  if (this->gamma_table_ == nullptr)
    return value;
  float scaled = value * 255.0f;
  auto idx = static_cast<uint8_t>(scaled);
  if (idx >= 255)
    return progmem_read_uint16(&this->gamma_table_[255]) / 65535.0f;
  float frac = scaled - idx;
  float a = progmem_read_uint16(&this->gamma_table_[idx]);
  float b = progmem_read_uint16(&this->gamma_table_[idx + 1]);
  return (a + frac * (b - a)) / 65535.0f;
}
float LightState::gamma_uncorrect_lut(float value) const {
  if (value <= 0.0f)
    return 0.0f;
  if (value >= 1.0f)
    return 1.0f;
  if (this->gamma_table_ == nullptr)
    return value;
  uint16_t target = static_cast<uint16_t>(value * 65535.0f);
  uint8_t lo = gamma_table_reverse_search(this->gamma_table_, target);
  if (lo >= 255)
    return 1.0f;
  // Interpolate between lo and lo+1
  uint16_t a = progmem_read_uint16(&this->gamma_table_[lo]);
  uint16_t b = progmem_read_uint16(&this->gamma_table_[lo + 1]);
  if (b == a)
    return lo / 255.0f;
  float frac = static_cast<float>(target - a) / static_cast<float>(b - a);
  return (lo + frac) / 255.0f;
}
#endif  // USE_LIGHT_GAMMA_LUT

bool LightState::is_transformer_active() { return this->is_transformer_active_; }

void LightState::start_effect_(uint32_t effect_index) {
  this->stop_effect_();
  if (effect_index == 0)
    return;

  this->active_effect_index_ = effect_index;
  auto *effect = this->get_active_effect_();
  effect->start_internal();
  // Enable loop while effect is active
  this->enable_loop();
}
LightEffect *LightState::get_active_effect_() {
  if (this->active_effect_index_ == 0) {
    return nullptr;
  } else {
    return this->effects_[this->active_effect_index_ - 1];
  }
}
void LightState::stop_effect_() {
  auto *effect = this->get_active_effect_();
  if (effect != nullptr) {
    effect->stop();
  }
  this->active_effect_index_ = 0;
  // Disable loop if idle (no effect and no transformer)
  this->disable_loop_if_idle_();
}

void LightState::start_transition_(const LightColorValues &target, uint32_t length, bool set_remote_values) {
  this->transformer_ = this->output_->create_default_transition();
  this->transformer_->setup(this->current_values, target, length);

  if (set_remote_values) {
    this->remote_values = target;
  }
  // Enable loop while transition is active
  this->enable_loop();
}

void LightState::start_flash_(const LightColorValues &target, uint32_t length, bool set_remote_values) {
  LightColorValues end_colors = this->remote_values;
  // If starting a flash if one is already happening, set end values to end values of current flash
  // Hacky but works
  if (this->transformer_ != nullptr)
    end_colors = this->transformer_->get_start_values();

  this->transformer_ = make_unique<LightFlashTransformer>(*this);
  this->transformer_->setup(end_colors, target, length);

  if (set_remote_values) {
    this->remote_values = target;
  };
  // Enable loop while flash is active
  this->enable_loop();
}

void LightState::set_immediately_(const LightColorValues &target, bool set_remote_values) {
  this->is_transformer_active_ = false;
  this->transformer_ = nullptr;
  this->current_values = target;
  if (set_remote_values) {
    this->remote_values = target;
  }
  this->output_->update_state(this);
  this->schedule_write_();
}

void LightState::disable_loop_if_idle_() {
  // Only disable loop if both transformer and effect are inactive, and no pending writes
  if (this->transformer_ == nullptr && this->get_active_effect_() == nullptr && !this->next_write_) {
    this->disable_loop();
  }
}

void LightState::save_remote_values_() {
  LightStateRTCState saved;
  saved.color_mode = this->remote_values.get_color_mode();
  switch (this->restore_mode_) {
    case LIGHT_RESTORE_AND_OFF:
    case LIGHT_RESTORE_AND_ON:
      saved.state = (this->restore_mode_ == LIGHT_RESTORE_AND_ON);
      break;
    default:
      saved.state = this->remote_values.is_on();
      break;
  }
  saved.brightness = this->remote_values.get_brightness();
  saved.color_brightness = this->remote_values.get_color_brightness();
  saved.red = this->remote_values.get_red();
  saved.green = this->remote_values.get_green();
  saved.blue = this->remote_values.get_blue();
  saved.white = this->remote_values.get_white();
  saved.color_temp = this->remote_values.get_color_temperature();
  saved.cold_white = this->remote_values.get_cold_white();
  saved.warm_white = this->remote_values.get_warm_white();
  saved.effect = this->active_effect_index_;
  this->rtc_.save(&saved);
}

}  // namespace esphome::light