// Copyright 2000-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
package com.intellij.openapi.actionSystem.impl;

import com.intellij.ide.DataManager;
import com.intellij.ide.impl.DataManagerImpl;
import com.intellij.ide.impl.dataRules.GetDataRule;
import com.intellij.openapi.actionSystem.CommonDataKeys;
import com.intellij.openapi.actionSystem.DataContext;
import com.intellij.openapi.actionSystem.PlatformDataKeys;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.wm.IdeFocusManager;
import org.jetbrains.annotations.NotNull;

import java.util.Map;

public final class SimpleDataContext implements DataContext {
  private final Map<String, Object> myDataId2Data;
  private final DataContext myParent;

  private SimpleDataContext(@NotNull Map<String, Object> dataId2data, DataContext parent) {
    myDataId2Data = dataId2data;
    myParent = parent;
  }

  @Override
  public Object getData(@NotNull String dataId) {
    Object result = getDataFromSelfOrParent(dataId);

    if (result == null && PlatformDataKeys.CONTEXT_COMPONENT.is(dataId)) {
      result = IdeFocusManager.getGlobalInstance().getFocusOwner();
    }

    if (result == null) {
      GetDataRule rule = ((DataManagerImpl)DataManager.getInstance()).getDataRule(dataId);
      if (rule != null) {
        return rule.getData(this::getDataFromSelfOrParent);
      }
    }

    return result;
  }

  private Object getDataFromSelfOrParent(@NotNull String dataId) {
    return myDataId2Data.containsKey(dataId) ? myDataId2Data.get(dataId) :
           myParent == null ? null : myParent.getData(dataId);
  }

  @NotNull
  public static DataContext getSimpleContext(@NotNull String dataId, @NotNull Object data, DataContext parent) {
    return new SimpleDataContext(Map.of(dataId, data), parent);
  }

  @NotNull
  public static DataContext getSimpleContext(@NotNull Map<String,Object> dataId2data, DataContext parent) {
    return new SimpleDataContext(dataId2data, parent);
  }

  @NotNull
  public static DataContext getSimpleContext(@NotNull String dataId, @NotNull Object data) {
    return getSimpleContext(dataId, data, null);
  }

  @NotNull
  public static DataContext getProjectContext(@NotNull Project project) {
    return getSimpleContext(CommonDataKeys.PROJECT.getName(), project);
  }
}

@@ -429,16 +429,66 @@
 ontext;%0A
+import com.intellij.openapi.actionSystem.DataKey;%0A
 import c
@@ -673,16 +673,107 @@
 otNull;%0A
+import org.jetbrains.annotations.Nullable;%0A%0Aimport javax.swing.*;%0Aimport java.util.HashMap;
 %0Aimport 
@@ -1803,32 +1803,153 @@
 a(dataId);%0A  %7D%0A%0A
+  /** @deprecated use %7B@link SimpleDataContext#getSimpleContext(DataKey, Object, DataContext)%7D instead. */%0A  @Deprecated%0A
   @NotNull%0A  pub
@@ -2117,32 +2117,248 @@
 , parent);%0A  %7D%0A%0A
+  @NotNull%0A  public static %3CT%3E DataContext getSimpleContext(@NotNull DataKey%3C? super T%3E dataKey, @NotNull T data, DataContext parent) %7B%0A    return new SimpleDataContext(Map.of(dataKey.getName(), data), parent);%0A  %7D%0A%0A
   @NotNull%0A  pub
@@ -2512,32 +2512,140 @@
 , parent);%0A  %7D%0A%0A
+  /** @deprecated use %7B@link SimpleDataContext#getSimpleContext(DataKey, Object)%7D instead. */%0A  @Deprecated%0A
   @NotNull%0A  pub
@@ -2817,137 +2817,1174 @@
 tic 
-DataContext getProjectContext(@NotNull Project project) %7B%0A    return getSimpleContext(CommonDataKeys.PROJECT.getName(), project);
+%3CT%3E DataContext getSimpleContext(@NotNull DataKey%3C? super T%3E dataKey, @NotNull T data) %7B%0A    return getSimpleContext(dataKey, data, null);%0A  %7D%0A%0A  @NotNull%0A  public static DataContext getProjectContext(@NotNull Project project) %7B%0A    return getSimpleContext(CommonDataKeys.PROJECT.getName(), project);%0A  %7D%0A%0A  @NotNull%0A  public static Builder builder() %7B%0A    return new Builder(null);%0A  %7D%0A%0A  @NotNull%0A  public static Builder builder(@Nullable DataContext parent) %7B%0A    return new Builder(parent);%0A  %7D%0A%0A  public final static class Builder %7B%0A    private final DataContext myParent;%0A    private Map%3CString, Object%3E myMap;%0A%0A    Builder(DataContext parent) %7B%0A      myParent = parent;%0A    %7D%0A%0A    @NotNull%0A    public %3CT%3E Builder add(@NotNull DataKey%3C? super T%3E dataKey, @Nullable T value) %7B%0A      if (value != null) %7B%0A        if (myMap == null) myMap = new HashMap%3C%3E();%0A        myMap.put(dataKey.getName(), value);%0A      %7D%0A      return this;%0A    %7D%0A%0A    @NotNull%0A    public Builder noUI() %7B%0A      return add(PlatformDataKeys.CONTEXT_COMPONENT, new JLabel());%0A    %7D%0A%0A    @NotNull%0A    public SimpleDataContext build() %7B%0A      return new SimpleDataContext(myMap, myParent);%0A    %7D
 %0A  %7D

introduce SimpleDataContext#builder