/*
  WiFi compat with ESP32

  Copyright (C) 2021  Theo Arends / Jörg Schüler-Maroldt

  This program is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/
#ifdef ESP32

#include "Arduino.h"
#include "WiFiHelper.h"
#include <esp_wifi.h>

extern void AddLog(uint32_t loglevel, PGM_P formatP, ...);
enum LoggingLevels {LOG_LEVEL_NONE, LOG_LEVEL_ERROR, LOG_LEVEL_INFO, LOG_LEVEL_DEBUG, LOG_LEVEL_DEBUG_MORE};


ip_addr_t dns_save4[2] = {};      // IPv4 DNS servers
#ifdef USE_IPV6
ip_addr_t dns_save6[2] = {};      // IPv6 DNS servers
#endif // USE_IPV6

#include "tasmota_options.h"
#include "lwip/dns.h"

#if CONFIG_ESP_WIFI_ENTERPRISE_SUPPORT
wl_status_t WiFiHelper::begin(const char *wpa2_ssid, wpa2_auth_method_t method, const char *wpa2_identity, const char *wpa2_username, const char *wpa2_password, const char *ca_pem, const char *client_crt, const char *client_key, int ttls_phase2_type, int32_t channel, const uint8_t *bssid, bool connect) {
  WiFiHelper::scrubDNS();
  wl_status_t ret = WiFi.begin(wpa2_ssid, method, wpa2_identity, wpa2_username, wpa2_password, ca_pem, client_crt, client_key, ttls_phase2_type, channel, bssid, connect);
  WiFiHelper::scrubDNS();
  return ret;
}
#endif // CONFIG_ESP_WIFI_ENTERPRISE_SUPPORT

wl_status_t WiFiHelper::begin(const char* ssid, const char *passphrase, int32_t channel, const uint8_t* bssid, bool connect) {
  WiFiHelper::scrubDNS();
  wl_status_t ret = WiFi.begin(ssid, passphrase, channel, bssid, connect);
  WiFiHelper::scrubDNS();
  return ret;
}

wl_status_t WiFiHelper::begin(char* ssid, char *passphrase, int32_t channel, const uint8_t* bssid, bool connect) {
  WiFiHelper::scrubDNS();
  wl_status_t ret = WiFi.begin(ssid, passphrase, channel, bssid, connect);
  WiFiHelper::scrubDNS();
  return ret;
}
wl_status_t WiFiHelper::begin() {
  WiFiHelper::scrubDNS();
  wl_status_t ret = WiFi.begin();
  WiFiHelper::scrubDNS();
  return ret;
}

// scrubDNS
//
// LWIP has a single DNS table for all interfaces and for v4/v6
// Unfortunately when trying to connect to Wifi, the dns server table is erased.
//
// We restore DNS previous values if they are empty
// We restore or erase DNS entries if they are unsupported (v4 vs v6)
extern bool WifiHasIPv4(void);
extern bool EthernetHasIPv4(void);
extern bool WifiHasIPv6(void);
extern bool EthernetHasIPv6(void);

void WiFiHelper::scrubDNS(void) {
  // AddLog(LOG_LEVEL_DEBUG, "IP>1: dns_save4 %s %s dns_save6 %s %s",
  //     IPAddress(&dns_save4[0]).toString().c_str(),IPAddress(&dns_save4[1]).toString().c_str(),
  //     IPAddress(&dns_save6[0]).toString().c_str(),IPAddress(&dns_save6[1]).toString().c_str());
  // String dns_entry0 = IPAddress(dns_getserver(0)).toString();
  // String dns_entry1 = IPAddress(dns_getserver(1)).toString();

  // scan DNS entries
  bool has_v4 = WifiHasIPv4() || EthernetHasIPv4();
#ifdef USE_IPV6
  bool has_v6 = WifiHasIPv6() || EthernetHasIPv6();
#endif
  // AddLog(LOG_LEVEL_DEBUG, "IP>1: DNS: (%s %s) has4/6:%i-%i", dns_entry0.c_str(), dns_entry1.c_str(), has_v4, has_v6);

  // First pass, save values
  for (uint32_t i=0; i<2; i++) {
    const ip_addr_t* ip_dns = dns_getserver(i);
    // Step 1. save valid values from DNS
    if (ip_dns && !ip_addr_isany_val(*ip_dns)) {
#ifdef USE_IPV6
      if (IP_IS_V4(ip_dns)) {
#endif // USE_IPV6
        dns_save4[i] = *ip_dns;       // dns entry is populated, save it in v4 slot
#ifdef USE_IPV6
      }
      else if (has_v6) {
        dns_save6[i] = *ip_dns;       // dns entry is populated, save it in v6 slot
      }
#endif // USE_IPV6
    }
  }

  // Step 2. scrub addresses not supported
#ifdef USE_IPV6
  // Note: do NOT clear dns_save4 when !has_v4 && has_v6.
  // IPv4 may be temporarily unavailable during WiFi reconnect while IPv6 RA
  // completes before DHCP. Clearing dns_save4 here would lose the saved IPv4
  // DNS server, causing all subsequent DNS lookups to fail until DHCP restores
  // it. Step 3 already uses dns_save6 exclusively when !has_v4 && has_v6,
  // so clearing dns_save4 here is both harmful and redundant.
  if (!has_v6) {
    dns_save6[0] = *IP_ADDR_ANY;
    dns_save6[1] = *IP_ADDR_ANY;
  }

  // Step 3. restore saved value
  if (has_v4 && has_v6) {   // if both IPv4 and IPv6 are active, prefer IPv4 for first and IPv6 for second
    if (!ip_addr_isany_val(dns_save4[0])) {
      // First DNS IPv4
      dns_setserver(0, &dns_save4[0]);
      if (!ip_addr_isany_val(dns_save6[0])) {
        dns_setserver(1, &dns_save6[0]);    // take first IPv6 as second DNS
      } else {
        dns_setserver(1, &dns_save4[1]);    // or revert to second IPv4
      }
    } else {
      // If no DNS IPv4, use IPv6
      dns_setserver(0, &dns_save6[0]);
      dns_setserver(1, &dns_save6[1]);
    }
  } else if (has_v6) {                      // v6 and no v4
    dns_setserver(0, &dns_save6[0]);
    dns_setserver(1, &dns_save6[1]);
  } else {                                  // no v6, we use v4 even if not connected
#endif // USE_IPV6
    dns_setserver(0, &dns_save4[0]);
    dns_setserver(1, &dns_save4[1]);
#ifdef USE_IPV6
  }
#endif // USE_IPV6
  // AddLog(LOG_LEVEL_DEBUG, "IP>2: DNS: from(%s %s) to (%s %s) has4/6:%i-%i", dns_entry0.c_str(), dns_entry1.c_str(), IPAddress(dns_getserver(0)).toString().c_str(),  IPAddress(dns_getserver(1)).toString().c_str(), has_v4, has_v6);

  // AddLog(LOG_LEVEL_DEBUG, "IP>2: dns_save4 %s %s dns_save6 %s %s",
  //     IPAddress(&dns_save4[0]).toString().c_str(),IPAddress(&dns_save4[1]).toString().c_str(),
  //     IPAddress(&dns_save6[0]).toString().c_str(),IPAddress(&dns_save6[1]).toString().c_str());
}

void WiFiHelper::hostname(const char* aHostname) {
  WiFi.setHostname(aHostname);
}

void WiFiHelper::setSleepMode(int iSleepMode) {
  // WIFI_LIGHT_SLEEP and WIFI_MODEM_SLEEP
  WiFi.setSleep(iSleepMode != WIFI_NONE_SLEEP);
}

int WiFiHelper::getPhyMode() {
  /*
    // IDF v5.2.6
    typedef enum {
      WIFI_PHY_MODE_LR,   // PHY mode for Low Rate
      WIFI_PHY_MODE_11B,  // PHY mode for 11b
      WIFI_PHY_MODE_11G,  // PHY mode for 11g
      WIFI_PHY_MODE_HT20, // PHY mode for Bandwidth HT20 (11n)
      WIFI_PHY_MODE_HT40, // PHY mode for Bandwidth HT40 (11n)
      WIFI_PHY_MODE_HE20, // PHY mode for Bandwidth HE20 (11ax)
    } wifi_phy_mode_t;

    // IDF v5.3.4
    typedef enum {
      WIFI_PHY_MODE_LR,   // PHY mode for Low Rate
      WIFI_PHY_MODE_11B,  // PHY mode for 11b
      WIFI_PHY_MODE_11G,  // PHY mode for 11g
      WIFI_PHY_MODE_11A,  // PHY mode for 11a
      WIFI_PHY_MODE_HT20, // PHY mode for Bandwidth HT20
      WIFI_PHY_MODE_HT40, // PHY mode for Bandwidth HT40
      WIFI_PHY_MODE_HE20, // PHY mode for Bandwidth HE20
      WIFI_PHY_MODE_VHT20,// PHY mode for Bandwidth VHT20
    } wifi_phy_mode_t;
  */
  int phy_mode = 0;  // "low rate|11b|11g|11a|HT20|HT40|HE20|VHT20"
  wifi_phy_mode_t WiFiMode;
  if (esp_wifi_sta_get_negotiated_phymode(&WiFiMode) == ESP_OK) {
    phy_mode = (int)WiFiMode;
    if (phy_mode > 7) {
      phy_mode = 7;
    }
  }
  return phy_mode;
}

