content: package org.aliyun.gsl_client.parser.schema;
  
  import org.aliyun.gsl_client.exception.UserException;
  import org.aliyun.gsl_client.status.ErrorCode;
  import org.json.JSONArray;
  import org.json.JSONObject;
  
  import java.io.IOException;
  import java.nio.file.Files;
  import java.nio.file.Paths;
  import java.util.*;
  import java.util.stream.Collectors;
  
  public class Schema {
      private Map<String, AttrDef> nameToAttrDef;
      private Map<Integer, AttrDef> idToAttrDef;
      private Map<String, TypeDef> nameToTypeDef;
      private Map<Integer, TypeDef> idToTypeDef;
      private Map<String, EdgeRelation> edgeNameToRelation;
      private Map<Integer, EdgeRelation> edgeIdToRelation;
  
      public AttrDef getAttrDef(String attrName) throws UserException {
          if (!this.nameToAttrDef.containsKey(attrName)) {
              throw new UserException(ErrorCode.PARSE_ERROR, "Missing definition of attr name: " + attrName);
          }
          return this.nameToAttrDef.get(attrName);
      }
  
      public AttrDef getAttrDef(int attrTypeId) throws UserException {
          if (!this.idToAttrDef.containsKey(attrTypeId)) {
              throw new UserException(ErrorCode.PARSE_ERROR, "Missing definition of attr type: " + attrTypeId);
          }
          return this.idToAttrDef.get(attrTypeId);
      }
  
      public List<AttrDef> getAllAttrDefinitions() {
          return new ArrayList<>(this.nameToAttrDef.values());
      }
  
      public TypeDef getTypeDef(String typeName) throws UserException {
          if (!this.nameToTypeDef.containsKey(typeName)) {
              throw new UserException(ErrorCode.PARSE_ERROR, "Missing definition of type name: " + typeName);
          }
          return this.nameToTypeDef.get(typeName);
      }
  
      public TypeDef getTypeDef(int typeId) throws UserException {
          if (!this.idToTypeDef.containsKey(typeId)) {
              throw new UserException(ErrorCode.PARSE_ERROR, "Missing definition of type: " + typeId);
          }
          return this.idToTypeDef.get(typeId);
      }
  
      public List<TypeDef> getAllTypeDefinitions() {
          return new ArrayList<>(this.idToTypeDef.values());
      }
  
      public EdgeRelation getRelation(String edgeName) throws UserException {
          if (!this.edgeNameToRelation.containsKey(edgeName)) {
              throw new UserException(ErrorCode.PARSE_ERROR, "Missing relation of edge name: " + edgeName);
          }
          return this.edgeNameToRelation.get(edgeName);
      }
  
      public EdgeRelation getRelation(int edgeTypeId) throws UserException {
          if (!this.edgeIdToRelation.containsKey(edgeTypeId)) {
              throw new UserException(ErrorCode.PARSE_ERROR, "Missing relation of edge type: " + edgeTypeId);
          }
          return this.edgeIdToRelation.get(edgeTypeId);
      }
  
      public List<EdgeRelation> getAllEdgeRelations() {
          return new ArrayList<>(this.edgeIdToRelation.values());
      }
  
      public static Schema parseFrom(String jsonFile) throws UserException, IOException {
          byte[] content = Files.readAllBytes(Paths.get(jsonFile));
          return parseFrom(content);
      }
  
      public static Schema parseFrom(byte[] content) throws UserException {
          Map<String, AttrDef> nameToAttrDef = new HashMap<>();
          Map<Integer, AttrDef> idToAttrDef = new HashMap<>();
          Map<String, TypeDef> nameToTypeDef = new HashMap<>();
          Map<Integer, TypeDef> idToTypeDef = new HashMap<>();
          Map<String, EdgeRelation> edgeNameToRelation = new HashMap<>();
          Map<Integer, EdgeRelation> edgeIdToRelation = new HashMap<>();
  
          JSONObject obj = new JSONObject(new String(content));
  
          JSONArray jsonAttrs = obj.getJSONArray("attr_defs");
          for (int i = 0; i < jsonAttrs.length(); ++i) {
              JSONObject jsonAttr = jsonAttrs.getJSONObject(i);
              int attrTypeId = jsonAttr.getInt("type");
              String attrName = jsonAttr.getString("name");
              DataType dataType = DataType.parseString(jsonAttr.getString("value_type"));
              AttrDef attrDef = new AttrDef(attrTypeId, attrName, dataType);
              nameToAttrDef.put(attrName, attrDef);
              idToAttrDef.put(attrTypeId, attrDef);
          }
  
          JSONArray jsonVertices = obj.getJSONArray("vertex_defs");
          for (int i = 0; i < jsonVertices.length(); ++i) {
              JSONObject jsonVertex = jsonVertices.getJSONObject(i);
              JSONArray jsonAttrTypes = jsonVertex.getJSONArray("attr_types");
              List<AttrDef> vertexAttrs = new ArrayList<>();
              for (int j = 0; j < jsonAttrTypes.length(); ++j) {
                  vertexAttrs.add(idToAttrDef.get(jsonAttrTypes.getInt(j)));
              }
              int vertexTypeId = jsonVertex.getInt("vtype");
              String vertexName = jsonVertex.getString("name");
              TypeDef vertexTypeDef = new TypeDef(TypeKind.VERTEX, vertexTypeId, vertexName, vertexAttrs);
              nameToTypeDef.put(vertexName, vertexTypeDef);
              idToTypeDef.put(vertexTypeId, vertexTypeDef);
          }
  
          JSONArray jsonEdges = obj.getJSONArray("edge_defs");
          for (int i = 0; i < jsonEdges.length(); ++i) {
              JSONObject jsonEdge = jsonEdges.getJSONObject(i);
              JSONArray jsonAttrTypes = jsonEdge.getJSONArray("attr_types");
              List<AttrDef> edgeAttrs = new ArrayList<>();
              for (int j = 0; j < jsonAttrTypes.length(); ++j) {
                  edgeAttrs.add(idToAttrDef.get(jsonAttrTypes.getInt(j)));
              }
              int edgeTypeId = jsonEdge.getInt("etype");
              String edgeName = jsonEdge.getString("name");
              TypeDef edgeTypeDef = new TypeDef(TypeKind.EDGE, edgeTypeId, edgeName, edgeAttrs);
              nameToTypeDef.put(edgeName, edgeTypeDef);
              idToTypeDef.put(edgeTypeId, edgeTypeDef);
          }
  
          JSONArray jsonRelations = obj.getJSONArray("edge_relation_defs");
          for (int i = 0; i < jsonRelations.length(); ++i) {
              JSONObject jsonRelation = jsonRelations.getJSONObject(i);
              int edgeTypeId = jsonRelation.getInt("etype");
              String edgeName = idToTypeDef.get(edgeTypeId).getName();
              int srcVertexTypeId = jsonRelation.getInt("src_vtype");
              String srcVertexName = idToTypeDef.get(srcVertexTypeId).getName();
              int dstVertexTypeId = jsonRelation.getInt("dst_vtype");
              String dstVertexName = idToTypeDef.get(dstVertexTypeId).getName();
              EdgeRelation edgeRelation = new EdgeRelation(
                      edgeTypeId, edgeName, srcVertexTypeId, srcVertexName, dstVertexTypeId, dstVertexName);
              edgeNameToRelation.put(edgeName, edgeRelation);
              edgeIdToRelation.put(edgeTypeId, edgeRelation);
          }
  
          return new Schema(nameToAttrDef, idToAttrDef, nameToTypeDef, idToTypeDef, edgeNameToRelation, edgeIdToRelation);
      }
  
      public void dumpTo(String file) throws IOException {
          JSONArray jsonAttrs = new JSONArray();
          for (AttrDef attrDef : this.getAllAttrDefinitions()) {
              JSONObject jsonAttr = new JSONObject();
              jsonAttr.put("type", attrDef.getTypeId());
              jsonAttr.put("name", attrDef.getName());
              jsonAttr.put("value_type", attrDef.getDataType().toString());
              jsonAttrs.put(jsonAttr);
          }
  
          JSONArray jsonVertices = new JSONArray();
          JSONArray jsonEdges = new JSONArray();
          for (TypeDef typeDef : this.getAllTypeDefinitions()) {
              JSONArray jsonAttrTypes = new JSONArray();
              for (AttrDef attrDef : typeDef.getAttributes()) {
                  jsonAttrTypes.put(attrDef.getTypeId());
              }
              if (typeDef.getTypeKind() == TypeKind.VERTEX) {
                  JSONObject jsonVertex = new JSONObject();
                  jsonVertex.put("vtype", typeDef.getTypeId());
                  jsonVertex.put("name", typeDef.getName());
                  jsonVertex.put("attr_types", jsonAttrTypes);
                  jsonVertices.put(jsonVertex);
              } else {
                  JSONObject jsonEdge = new JSONObject();
                  jsonEdge.put("etype", typeDef.getTypeId());
                  jsonEdge.put("name", typeDef.getName());
                  jsonEdge.put("attr_types", jsonAttrTypes);
                  jsonEdges.put(jsonEdge);
              }
          }
  
          JSONArray jsonRelations = new JSONArray();
          for (EdgeRelation edgeRelation : this.getAllEdgeRelations()) {
              JSONObject jsonRelation = new JSONObject();
              jsonRelation.put("etype", edgeRelation.getEdgeTypeId());
              jsonRelation.put("src_vtype", edgeRelation.getSrcTypeId());
              jsonRelation.put("dst_vtype", edgeRelation.getDstTypeId());
              jsonRelations.put(jsonRelation);
          }
  
          JSONObject obj = new JSONObject();
          obj.put("attr_defs", jsonAttrs);
          obj.put("vertex_defs", jsonVertices);
          obj.put("edge_defs", jsonEdges);
          obj.put("edge_relation_defs", jsonRelations);
  
          Files.write(Paths.get(file), obj.toString().getBytes());
      }
  
      private Schema(Map<String, AttrDef> nameToAttrDef,
                     Map<Integer, AttrDef> idToAttrDef,
                     Map<String, TypeDef> nameToTypeDef,
                     Map<Integer, TypeDef> idToTypeDef,
                     Map<String, EdgeRelation> edgeNameToRelation,
                     Map<Integer, EdgeRelation> edgeIdToRelation) {
          this.nameToAttrDef = new HashMap<>(nameToAttrDef);
          this.idToAttrDef = new HashMap<>(idToAttrDef);
          this.nameToTypeDef = new HashMap<>(nameToTypeDef);
          this.idToTypeDef = new HashMap<>(idToTypeDef);
          this.edgeNameToRelation = new HashMap<>(edgeNameToRelation);
          this.edgeIdToRelation = new HashMap<>(edgeIdToRelation);
      }
  
      public static Builder newBuilder() {
          return new Builder();
      }
  
      public static Builder newBuilder(Schema schema) {
          return new Builder(schema);
      }
  
      public static class Builder {
          private Map<String, DataType> attrToDataType;
          private Map<String, Set<String>> vertexToAttr;
          private Map<String, Set<String>> edgeToAttr;
          private Map<String, RawEdgeRelation> edgeToRawRelation;
  
          private Builder() {
              this.attrToDataType = new HashMap<>();
              this.vertexToAttr = new HashMap<>();
              this.edgeToAttr = new HashMap<>();
              this.edgeToRawRelation = new HashMap<>();
          }
  
          private Builder(Schema schema) {
              this.attrToDataType = new HashMap<>();
              this.vertexToAttr = new HashMap<>();
              this.edgeToAttr = new HashMap<>();
              this.edgeToRawRelation = new HashMap<>();
              for (AttrDef attrDef : schema.getAllAttrDefinitions()) {
                  this.attrToDataType.put(attrDef.getName(), attrDef.getDataType());
              }
              for (TypeDef typeDef : schema.getAllTypeDefinitions()) {
                  Set<String> attrNames =
                          typeDef.getAttributes().stream().map(AttrDef::getName).collect(Collectors.toSet());
                  if (typeDef.getTypeKind() == TypeKind.VERTEX) {
                      this.vertexToAttr.put(typeDef.getName(), attrNames);
                  } else {
                      this.edgeToAttr.put(typeDef.getName(), attrNames);
                  }
              }
              for (EdgeRelation edgeRelation : schema.getAllEdgeRelations()) {
                  this.edgeToRawRelation.put(edgeRelation.getEdgeName(), new RawEdgeRelation(
                          edgeRelation.getEdgeName(), edgeRelation.getSrcName(), edgeRelation.getDstName()));
              }
          }
  
          public Builder addAttribute(String attrName, DataType dataType) {
              this.attrToDataType.put(attrName, dataType);
              return this;
          }
  
          public Builder removeAttribute(String attrName) {
              this.attrToDataType.remove(attrName);
              return this;
          }
  
          public Builder addVertex(String vertexName, List<String> vertexAttrNames) {
              this.vertexToAttr.put(vertexName, new HashSet<>(vertexAttrNames));
              return this;
          }
  
          public Builder removeVertex(String vertexName) {
              this.vertexToAttr.remove(vertexName);
              return this;
          }
  
          public Builder addVertexAttribute(String vertexName, String attrName) {
              this.vertexToAttr.get(vertexName).add(attrName);
              return this;
          }
  
          public Builder removeVertexAttribute(String vertexName, String attrName) {
              this.vertexToAttr.get(vertexName).remove(attrName);
              return this;
          }
  
          public Builder addEdge(String edgeName, String srcName, String dstName, List<String> edgeAttrNames) {
              this.edgeToAttr.put(edgeName, new HashSet<>(edgeAttrNames));
              this.edgeToRawRelation.put(edgeName, new RawEdgeRelation(edgeName, srcName, dstName));
              return this;
          }
  
          public Builder removeEdge(String edgeName) {
              this.edgeToAttr.remove(edgeName);
              this.edgeToRawRelation.remove(edgeName);
              return this;
          }
  
          public Builder addEdgeAttribute(String edgeName, String attrName) {
              this.edgeToAttr.get(edgeName).add(attrName);
              return this;
          }
  
          public Builder removeEdgeAttribute(String edgeName, String attrName) {
              this.edgeToAttr.get(edgeName).remove(attrName);
              return this;
          }
  
          public Schema build() {
              Map<String, AttrDef> nameToAttrDef = new HashMap<>();
              Map<Integer, AttrDef> idToAttrDef = new HashMap<>();
              Map<String, TypeDef> nameToTypeDef = new HashMap<>();
              Map<Integer, TypeDef> idToTypeDef = new HashMap<>();
              Map<String, EdgeRelation> edgeNameToRelation = new HashMap<>();
              Map<Integer, EdgeRelation> edgeIdToRelation = new HashMap<>();
  
              int attrTypeId = 0;
              for (Map.Entry<String, DataType> entry : this.attrToDataType.entrySet()) {
                  AttrDef attrDef = new AttrDef(attrTypeId, entry.getKey(), entry.getValue());
                  nameToAttrDef.put(entry.getKey(), attrDef);
                  idToAttrDef.put(attrTypeId, attrDef);
                  attrTypeId++;
              }
  
              int typeId = 0;
              for (Map.Entry<String, Set<String>> entry : this.vertexToAttr.entrySet()) {
                  String vertexName = entry.getKey();
                  List<AttrDef> attrList = new ArrayList<>();
                  for (String attrName : entry.getValue()) {
                      if (nameToAttrDef.containsKey(attrName)) {
                          attrList.add(nameToAttrDef.get(attrName));
                      }
                  }
                  TypeDef typeDef = new TypeDef(TypeKind.VERTEX, typeId, vertexName, attrList);
                  nameToTypeDef.put(vertexName, typeDef);
                  idToTypeDef.put(typeId, typeDef);
                  typeId++;
              }
              for (Map.Entry<String, Set<String>> entry : this.edgeToAttr.entrySet()) {
                  String edgeName = entry.getKey();
                  List<AttrDef> attrList = new ArrayList<>();
                  for (String attrName : entry.getValue()) {
                      if (nameToAttrDef.containsKey(attrName)) {
                          attrList.add(nameToAttrDef.get(attrName));
                      }
                  }
                  TypeDef typeDef = new TypeDef(TypeKind.EDGE, typeId, edgeName, attrList);
                  nameToTypeDef.put(edgeName, typeDef);
                  idToTypeDef.put(typeId, typeDef);
                  typeId++;
              }
  
              for (RawEdgeRelation rawRelation : this.edgeToRawRelation.values()) {
                  String edgeName = rawRelation.getEdgeName();
                  if (!nameToTypeDef.containsKey(edgeName)) {
                      continue;
                  }
                  int edgeTid = nameToTypeDef.get(edgeName).getTypeId();
  
                  String srcName = rawRelation.getSrcName();
                  if (!nameToTypeDef.containsKey(srcName)) {
                      continue;
                  }
                  int srcTid = nameToTypeDef.get(srcName).getTypeId();
  
                  String dstName = rawRelation.getDstName();
                  if (!nameToTypeDef.containsKey(dstName)) {
                      continue;
                  }
                  int dstTid = nameToTypeDef.get(dstName).getTypeId();
  
                  EdgeRelation edgeRelation = new EdgeRelation(edgeTid, edgeName, srcTid, srcName, dstTid, dstName);
                  edgeNameToRelation.put(edgeName, edgeRelation);
                  edgeIdToRelation.put(edgeTid, edgeRelation);
              }
  
              return new Schema(nameToAttrDef, idToAttrDef, nameToTypeDef, idToTypeDef, edgeNameToRelation, edgeIdToRelation);
          }
      }
  }
  