#include "nr_antenna_state.h"
#include "common/utils/LOG/log.h"
#include <string.h>
#include <time.h>

// Helper to count set bits in a byte
static inline int popcount8(uint8_t x) {
  int count = 0;
  while (x) {
    count += x & 1;
    x >>= 1;
  }
  return count;
}

// Get current timestamp in microseconds
static inline uint64_t get_timestamp_us(void) {
  struct timespec ts;
  clock_gettime(CLOCK_MONOTONIC, &ts);
  return (uint64_t)ts.tv_sec * 1000000ULL + (uint64_t)ts.tv_nsec / 1000ULL;
}

void nr_init_antenna_state(nr_antenna_state_t *state, 
                           uint8_t num_tx, 
                           uint8_t num_rx,
                           const char *tx_array_name,
                           const char *rx_array_name)
{
  if (state == NULL) return;
  
  memset(state, 0, sizeof(nr_antenna_state_t));
  
  state->num_tx_antennas = (num_tx > NR_MAX_ANTENNAS) ? NR_MAX_ANTENNAS : num_tx;
  state->num_rx_antennas = (num_rx > NR_MAX_ANTENNAS) ? NR_MAX_ANTENNAS : num_rx;
  
  // Initialize with all configured antennas active
  // Create bitmask with num_tx/num_rx lowest bits set
  state->tx_antenna_bitmask = (1 << state->num_tx_antennas) - 1;
  state->requested_tx_antenna_bitmask = state->tx_antenna_bitmask;
  state->tx_mask_update_pending = false;
  state->rx_antenna_bitmask = (1 << state->num_rx_antennas) - 1;
  
  // Set O-RAN array names
  if (tx_array_name) {
    strncpy(state->tx_array_name, tx_array_name, sizeof(state->tx_array_name) - 1);
  } else {
    strncpy(state->tx_array_name, "tx-array-0", sizeof(state->tx_array_name) - 1);
  }
  
  if (rx_array_name) {
    strncpy(state->rx_array_name, rx_array_name, sizeof(state->rx_array_name) - 1);
  } else {
    strncpy(state->rx_array_name, "rx-array-0", sizeof(state->rx_array_name) - 1);
  }
  
  pthread_mutex_init(&state->lock, NULL);
  
  state->stats.last_update_timestamp = get_timestamp_us();
  
  // Initialize recovery mode
  state->recovery.in_recovery = false;
  state->recovery.recovery_start_ts = 0;
  state->recovery.recovery_duration_ms = 500; // 500ms default recovery time
  
  // Protect antenna 0 by default (needed for SSB/PSS/SSS)
  state->config.protect_antenna0 = true;
  
  LOG_I(PHY, "[ANTENNA_STATE] Initialized: TX antennas=%d (mask=0x%02X, array=%s), "
             "RX antennas=%d (mask=0x%02X, array=%s), antenna0_protected=%d\n",
        state->num_tx_antennas, state->tx_antenna_bitmask, state->tx_array_name,
        state->num_rx_antennas, state->rx_antenna_bitmask, state->rx_array_name,
        state->config.protect_antenna0);
}

void nr_free_antenna_state(nr_antenna_state_t *state)
{
  if (state == NULL) return;
  pthread_mutex_destroy(&state->lock);
}

static uint8_t valid_tx_mask(const nr_antenna_state_t *state)
{
  return (uint8_t)((1u << state->num_tx_antennas) - 1u);
}

/* Caller holds state->lock. */
static void set_tx_antenna_bitmask_locked(nr_antenna_state_t *state, uint8_t bitmask)
{
  const uint8_t old_bitmask = state->tx_antenna_bitmask;
  state->tx_antenna_bitmask = bitmask & valid_tx_mask(state);

  if (state->tx_antenna_bitmask == old_bitmask)
    return;

  state->stats.tx_mute_events++;
  state->stats.last_update_timestamp = get_timestamp_us();
  LOG_I(PHY, "[ANTENNA_STATE] TX bitmask changed: 0x%02X -> 0x%02X (array=%s, active=%d/%d)\n",
        old_bitmask,
        state->tx_antenna_bitmask,
        state->tx_array_name,
        popcount8(state->tx_antenna_bitmask),
        state->num_tx_antennas);
}

uint8_t nr_set_tx_antenna_bitmask(nr_antenna_state_t *state, uint8_t bitmask)
{
  if (state == NULL) return 0;
  
  pthread_mutex_lock(&state->lock);
  const uint8_t old_bitmask = state->tx_antenna_bitmask;
  state->requested_tx_antenna_bitmask = bitmask & valid_tx_mask(state);
  state->tx_mask_update_pending = false;
  set_tx_antenna_bitmask_locked(state, state->requested_tx_antenna_bitmask);
  pthread_mutex_unlock(&state->lock);
  return old_bitmask;
}

