// 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.remote.internal;

import com.google.common.base.Function;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;

import org.openqa.selenium.WebElement;
import org.openqa.selenium.remote.Dialect;
import org.openqa.selenium.remote.RemoteWebDriver;
import org.openqa.selenium.remote.RemoteWebElement;

import java.util.Collection;
import java.util.Map;

/**
 * Reconstitutes {@link WebElement}s from their JSON representation. Will recursively convert Lists
 * and Maps to catch nested references. All other values pass through the converter unchanged.
 */
public class JsonToWebElementConverter implements Function<Object, Object> {

  private final RemoteWebDriver driver;

  public JsonToWebElementConverter(RemoteWebDriver driver) {
    this.driver = driver;
  }

  public Object apply(Object result) {
    if (result instanceof Collection<?>) {
      Collection<?> results = (Collection<?>) result;
      return Lists.newArrayList(Iterables.transform(results, this));
    }

    if (result instanceof Map<?, ?>) {
      Map<?, ?> resultAsMap = (Map<?, ?>) result;
      if (resultAsMap.containsKey(Dialect.OSS.getEncodedElementKey())) {
        RemoteWebElement element = newRemoteWebElement();
        element.setId(String.valueOf(resultAsMap.get(Dialect.OSS.getEncodedElementKey())));
        return element;
      } else if (resultAsMap.containsKey(Dialect.W3C.getEncodedElementKey())) {
        RemoteWebElement element = newRemoteWebElement();
        element.setId(String.valueOf(resultAsMap.get(Dialect.W3C.getEncodedElementKey())));
        return element;
      } else {
        return Maps.transformValues(resultAsMap, this);
      }
    }

    if (result instanceof RemoteWebElement) {
      return setOwner((RemoteWebElement) result);
    }

    if (result instanceof Number) {
      if (result instanceof Float || result instanceof Double) {
        return ((Number) result).doubleValue();
      }
      return ((Number) result).longValue();
    }

    return result;
  }

  private RemoteWebElement newRemoteWebElement() {
    return setOwner(new RemoteWebElement());
  }

  private RemoteWebElement setOwner(RemoteWebElement element) {
    if (driver != null) {
      element.setParent(driver);
      element.setFileDetector(driver.getFileDetector());
    }
    return element;
  }
}

@@ -1970,313 +1970,71 @@
     
-if (resultAsMap.containsKey(Dialect.OSS.getEncodedElementKey())) %7B%0A        RemoteWebElement element = newRemoteWebElement();%0A        element.setId(String.valueOf(resultAsMap.get(Dialect.OSS.getEncodedElementKey())));%0A        return element;%0A      %7D else if (resultAsMap.containsKey(Dialect.W3C.getEncodedE
+String elementKey = getElementKey(resultAsMap);%0A%09%09  if (null!=e
 leme
@@ -2042,21 +2042,15 @@
 tKey
-())
 ) %7B%0A
-      
+%09%09%09
   Re
@@ -2097,22 +2097,19 @@
 ment();%0A
-      
+%09%09%09
   elemen
@@ -2151,31 +2151,9 @@
 get(
-Dialect.W3C.getEncodedE
+e
 leme
@@ -2161,21 +2161,16 @@
 tKey
-()
 )));%0A
-      
+%09%09%09
   re
@@ -2600,37 +2600,39 @@
 esult;%0A  %7D%0A%0A  pr
-ivate
+otected
  RemoteWebElemen
@@ -2913,14 +2913,291 @@
 lement;%0A
-
   %7D%0A
+  private static String getElementKey(Map%3C?, ?%3E resultAsMap) %7B%0A%09%09for (Dialect d : Dialect.values()) %7B%0A%09%09%09String elementKeyForDialect = d.getEncodedElementKey();%0A%09%09%09if (resultAsMap.containsKey(elementKeyForDialect)) %7B%0A%09%09%09%09return elementKeyForDialect;%0A%09%09%09%7D%0A%09%09%7D%0A%09%09return null;%0A%09%7D%0A
 %7D%0A

Make `newRemoteWebElement` `protected`