#include <HTTPClient.h>
#include "HermitCrab.h"
#include "UPnpClient.h"
#include "Config.h"
#include "WiFiHost.h"
#define TAG_UPNP "UPnP"


bool CUpnpClient::registerUPnP(uint32_t *pip, uint16_t *pport) {
    routerIP = *pip;        // Gateway IP address
    publicPort = *pport;  // host.m_nPublicPort = config.m_nPublicPort
    publicIP = 0UL;
    bool bSuccess = false;
    ESP_LOGI(TAG_UPNP,"UPnP %s(%D)\n", routerIP.toString().c_str(), publicPort);

    // Step 1 - Discover UPnP service.
    if (discoverUPnP()) {
        // Step 2 - Check if mapping already exists
        if (!(bSuccess = requestPortMappingEntry())) {
            // Step 3 - Request new mapping
            bSuccess = requestPortForwarding();
        }
    } else {
        ESP_LOGI(TAG_UPNP,"    UPnP discovery failed.");
    }

    if (bSuccess) {
        // Extract external IP
        requestExternalIP();
        ESP_LOGI(TAG_UPNP,"    Public IP(Port): %s(%d)\n", publicIP.toString().c_str(), publicPort);
  

        // Extract external port assigned by UPnP
        // requestExternalPort();
        // ESP_LOGI(TAG_UPNP,"Assigned External Port: %d\n", publicPort);

        *pport = publicPort;  // External server port
        *pip   = publicIP;    // External IP address
        return true;
    } else {
        ESP_LOGI(TAG_UPNP,"    UPnP PortForwarding failed.");
    }
    return bSuccess;
}

bool CUpnpClient::discoverUPnP() {
    //ESP_LOGI(TAG_UPNP,"\n\n**  Sending SSDP discovery request...");
    WiFiUDP udp;
    bool ret = false;

    // SSDP M-SEARCH Request
    const char *ssdpRequest =
        "M-SEARCH * HTTP/1.1\r\n"
        "HOST: 239.255.255.250:1900\r\n"
        "MAN: \"ssdp:discover\"\r\n"
        "MX: 3\r\n"
        "ST: urn:schemas-upnp-org:device:InternetGatewayDevice:1\r\n"
        "\r\n";

    udp.beginMulticast(SSDP_MULTICAST_IP, SSDP_PORT);
    udp.beginPacket(SSDP_MULTICAST_IP, SSDP_PORT);
    udp.write((const uint8_t *)ssdpRequest, strlen(ssdpRequest)); // Fix length
    udp.endPacket();

    unsigned long startTime = millis();
    while (millis() - startTime < SEARCH_TIMEOUT) {
        int packetSize = udp.parsePacket();
        if (packetSize > 0) {
            udp.read(buffer, sizeof(buffer) - 1);
            buffer[packetSize] = '\0';

            //ESP_LOGI(TAG_UPNP,"SSDP response received:");
            //ESP_LOGI(TAG_UPNP,buffer);

            // Check if the response contains "ST: urn:schemas-upnp-org:device:InternetGatewayDevice:1"
            char *stField = strstr(buffer, "ST: urn:schemas-upnp-org:device:InternetGatewayDevice:1");
            if (stField) {
                // Extract router's UPnP service URL
                char *location = strstr(buffer, "LOCATION: ");
                if (location) {
                    location += 9; // Move past "LOCATION: "
                    while (*location == ' ') location++; // skip blanks
                    char *end = strchr(location, '\r');
                    if (end) {
                        *end = '\0';
                        //ESP_LOGI(TAG_UPNP,"Router UPnP URL: %s\n", location);
                        routerLocation = String(location);

                        // Extract IP and Port from the Location URL
                        int firstColonPos = routerLocation.indexOf(':');  // Find the first colon (for the protocol)
                        int secondColonPos = routerLocation.indexOf(':', firstColonPos + 1);  // Find the second colon (IP:PORT)

                        if (secondColonPos != -1) {
                            routerIPString = routerLocation.substring(routerLocation.indexOf("://") + 3, secondColonPos);  // Extract IP
                            routerPort = routerLocation.substring(secondColonPos + 1, routerLocation.indexOf('/', secondColonPos)).toInt();  // Extract Port
                        }

                        routerIP.fromString(routerIPString);
                        //ESP_LOGI(TAG_UPNP,"Router UPnP IP(Port): %s(%d)\n", routerIP.toString().c_str(), routerPort);

                        String xml = fetchUPnPDescription(routerLocation);
                        if (!xml.isEmpty()) {
                            // ESP_LOGI(TAG_UPNP,xml);
                            // ESP_LOGI(TAG_UPNP,"\n--- End Of XML ----\n");

                            // Parse the XML and get the controlURL
                            // String parseXML(const String &xml);
                            controlURL = parseXML(xml);
                            if (!controlURL.isEmpty()) {
                                ret = true;
                                break;
                            }
                        }
                    }
                }
            }
        }
    }
    udp.stop();

    if (!ret) {
        DPRINTLN("No SSDP response from UPnP device.");
    }
    return ret;
}