uint8_t nr_request_tx_antenna_bitmask(nr_antenna_state_t *state, uint8_t bitmask)
{
  if (state == NULL) return 0;

  pthread_mutex_lock(&state->lock);
  const uint8_t old_requested = state->requested_tx_antenna_bitmask;
  const uint8_t requested = bitmask & valid_tx_mask(state);
  state->requested_tx_antenna_bitmask = requested;

  /* Adding chains cannot invalidate an in-flight PDSCH. Removing a chain can,
   * so leave the physical mask unchanged until MAC has drained HARQ processes
   * whose original layer count exceeds the new target. */
  if ((state->tx_antenna_bitmask & (uint8_t)~requested) == 0) {
    state->tx_mask_update_pending = false;
    set_tx_antenna_bitmask_locked(state, requested);
  } else {
    state->tx_mask_update_pending = true;
    LOG_I(PHY,
          "[ANTENNA_STATE] TX downshift requested: active=0x%02X target=0x%02X; waiting for DL HARQ drain\n",
          state->tx_antenna_bitmask,
          requested);
  }
  pthread_mutex_unlock(&state->lock);
  return old_requested;
}

uint8_t nr_get_requested_tx_antenna_bitmask(const nr_antenna_state_t *state)
{
  if (state == NULL) return 0xFF;
  pthread_mutex_lock((pthread_mutex_t *)&state->lock);
  const uint8_t requested = state->requested_tx_antenna_bitmask;
  pthread_mutex_unlock((pthread_mutex_t *)&state->lock);
  return requested;
}

bool nr_get_pending_tx_antenna_bitmask(const nr_antenna_state_t *state, uint8_t *bitmask)
{
  if (state == NULL) return false;
  pthread_mutex_lock((pthread_mutex_t *)&state->lock);
  const bool pending = state->tx_mask_update_pending;
  if (pending && bitmask != NULL)
    *bitmask = state->requested_tx_antenna_bitmask;
  pthread_mutex_unlock((pthread_mutex_t *)&state->lock);
  return pending;
}

bool nr_apply_requested_tx_antenna_bitmask(nr_antenna_state_t *state, uint8_t expected_bitmask)
{
  if (state == NULL) return false;
  pthread_mutex_lock(&state->lock);
  const bool matches = state->tx_mask_update_pending
                       && state->requested_tx_antenna_bitmask == expected_bitmask;
  if (matches) {
    set_tx_antenna_bitmask_locked(state, expected_bitmask);
    state->tx_mask_update_pending = false;
  }
  pthread_mutex_unlock(&state->lock);
  return matches;
}

uint8_t nr_set_rx_antenna_bitmask(nr_antenna_state_t *state, uint8_t bitmask)
{
  if (state == NULL) return 0;
  
  pthread_mutex_lock(&state->lock);
  
  uint8_t old_bitmask = state->rx_antenna_bitmask;
  
  // Mask to only valid antenna bits
  uint8_t valid_mask = (1 << state->num_rx_antennas) - 1;
  state->rx_antenna_bitmask = bitmask & valid_mask;
  
  // Update statistics
  if (state->rx_antenna_bitmask != old_bitmask) {
    state->stats.rx_mute_events++;
    state->stats.last_update_timestamp = get_timestamp_us();
    
    LOG_I(PHY, "[ANTENNA_STATE] RX bitmask changed: 0x%02X -> 0x%02X (array=%s, active=%d/%d)\n",
          old_bitmask, state->rx_antenna_bitmask, state->rx_array_name,
          popcount8(state->rx_antenna_bitmask), state->num_rx_antennas);
  }
  
  pthread_mutex_unlock(&state->lock);
  
  return old_bitmask;
}

uint8_t nr_get_tx_antenna_bitmask(const nr_antenna_state_t *state)
{
  if (state == NULL) return 0xFF;
  pthread_mutex_lock((pthread_mutex_t*)&state->lock);
  uint8_t m = state->tx_antenna_bitmask;
  pthread_mutex_unlock((pthread_mutex_t*)&state->lock);
  return m;
}

uint8_t nr_get_rx_antenna_bitmask(const nr_antenna_state_t *state)
{
  if (state == NULL) return 0xFF;
  pthread_mutex_lock((pthread_mutex_t*)&state->lock);
  uint8_t m = state->rx_antenna_bitmask;
  pthread_mutex_unlock((pthread_mutex_t*)&state->lock);
  return m;
}

bool nr_is_tx_antenna_active(const nr_antenna_state_t *state, uint8_t ant_id)
{
  if (state == NULL || ant_id >= state->num_tx_antennas) return false;
  return (state->tx_antenna_bitmask & (1 << ant_id)) != 0;
}