bool WiFiHelper::setPhyMode(WiFiPhyMode_t mode) {
  /*
    mode 1 = B
    mode 2 = BG
    mode 3 = BGN
    mode 4 = BGNAX
  */
  /*
    ESPEasy:
    HT20 = 20 MHz channel width.
    HT40 = 40 MHz channel width.
    In theory, HT40 can offer upto 150 Mbps connection speed.
    However since HT40 is using nearly all channels on 2.4 GHz WiFi,
    Thus you are more likely to experience disturbances.
    The response speed and stability is better at HT20 for ESP units.
  */
  esp_wifi_set_bandwidth(WIFI_IF_STA, WIFI_BW_HT20);
  /*
    ESPEasy:
    Set to use "Long GI" making it more resilliant to reflections
    See: https://www.tp-link.com/us/configuration-guides/q_a_basic_wireless_concepts/?configurationId=2958#_idTextAnchor038
  */
  esp_wifi_config_80211_tx_rate(WIFI_IF_STA, WIFI_PHY_RATE_MCS3_LGI);

  uint8_t protocol_bitmap = WIFI_PROTOCOL_11B;      // 1
  switch (mode) {
    case 4: protocol_bitmap |= WIFI_PROTOCOL_11AX;  // 16
    case 3: protocol_bitmap |= WIFI_PROTOCOL_11N;   // 4
    case 2: protocol_bitmap |= WIFI_PROTOCOL_11G;   // 2
  }
  return (ESP_OK == esp_wifi_set_protocol(WIFI_IF_STA, protocol_bitmap));
}