String CUpnpClient::fetchUPnPDescription(const String &location) {
    HTTPClient http;
    WiFiClient client;
    
    //ESP_LOGI(TAG_UPNP,"Fetching UPnP XML from: \"%s\"\n", location.c_str());

    http.begin(client, location);
    int httpCode = http.GET();

    if (httpCode > 0) {
        //ESP_LOGI(TAG_UPNP,"HTTP Response Code: %d\n", httpCode);
        if (httpCode == HTTP_CODE_OK) {
            String xmlContent = http.getString();
            http.end();
            return xmlContent;
        }
    } 
    else {
        //ESP_LOGI(TAG_UPNP,"HTTP Request failed, error: %s\n", http.errorToString(httpCode).c_str());
    }

    http.end();
    return "";
}

String CUpnpClient::parseXML(const String &xml) {
    // Locate the WANIPConnection service
    int servicePos = xml.indexOf("<serviceType>urn:schemas-upnp-org:service:WANIPConnection:1</serviceType>");
    if (servicePos == -1) {
        ESP_LOGI(TAG_UPNP,"WANIPConnection service not found in XML.");
        return "";
    }

    // Find the start of the controlURL tag within the <service> block
    int controlStart = xml.indexOf("<controlURL>", servicePos);
    if (controlStart == -1) {
        ESP_LOGI(TAG_UPNP,"controlURL not found in service block.");
        return "";
    }
    controlStart += 12;  // Move past "<controlURL>"

    // Find the closing tag
    int controlEnd = xml.indexOf("</controlURL>", controlStart);
    if (controlEnd == -1) {
        ESP_LOGI(TAG_UPNP,"Malformed XML: Missing </controlURL>.");
        return "";
    }

    // Extract and clean the controlURL
    String controlURL = xml.substring(controlStart, controlEnd);
    controlURL.trim();  // Remove any spaces or newlines

    //ESP_LOGI(TAG_UPNP,"Extracted controlURL: %s\n", controlURL.c_str());
    return controlURL;
}

int CUpnpClient::sendSoapRequest(const char *request, char *response, size_t responseSize) {
    WiFiClient client;
    if (!client.connect(routerIP, routerPort)) {  // UPnP uses port 1900 for SOAP requests
        ESP_LOGI(TAG_UPNP,"Failed to connect to router: %s\n", routerIP.toString().c_str());
        return -1;
    }

    client.print("POST /control?WANIPConn1 HTTP/1.1\r\n");
    client.print("Host: ");
    client.print(routerIP);
    client.print("\r\n");
    client.print("Content-Type: text/xml; charset=\"utf-8\"\r\n");
    client.print("SOAPAction: \"urn:schemas-upnp-org:service:WANIPConnection:1#GetSpecificPortMappingEntry\"\r\n");
    client.print("Content-Length: ");
    client.print(strlen(request));
    client.print("\r\n\r\n");
    client.print(request);

    unsigned long startMillis = millis();
    while (!client.available() && millis() - startMillis < 5000) {
        delay(10);  // Wait for response
    }

    int index = 0;
    while (client.available() && index < responseSize - 1) {
        response[index++] = client.read();
    }
    response[index] = '\0';

    client.stop();
    return index;
}

bool CUpnpClient::requestPortMappingEntry() {
    ESP_LOGI(TAG_UPNP,"**  Requesting UPnP Port Mapping Entry...");
    char soapRequest[512];
    char mappingName[32];
    sprintf(mappingName, "HC_%04X", publicPort);

    snprintf(soapRequest, sizeof(soapRequest),
        "<?xml version=\"1.0\"?>"
        "<s:Envelope xmlns:s=\"http://schemas.xmlsoap.org/soap/envelope/\" "
        "s:encodingStyle=\"http://schemas.xmlsoap.org/soap/encoding/\">"
        "<s:Body>"
        "<u:GetSpecificPortMappingEntry xmlns:u=\"urn:schemas-upnp-org:service:WANIPConnection:1\">"
        "<NewRemoteHost></NewRemoteHost>"
        "<NewExternalPort>%d</NewExternalPort>"
        "<NewProtocol>TCP</NewProtocol>"
        "<NewInternalClient>%s</NewInternalClient>"
        "<NewPortMappingDescription>%s</NewPortMappingDescription>"
        "</u:GetSpecificPortMappingEntry>"
        "</s:Body>"
        "</s:Envelope>", 
        publicPort,  
        WiFi.localIP().toString().c_str(), 
        mappingName);

    char response[1024];
    int responseLen = sendSoapRequest(soapRequest, response, sizeof(response));

    if (responseLen > 0) {
        if (strstr(response, mappingName)) {
            //ESP_LOGI(TAG_UPNP,"    Port mapping already exists.");
            return true;
        }
    }
    ESP_LOGI(TAG_UPNP,"    Port mapping not found.");
    return false;
}

