package com.example;

import javax.xml.XMLConstants;
import javax.xml.namespace.NamespaceContext;
import java.util.*;

/**
 * Minimal implementation of {@link NamespaceContext}.
 */
public class NamespaceContextImpl implements NamespaceContext {
    private final Map<String, String> prefixToUri = new HashMap<>();
    private final Map<String, Set<String>> uriToPrefixes = new HashMap<>();

    public NamespaceContextImpl(String... prefixNamespaceUriPairs) {
        if (prefixNamespaceUriPairs == null || prefixNamespaceUriPairs.length % 2 != 0) {
            throw new IllegalArgumentException("Must provide prefix/URI pairs.");
        }
        for (int i = 0; i < prefixNamespaceUriPairs.length; i += 2) {
            String prefix = prefixNamespaceUriPairs[i];
            String uri = prefixNamespaceUriPairs[i + 1];
            prefixToUri.put(prefix, uri);
            uriToPrefixes.computeIfAbsent(uri, k -> new HashSet<>()).add(prefix);
        }
    }

    @Override
    public String getNamespaceURI(String prefix) {
        if (prefix == null) throw new IllegalArgumentException("Prefix cannot be null.");
        if (prefix.equals(XMLConstants.XML_NS_PREFIX)) return XMLConstants.XML_NS_URI;
        if (prefix.equals(XMLConstants.XMLNS_ATTRIBUTE)) return XMLConstants.XMLNS_ATTRIBUTE_NS_URI;
        return prefixToUri.getOrDefault(prefix, XMLConstants.NULL_NS_URI);
    }

    @Override
    public String getPrefix(String namespaceURI) {
        if (namespaceURI == null) throw new IllegalArgumentException("NamespaceURI cannot be null.");
        if (namespaceURI.equals(XMLConstants.XML_NS_URI)) return XMLConstants.XML_NS_PREFIX;
        if (namespaceURI.equals(XMLConstants.XMLNS_ATTRIBUTE_NS_URI)) return XMLConstants.XMLNS_ATTRIBUTE;
        Set<String> prefixes = uriToPrefixes.get(namespaceURI);
        return (prefixes == null || prefixes.isEmpty()) ? null : prefixes.iterator().next();
    }

    @Override
    public Iterator<String> getPrefixes(String namespaceURI) {
        if (namespaceURI == null) throw new IllegalArgumentException("NamespaceURI cannot be null.");
        if (namespaceURI.equals(XMLConstants.XML_NS_URI)) {
            return Collections.singleton(XMLConstants.XML_NS_PREFIX).iterator();
        }
        if (namespaceURI.equals(XMLConstants.XMLNS_ATTRIBUTE_NS_URI)) {
            return Collections.singleton(XMLConstants.XMLNS_ATTRIBUTE).iterator();
        }
        Set<String> prefixes = uriToPrefixes.get(namespaceURI);
        return prefixes == null ? Collections.emptyIterator() : Collections.unmodifiableSet(prefixes).iterator();
    }
}

package com.example;

import static java.util.Collections.unmodifiableCollection;
import static java.util.Objects.requireNonNull;
import static javax.xml.XMLConstants.NULL_NS_URI;
import static javax.xml.XMLConstants.XMLNS_ATTRIBUTE;
import static javax.xml.XMLConstants.XMLNS_ATTRIBUTE_NS_URI;
import static javax.xml.XMLConstants.XML_NS_PREFIX;
import static javax.xml.XMLConstants.XML_NS_URI;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import javax.xml.namespace.NamespaceContext;

public class NamespaceContextImpl implements NamespaceContext {

  private final List<String> prefixes;
  private final Map<String, String> prefixToNamespaceUri;
  private final Map<String, String> namespaceUriToPrefix;

  public NamespaceContextImpl(String... prefixNamespaceUriPairs) {
    if (prefixNamespaceUriPairs.length % 2 != 0) {
      throw new IllegalArgumentException("prefixNamespaceUriPairs must be even");
    }

    var capacity = prefixNamespaceUriPairs.length / 2;

    prefixes = new ArrayList<>(capacity);
    prefixToNamespaceUri = HashMap.newHashMap(capacity);
    namespaceUriToPrefix = HashMap.newHashMap(capacity);

    for (int i = 0; i < prefixNamespaceUriPairs.length; i += 2) {
      var prefix = prefixNamespaceUriPairs[i];
      var namespaceUri = prefixNamespaceUriPairs[i + 1];

      requireNonNull(prefix, "prefix");
      requireNonNull(namespaceUri, "namespaceUri");

      prefixes.add(prefix);

      if (prefixToNamespaceUri.put(prefix, namespaceUri) != null) {
        throw new IllegalArgumentException("duplicate prefix: " + prefix);
      }

      var oldPrefix = namespaceUriToPrefix.put(namespaceUri, prefix);
      if (oldPrefix != null && oldPrefix.isEmpty()) {
        namespaceUriToPrefix.put(namespaceUri, oldPrefix);
      }
    }
  }

  @Override
  public String getNamespaceURI(String prefix) {
    return switch (prefix) {
      case null -> throw new IllegalArgumentException("prefix is null");
      case XML_NS_PREFIX -> XML_NS_URI;
      case XMLNS_ATTRIBUTE -> XMLNS_ATTRIBUTE_NS_URI;
      default -> prefixToNamespaceUri.getOrDefault(prefix, NULL_NS_URI);
    };
  }

  @Override
  public String getPrefix(String namespaceUri) {
    return switch (namespaceUri) {
      case null -> throw new IllegalArgumentException("namespaceUri is null");
      case XML_NS_URI -> XML_NS_PREFIX;
      case XMLNS_ATTRIBUTE_NS_URI -> XMLNS_ATTRIBUTE;
      default -> namespaceUriToPrefix.getOrDefault(namespaceUri, null);
    };
  }

  @Override
  public Iterator<String> getPrefixes(String namespaceUri) {
    return unmodifiableCollection(prefixes).iterator();
  }

}