void WiFiHelper::setOutputPower(int n) {
    wifi_power_t p = WIFI_POWER_2dBm;
    if (n > 19)
        p = WIFI_POWER_19_5dBm;
    else if (n > 18)
        p = WIFI_POWER_18_5dBm;
    else if (n >= 17)
        p = WIFI_POWER_17dBm;
    else if (n >= 15)
        p = WIFI_POWER_15dBm;
    else if (n >= 13)
        p = WIFI_POWER_13dBm;
    else if (n >= 11)
        p = WIFI_POWER_11dBm;
    else if (n >= 8)
        p = WIFI_POWER_8_5dBm;
    else if (n >= 7)
        p = WIFI_POWER_7dBm;
    else if (n >= 5)
        p = WIFI_POWER_5dBm;
    WiFi.setTxPower(p);
}

void WiFiHelper::forceSleepBegin() {
}

void WiFiHelper::forceSleepWake() {
}

bool WiFiHelper::getNetworkInfo(uint8_t i, String &ssid, uint8_t &encType, int32_t &rssi, uint8_t *&bssid, int32_t &channel, bool &hidden_scan) {
    hidden_scan = false;
    return WiFi.getNetworkInfo(i, ssid, encType, rssi, bssid, channel);
}

//
// Manage dns callbacks from lwip DNS resolver.
// We need a trick here, because the callback may be called after we timed-out and
// launched a new request. We need to discard outdated responses.
//
// We use a static ip_addr_t (anyways we don't support multi-threading)
// and use a counter so the callback can check if it's responding to the current
// request or to an old one (hence discard)
//
// It's not an issue to have old requests in flight. LWIP has a default limit of 4
// DNS requests in flight.
// If the buffer for in-flight requests is full, LWIP removes the oldest from the list.
// (it does not block new DNS resolutions)
static ip_addr_t dns_ipaddr;
static volatile uint32_t ip_addr_counter = 0;   // counter for requests
extern int32_t WifiDNSGetTimeout(void);
extern bool WifiDNSGetIPv6Priority(void);