bool CUpnpClient::requestPortForwarding() {
    // Sending port-forwarding request to the router's external IP or gateway
    ESP_LOGI(TAG_UPNP,"**  Requesting UPnP Port Forwarding...");

    if (controlURL.isEmpty()) {
        ESP_LOGI(TAG_UPNP,"Error: controlURL is empty.");
        return false;
    }

    // Ensure controlURL is correctly formatted (handle absolute/relative cases)
    String postURL = controlURL;
    if (postURL.startsWith("http://") || postURL.startsWith("https://")) {
        int pathStart = postURL.indexOf('/', 8);  // Skip "http://"
        if (pathStart != -1) {
            postURL = postURL.substring(pathStart);  // Extract path only
            ESP_LOGI(TAG_UPNP,"PostURL: \"%s\"\n", postURL.c_str());
        } else {
            ESP_LOGI(TAG_UPNP,"Invalid controlURL format.");
            return false;
        }
    }

    // XML body
    char xmlBody[800]; 
    snprintf(xmlBody, sizeof(xmlBody),
        "<?xml version=\"1.0\"?>"
        "<s:Envelope xmlns:s=\"http://schemas.xmlsoap.org/soap/envelope/\" s:encodingStyle=\"http://schemas.xmlsoap.org/soap/encoding/\">"
        "<s:Body><u:AddPortMapping xmlns:u=\"urn:schemas-upnp-org:service:WANIPConnection:1\">"
        "<NewRemoteHost></NewRemoteHost>"
        "<NewExternalPort>%d</NewExternalPort>"
        "<NewProtocol>TCP</NewProtocol>"
        "<NewInternalPort>%d</NewInternalPort>"
        "<NewInternalClient>%s</NewInternalClient>"
        "<NewEnabled>1</NewEnabled>"
        "<NewPortMappingDescription>HC_%04X</NewPortMappingDescription>"
        "<NewLeaseDuration>0</NewLeaseDuration>"
        "</u:AddPortMapping></s:Body></s:Envelope>",
        publicPort,  publicPort, 
        WiFi.localIP().toString().c_str(),
        publicPort);
    //ESP_LOGI(TAG_UPNP,"XML Body: ");
    //ESP_LOGI(TAG_UPNP,xmlBody);


    // Calculate the correct content length
    int contentLength = strlen(xmlBody);

    snprintf(buffer, sizeof(buffer),
        "POST %s HTTP/1.1\r\n"
        "Host: %s:%d\r\n"
        "Content-Type: text/xml; charset=\"utf-8\"\r\n"
        "SOAPAction: \"urn:schemas-upnp-org:service:WANIPConnection:1#AddPortMapping\"\r\n"
        "Content-Length: %d\r\n"
        "Connection: Close\r\n"
        "\r\n"
        "%s",  // Append XML body after headers
        controlURL.c_str(),
        routerIP.toString().c_str(), routerPort, 
        contentLength,
        xmlBody);

    //ESP_LOGI(TAG_UPNP,"==== HTTP Request ===\n\n%s\n\n--- End Of HTTP ---\n", buffer);

    // Connect and send
    WiFiClient client;
    if (!client.connect(routerIP, routerPort)) {
        ESP_LOGI(TAG_UPNP,"Failed to connect to the router. %s:%d\n", routerIP.toString().c_str(), routerPort);
        return false;
    }

    client.print(buffer);
    
    // Read response
    String response;
    uint32_t timeout = millis() + 5000;
    while (client.available() == 0) {
        if (millis() > timeout) {
            ESP_LOGI(TAG_UPNP,"Router did not respond to port mapping request.");
            client.stop();
            return false;
        }
    }

    while (client.available()) {
        response += client.readString();
    }
    client.stop();
    return true;
};

