// Licensed to the Software Freedom Conservancy (SFC) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The SFC licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

package org.openqa.selenium.grid.node.config;

import com.google.common.collect.ImmutableMap;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.function.Function;
import org.openqa.selenium.Capabilities;
import org.openqa.selenium.ImmutableCapabilities;
import org.openqa.selenium.PersistentCapabilities;

public class SessionCapabilitiesMutator implements Function<Capabilities, Capabilities> {

  private static final ImmutableMap<String, String> BROWSER_OPTIONS =
      ImmutableMap.of(
          "chrome", "goog:chromeOptions",
          "firefox", "moz:firefoxOptions",
          "microsoftedge", "ms:edgeOptions");
  private static final String SE_VNC_ENABLED = "se:vncEnabled";
  private static final String SE_NO_VNC_PORT = "se:noVncPort";
  private final Capabilities slotStereotype;

  public SessionCapabilitiesMutator(Capabilities slotStereotype) {
    this.slotStereotype = slotStereotype;
  }

  @Override
  public Capabilities apply(Capabilities capabilities) {
    if (!Objects.equals(slotStereotype.getBrowserName(), capabilities.getBrowserName())) {
      return capabilities;
    }

    if (slotStereotype.getCapability(SE_VNC_ENABLED) != null) {
      capabilities =
          new PersistentCapabilities(capabilities)
              .setCapability(SE_VNC_ENABLED, slotStereotype.getCapability(SE_VNC_ENABLED))
              .setCapability(SE_NO_VNC_PORT, slotStereotype.getCapability(SE_NO_VNC_PORT));
    }

    String browserName = capabilities.getBrowserName().toLowerCase();

    if ("internet explorer".equalsIgnoreCase(browserName)) {
      return new ImmutableCapabilities(removeUnknownExtensionsForIE(capabilities));
    }

    if (!BROWSER_OPTIONS.containsKey(browserName)) {
      return capabilities;
    }

    String options = BROWSER_OPTIONS.get(browserName);
    if (slotStereotype.asMap().containsKey(options) && capabilities.asMap().containsKey(options)) {

      @SuppressWarnings("unchecked")
      Map<String, Object> stereotypeOptions =
          new HashMap<>((Map<String, Object>) slotStereotype.asMap().get(options));

      @SuppressWarnings("unchecked")
      Map<String, Object> capsOptions =
          new HashMap<>((Map<String, Object>) capabilities.asMap().get(options));

      // Merge top level capabilities, excluding browser specific options.
      // This will overwrite the browser options too, but it does not matter since we tackle it
      // separately just after this.
      Map<String, Object> toReturn = new HashMap<>(slotStereotype.merge(capabilities).asMap());

      // Merge browser specific stereotype and capabilities options
      switch (browserName.toLowerCase()) {
        case "chrome":
        case "microsoftedge":
        case "msedge":
          toReturn.put(options, mergeChromiumOptions(stereotypeOptions, capsOptions));
          break;
        case "firefox":
          toReturn.put(options, mergeFirefoxOptions(stereotypeOptions, capsOptions));
          break;
        default:
          break;
      }

      return new ImmutableCapabilities(toReturn);
    }

    return slotStereotype.merge(capabilities);
  }

  private Map<String, Object> removeUnknownExtensionsForIE(Capabilities capabilities) {
    Map<String, Object> toReturn = new HashMap<>(capabilities.asMap());
    capabilities.asMap().keySet().stream()
        .filter(key -> key.contains(":"))
        .filter(key -> !"se:ieOptions".equalsIgnoreCase(key))
        .forEach(toReturn::remove);
    return toReturn;
  }

  private Map<String, Object> mergeChromiumOptions(
      Map<String, Object> stereotypeOptions, Map<String, Object> capsOptions) {
    Map<String, Object> toReturn = new HashMap<>(stereotypeOptions);

    for (Map.Entry<String, Object> entry : capsOptions.entrySet()) {
      String name = entry.getKey();
      Object value = entry.getValue();
      if (name.equals("args")) {
        @SuppressWarnings("unchecked")
        List<String> arguments = new ArrayList<>((List<String>) value);

        @SuppressWarnings("unchecked")
        List<String> stereotypeArguments =
            new ArrayList<>(
                (List<String>) (stereotypeOptions.getOrDefault(("args"), new ArrayList<>())));

        arguments.forEach(
            arg -> {
              if (!stereotypeArguments.contains(arg)) {
                stereotypeArguments.add(arg);
              }
            });
        toReturn.put("args", stereotypeArguments);
      }

      if (name.equals("extensions")) {
        @SuppressWarnings("unchecked")
        List<String> extensionList = new ArrayList<>((List<String>) value);

        @SuppressWarnings("unchecked")
        List<String> stereotypeExtensions =
            new ArrayList<>(
                (List<String>) (stereotypeOptions.getOrDefault(("extensions"), new ArrayList<>())));

        extensionList.forEach(
            extension -> {
              if (!stereotypeExtensions.contains(extension)) {
                stereotypeExtensions.add(extension);
              }
            });

        toReturn.put("extensions", stereotypeExtensions);
      }

      if (name.equals("binary") && !stereotypeOptions.containsKey("binary")) {
        toReturn.put(name, value);
      }

      if (!name.equals("binary") && !name.equals("extensions") && !name.equals("args")) {
        toReturn.put(name, value);
      }
    }

    return toReturn;
  }

  private Map<String, Object> mergeFirefoxOptions(
      Map<String, Object> stereotypeOptions, Map<String, Object> capsOptions) {
    Map<String, Object> toReturn = new HashMap<>(stereotypeOptions);

    List<String> handledOptions = List.of("args", "prefs", "profile", "log", "binary");

    for (Map.Entry<String, Object> entry : capsOptions.entrySet()) {
      String name = entry.getKey();
      Object value = entry.getValue();
      if (name.equals("args")) {
        @SuppressWarnings("unchecked")
        List<String> arguments = new ArrayList<>((List<String>) value);

        @SuppressWarnings("unchecked")
        List<String> stereotypeArguments =
            new ArrayList<>(
                (List<String>) (stereotypeOptions.getOrDefault(("args"), new ArrayList<>())));

        arguments.forEach(
            arg -> {
              if (!stereotypeArguments.contains(arg)) {
                stereotypeArguments.add(arg);
              }
            });
        toReturn.put("args", stereotypeArguments);
      }

      if (name.equals("prefs")) {
        @SuppressWarnings("unchecked")
        Map<String, Object> prefs = new HashMap<>((Map<String, Object>) value);

        @SuppressWarnings("unchecked")
        Map<String, Object> stereotypePrefs =
            new HashMap<>(
                (Map<String, Object>) (stereotypeOptions.getOrDefault(("prefs"), new HashMap<>())));

        stereotypePrefs.putAll(prefs);
        toReturn.put("prefs", stereotypePrefs);
      }

      if (name.equals("profile")) {
        String rawProfile = (String) value;
        toReturn.put("profile", rawProfile);
      }

      if (name.equals("log")) {
        @SuppressWarnings("unchecked")
        Map<String, Object> logLevelMap = (Map<String, Object>) value;
        toReturn.put("log", logLevelMap);
      }

      if (name.equals("binary") && !stereotypeOptions.containsKey("binary")) {
        toReturn.put(name, value);
      }

      if (!handledOptions.contains(name)) {
        toReturn.put(name, value);
      }
    }

    return toReturn;
  }
}