static volatile bool dns_found = false;

bool DNS_TimeReached(uint32_t timer)
{
  // Check if a certain timeout has been reached.
  int32_t passed = ((int32_t) (millis() - timer));
  return (passed >= 0);
}


static void wifi32_dns_found_callback(const char *name, const ip_addr_t *ipaddr, void *callback_arg)
{
  // Serial.printf("DNS: cb name=%s ipaddr=%s arg=%i counter=%i\n", name ? name : "<null>", IPAddress(ipaddr).toString().c_str(), (int) callback_arg, ip_addr_counter);
  uint32_t cb_counter = (uint32_t) callback_arg;
  if (cb_counter != ip_addr_counter) { return; }    // the response is from a previous request, ignore

  if (ipaddr != nullptr) {
    dns_ipaddr = *ipaddr;
  } else {
    dns_ipaddr = *IP4_ADDR_ANY;   // set to IPv4 0.0.0.0
  }
  dns_found = true;
  // AddLog(LOG_LEVEL_DEBUG, "WIF: dns_found=%s", ipaddr ? IPAddress(*ipaddr).toString().c_str() : "<null>");
}

// Auto-fix zone
//
// After a reconnect, the zone id may not be valid anymore
// In such case we detect any "%st<n>" or "%en<n>" zone identifier
// and replace with the current zone id
extern bool WifiGetIPv6LinkLocal(IPAddress *ip);
extern bool EthernetGetIPv6LinkLocal(IPAddress *ip);
void WiFiHelper::IPv6ZoneAutoFix(IPAddress &addr, const char* aHostname) {
#ifdef USE_IPV6
  if ((addr.type() == IPv6) && (addr.zone() == 0)) {
    // check if hostname contains '%'
    const char *zone_identifier = strchr(aHostname, '%');
    if (zone_identifier != nullptr) {
      uint8_t zone_id = 0;
      // check if zone id is valid
      if (strncmp(zone_identifier, "%st", 3) == 0) {
        IPAddress wifi_link_local;
        if (WifiGetIPv6LinkLocal(&wifi_link_local)) {
          zone_id = wifi_link_local.zone();
        }
      } else if (strncmp(zone_identifier, "%en", 3) == 0) {
        IPAddress eth_link_local;
        if (EthernetGetIPv6LinkLocal(&eth_link_local)) {
          zone_id = eth_link_local.zone();
        }
      }
      if (zone_id) {
        // convert to ip_addr_t which is currently the only way to change the zone_id
        ip_addr_t ip_addr;
        addr.to_ip_addr_t(&ip_addr);
        ip_addr.u_addr.ip6.zone = zone_id;
        addr = IPAddress(&ip_addr);
      }
    }
  }
#endif
}