bool nr_is_rx_antenna_active(const nr_antenna_state_t *state, uint8_t ant_id)
{
  if (state == NULL || ant_id >= state->num_rx_antennas) return false;
  return (state->rx_antenna_bitmask & (1 << ant_id)) != 0;
}

int nr_get_active_tx_antenna_count(const nr_antenna_state_t *state)
{
  if (state == NULL) return 0;
  return popcount8(state->tx_antenna_bitmask);
}

int nr_get_active_rx_antenna_count(const nr_antenna_state_t *state)
{
  if (state == NULL) return 0;
  return popcount8(state->rx_antenna_bitmask);
}

void nr_mute_tx_antenna(nr_antenna_state_t *state, uint8_t ant_id)
{
  if (state == NULL || ant_id >= state->num_tx_antennas) return;
  
  // Protect antenna 0 if configured (needed for SSB/PSS/SSS transmission)
  if (ant_id == 0 && state->config.protect_antenna0) {
    LOG_W(PHY, "[ANTENNA_STATE] Cannot mute TX antenna 0: protected for SSB transmission\n");
    LOG_W(PHY, "[ANTENNA_STATE] Use 'SET TX MASK' to override protection if needed\n");
    return;
  }
  
  uint8_t new_mask = nr_get_requested_tx_antenna_bitmask(state) & ~(1 << ant_id);
  nr_request_tx_antenna_bitmask(state, new_mask);
}

void nr_unmute_tx_antenna(nr_antenna_state_t *state, uint8_t ant_id)
{
  if (state == NULL || ant_id >= state->num_tx_antennas) return;
  
  uint8_t new_mask = nr_get_requested_tx_antenna_bitmask(state) | (1 << ant_id);
  nr_request_tx_antenna_bitmask(state, new_mask);
}

void nr_mute_rx_antenna(nr_antenna_state_t *state, uint8_t ant_id)
{
  if (state == NULL || ant_id >= state->num_rx_antennas) return;
  
  uint8_t new_mask = state->rx_antenna_bitmask & ~(1 << ant_id);
  nr_set_rx_antenna_bitmask(state, new_mask);
}

void nr_unmute_rx_antenna(nr_antenna_state_t *state, uint8_t ant_id)
{
  if (state == NULL || ant_id >= state->num_rx_antennas) return;
  
  uint8_t new_mask = state->rx_antenna_bitmask | (1 << ant_id);
  nr_set_rx_antenna_bitmask(state, new_mask);
}

int nr_get_effective_mimo_layers(const nr_antenna_state_t *state, int requested_layers)
{
  if (state == NULL)
    return requested_layers;

  /* Cap new transmissions to the target immediately. During a downshift the
   * physical mask can remain wider for a few HARQ rounds, which is safe; the
   * reverse ordering would erase layers from buffered retransmissions. */
  const int target = popcount8(nr_get_requested_tx_antenna_bitmask(state));
  if (target > 0 && target < requested_layers)
    return target;
  return requested_layers;
}

void nr_log_antenna_state(const nr_antenna_state_t *state)
{
  if (state == NULL) {
    LOG_I(PHY, "[ANTENNA_STATE] State is NULL\n");
    return;
  }
  
  LOG_I(PHY, "[ANTENNA_STATE] === Antenna State Report ===\n");
  LOG_I(PHY, "[ANTENNA_STATE] TX Array: %s\n", state->tx_array_name);
  LOG_I(PHY, "[ANTENNA_STATE]   Configured: %d antennas\n", state->num_tx_antennas);
  LOG_I(PHY, "[ANTENNA_STATE]   Bitmask: 0x%02X\n", state->tx_antenna_bitmask);
  LOG_I(PHY, "[ANTENNA_STATE]   Active: %d antennas\n", nr_get_active_tx_antenna_count(state));
  LOG_I(PHY, "[ANTENNA_STATE]   Individual: ");
  for (int i = 0; i < state->num_tx_antennas; i++) {
    LOG_I(PHY, "TX%d=%s ", i, nr_is_tx_antenna_active(state, i) ? "ON" : "OFF");
  }
  LOG_I(PHY, "\n");
  
  LOG_I(PHY, "[ANTENNA_STATE] RX Array: %s\n", state->rx_array_name);
  LOG_I(PHY, "[ANTENNA_STATE]   Configured: %d antennas\n", state->num_rx_antennas);
  LOG_I(PHY, "[ANTENNA_STATE]   Bitmask: 0x%02X\n", state->rx_antenna_bitmask);
  LOG_I(PHY, "[ANTENNA_STATE]   Active: %d antennas\n", nr_get_active_rx_antenna_count(state));
  
  LOG_I(PHY, "[ANTENNA_STATE] Statistics:\n");
  LOG_I(PHY, "[ANTENNA_STATE]   TX mute events: %lu\n", state->stats.tx_mute_events);
  LOG_I(PHY, "[ANTENNA_STATE]   RX mute events: %lu\n", state->stats.rx_mute_events);
  LOG_I(PHY, "[ANTENNA_STATE]   Total muted slots: %lu\n", state->stats.total_muted_slots);
  LOG_I(PHY, "[ANTENNA_STATE] =============================\n");
}