bool CUpnpClient::requestExternalIP() {
    //ESP_LOGI(TAG_UPNP,"\n\nRequesting External IP via UPnP...");

    // XML Request Body
    char xmlBody[300];
    snprintf(xmlBody, sizeof(xmlBody),
        "<?xml version=\"1.0\"?>"
        "<s:Envelope xmlns:s=\"http://schemas.xmlsoap.org/soap/envelope/\" s:encodingStyle=\"http://schemas.xmlsoap.org/soap/encoding/\">"
        "<s:Body>"
        "<u:GetExternalIPAddress xmlns:u=\"urn:schemas-upnp-org:service:WANIPConnection:1\"/>"
        "</s:Body>"
        "</s:Envelope>");

    int contentLength = strlen(xmlBody);

    // HTTP Request
    snprintf(buffer, sizeof(buffer),
        "POST %s HTTP/1.1\r\n"
        "Host: %s:%d\r\n"
        "Content-Type: text/xml; charset=\"utf-8\"\r\n"
        "SOAPAction: \"urn:schemas-upnp-org:service:WANIPConnection:1#GetExternalIPAddress\"\r\n"
        "Content-Length: %d\r\n"
        "Connection: Close\r\n"
        "\r\n"
        "%s",
        controlURL.c_str(),
        routerIP.toString().c_str(), routerPort,
        contentLength,
        xmlBody);

    // Connect and Send
    WiFiClient client;
    if (!client.connect(routerIP, routerPort)) {
        DPRINTLN("Failed to connect to router for External IP request.");
        return false;
    }

    client.print(buffer);

    // Read Response
    String response;
    uint32_t timeout = millis() + 5000;
    while (client.available() == 0) {
        if (millis() > timeout) {
            DPRINTLN("Router did not respond to External IP request.");
            client.stop();
            return false;
        }
    }

    while (client.available()) {
        response += client.readString();
    }
    client.stop();

    //ESP_LOGI(TAG_UPNP,"External IP Response:");
    //ESP_LOGI(TAG_UPNP,response);

    // Extract External IP from XML
    char *sz = (char *)(response.c_str());
    char *extIP = strstr(sz, "<NewExternalIPAddress>");
    if (extIP) {
        extIP += 22;
        char *end = strchr(extIP, '<');
        if (end) *end = '\0';

        publicIP = IPAddress(extIP);
        DPRINTF("UPnP - External IP: %s\n", extIP);
        return true;
    }

    DPRINTLN("Failed to extract External IP.");
    return false;
}

bool CUpnpClient::requestExternalPort() {
    //ESP_LOGI(TAG_UPNP,"\n\nRequesting External Port via UPnP...");

    // XML Request Body for getting External Port
    char xmlBody[300];
    snprintf(xmlBody, sizeof(xmlBody),
        "<?xml version=\"1.0\"?>"
        "<s:Envelope xmlns:s=\"http://schemas.xmlsoap.org/soap/envelope/\" s:encodingStyle=\"http://schemas.xmlsoap.org/soap/encoding/\">"
        "<s:Body>"
        "<u:GetExternalPort xmlns:u=\"urn:schemas-upnp-org:service:WANIPConnection:1\"/>"
        "</s:Body>"
        "</s:Envelope>");

    int contentLength = strlen(xmlBody);

    // HTTP Request to get External Port
    snprintf(buffer, sizeof(buffer),
        "POST %s HTTP/1.1\r\n"
        "Host: %s:%d\r\n"
        "Content-Type: text/xml; charset=\"utf-8\"\r\n"
        "SOAPAction: \"urn:schemas-upnp-org:service:WANIPConnection:1#GetExternalPort\"\r\n"
        "Content-Length: %d\r\n"
        "Connection: Close\r\n"
        "\r\n"
        "%s",
        controlURL.c_str(),
        routerIP.toString().c_str(), routerPort,
        contentLength,
        xmlBody);

    // Connect and Send
    WiFiClient client;
    if (!client.connect(routerIP, routerPort)) {
        ESP_LOGI(TAG_UPNP,"Failed to connect to router for External Port request.");
        return false;
    }

    client.print(buffer);

    // Read Response
    String response;
    uint32_t timeout = millis() + 5000;
    while (client.available() == 0) {
        if (millis() > timeout) {
            ESP_LOGI(TAG_UPNP,"Router did not respond to External Port request.");
            client.stop();
            return false;
        }
    }

    while (client.available()) {
        response += client.readString();
    }
    client.stop();

    //ESP_LOGI(TAG_UPNP,"External Port Response:");
    //ESP_LOGI(TAG_UPNP,response);

    // Extract External Port from XML
    char *sz = (char *)(response.c_str());
    char *extPort = strstr(sz, "<NewExternalPort>");
    if (extPort) {
        extPort += 17;
        char *end = strchr(extPort, '<');
        if (end) *end = '\0';

        publicPort = atoi(extPort);
        ESP_LOGI(TAG_UPNP,"External Port: %d\n", publicPort);
        return true;
    }

    ESP_LOGI(TAG_UPNP,"Failed to extract External Port.");
    return false;
};