/**
 * Resolve the given hostname to an IP address.
 * @param aHostname     Name to be resolved
 * @param aResult       IPAddress structure to store the returned IP address
 * @return 1 if aIPAddrString was successfully converted to an IP address,
 *          else error code
 */
int WiFiHelper::hostByName(const char* aHostname, IPAddress& aResult, int32_t timer_ms)
{
//   return WiFi.hostByName(aHostname, aResult);
// #if 0
  ip_addr_t addr;
  aResult = (uint32_t) 0;     // by default set to IPv4 0.0.0.0
  dns_ipaddr = *IP4_ADDR_ANY;  // by default set to IPv4 0.0.0.0

  WiFiHelper::scrubDNS();    // internal calls to reconnect can zero the DNS servers, save DNS for future use
  ip_addr_counter++;      // increase counter, from now ignore previous responses
  // clearStatusBits(NET_DNS_IDLE_BIT | NET_DNS_DONE_BIT);
  uint8_t v4v6priority = LWIP_DNS_ADDRTYPE_IPV4;
#ifdef USE_IPV6
  v4v6priority = WifiDNSGetIPv6Priority() ? LWIP_DNS_ADDRTYPE_IPV6_IPV4 : LWIP_DNS_ADDRTYPE_IPV4_IPV6;
#endif // USE_IPV6
  dns_found = false;
  err_t err = dns_gethostbyname_addrtype(aHostname, &dns_ipaddr, &wifi32_dns_found_callback, (void*) ip_addr_counter, v4v6priority);
  // Serial.printf("DNS: dns_gethostbyname_addrtype errg=%i counter=%i\n", err, ip_addr_counter);
  if(err == ERR_OK && !ip_addr_isany_val(dns_ipaddr)) {
#ifdef USE_IPV6
    aResult.from_ip_addr_t(&dns_ipaddr);
#else // USE_IPV6
    aResult = ip_addr_get_ip4_u32(&dns_ipaddr);
#endif // USE_IPV6
  } else if(err == ERR_INPROGRESS) {
    uint32_t deadline = millis() + timer_ms;
    while ((!DNS_TimeReached(deadline)) && !dns_found) {
      delay(1);
    }
  }

  if (!ip_addr_isany_val(dns_ipaddr)) {
#ifdef USE_IPV6
    aResult.from_ip_addr_t(&dns_ipaddr);
    WiFiHelper::IPv6ZoneAutoFix(aResult, aHostname);
#else // USE_IPV6
    aResult = ip_addr_get_ip4_u32(&dns_ipaddr);
#endif // USE_IPV6
    return true;
  }
  return false;
// #endif
}

int WiFiHelper::hostByName(const char* aHostname, IPAddress& aResult)
{
  return WiFiHelper::hostByName(aHostname, aResult, WifiDNSGetTimeout());
}

#include "esp_mac.h"

String WiFiHelper::macAddress(void) {
  uint8_t mac[6] = {0,0,0,0,0,0};
  char macStr[18] = { 0 };
#ifdef CONFIG_SOC_HAS_WIFI
  esp_read_mac(mac, ESP_MAC_WIFI_STA);  // Local WiFi station MAC address
#elif CONFIG_ESP_WIFI_REMOTE_ENABLED
  WiFi.macAddress(mac);                 // Remote WiFi station MAC address (devices without WiFi but hostedMCU)
#else   // No CONFIG_SOC_HAS_WIFI
  esp_read_mac(mac, ESP_MAC_BASE);      // Local hardware base MAC address
#endif  // CONFIG_SOC_HAS_WIFI
  snprintf(macStr, sizeof(macStr), "%02X:%02X:%02X:%02X:%02X:%02X", mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
  return String(macStr);
}

#endif  // ESP32