void nr_get_antenna_stats(const nr_antenna_state_t *state,
                          uint64_t *tx_mute_events,
                          uint64_t *rx_mute_events,
                          uint64_t *total_muted_slots)
{
  if (state == NULL) return;
  
  if (tx_mute_events) *tx_mute_events = state->stats.tx_mute_events;
  if (rx_mute_events) *rx_mute_events = state->stats.rx_mute_events;
  if (total_muted_slots) *total_muted_slots = state->stats.total_muted_slots;
}

void nr_reset_antenna_stats(nr_antenna_state_t *state)
{
  if (state == NULL) return;
  
  pthread_mutex_lock(&state->lock);
  
  state->stats.tx_mute_events = 0;
  state->stats.rx_mute_events = 0;
  state->stats.total_muted_slots = 0;
  state->stats.energy_saved_estimate = 0;
  state->stats.last_update_timestamp = get_timestamp_us();
  
  LOG_I(PHY, "[ANTENNA_STATE] Statistics reset\n");
  
  pthread_mutex_unlock(&state->lock);
}

void nr_unmute_all_tx_antennas(nr_antenna_state_t *state)
{
  if (state == NULL) return;
  
  pthread_mutex_lock(&state->lock);
  
  uint8_t old_mask = state->tx_antenna_bitmask;
  uint8_t all_on = (1 << state->num_tx_antennas) - 1;
  state->requested_tx_antenna_bitmask = all_on;
  state->tx_mask_update_pending = false;
  state->tx_antenna_bitmask = all_on;

  // Enter recovery mode if we were muted
  if (old_mask != all_on) {
    state->recovery.in_recovery = true;
    state->recovery.recovery_start_ts = get_timestamp_us();
    LOG_I(PHY, "[ANTENNA_STATE] Entering recovery mode for %u ms after unmuting TX antennas\n",
          state->recovery.recovery_duration_ms);
  }
  
  state->stats.tx_mute_events++;
  state->stats.last_update_timestamp = get_timestamp_us();
  
  LOG_I(PHY, "[ANTENNA_STATE] All TX antennas unmuted: 0x%02X -> 0x%02X (array=%s)\n",
        old_mask, all_on, state->tx_array_name);
  
  pthread_mutex_unlock(&state->lock);
}

void nr_unmute_all_rx_antennas(nr_antenna_state_t *state)
{
  if (state == NULL) return;
  
  pthread_mutex_lock(&state->lock);
  
  uint8_t old_mask = state->rx_antenna_bitmask;
  uint8_t all_on = (1 << state->num_rx_antennas) - 1;
  state->rx_antenna_bitmask = all_on;
  
  state->stats.rx_mute_events++;
  state->stats.last_update_timestamp = get_timestamp_us();
  
  LOG_I(PHY, "[ANTENNA_STATE] All RX antennas unmuted: 0x%02X -> 0x%02X (array=%s)\n",
        old_mask, all_on, state->rx_array_name);
  
  pthread_mutex_unlock(&state->lock);
}

bool nr_is_in_recovery_mode(const nr_antenna_state_t *state)
{
  if (state == NULL) return false;
  
  if (!state->recovery.in_recovery) return false;
  
  // Check if recovery period has elapsed
  uint64_t now = get_timestamp_us();
  uint64_t elapsed_us = now - state->recovery.recovery_start_ts;
  uint64_t duration_us = (uint64_t)state->recovery.recovery_duration_ms * 1000ULL;
  
  if (elapsed_us >= duration_us) {
    // Recovery complete - cast away const for state update
    nr_antenna_state_t *s = (nr_antenna_state_t *)state;
    s->recovery.in_recovery = false;
    LOG_I(PHY, "[ANTENNA_STATE] Recovery mode complete after %lu ms\n", elapsed_us / 1000);
    return false;
  }
  
  return true;
}

void nr_set_antenna0_protection(nr_antenna_state_t *state, bool protect)
{
  if (state == NULL) return;
  
  pthread_mutex_lock(&state->lock);
  state->config.protect_antenna0 = protect;
  LOG_I(PHY, "[ANTENNA_STATE] Antenna 0 protection %s\n", protect ? "enabled" : "disabled");
  pthread_mutex_unlock(&state->lock);
}