From 5a933a2a27a242593256978d429f2f482ac4d239 Mon Sep 17 00:00:00 2001
From: Yousef Haggy <yousef.r.haggy@gmail.com>
Date: Mon, 25 Nov 2024 11:25:01 -0500
Subject: [PATCH] Support explicit mappings

---
 .../components/model-wrapper.jsx              | 64 ++++++++++++++++---
 1 file changed, 54 insertions(+), 10 deletions(-)

diff --git a/src/core/plugins/json-schema-5/components/model-wrapper.jsx b/src/core/plugins/json-schema-5/components/model-wrapper.jsx
index 31979ae3cff..2613d7dd689 100644
--- a/src/core/plugins/json-schema-5/components/model-wrapper.jsx
+++ b/src/core/plugins/json-schema-5/components/model-wrapper.jsx
@@ -61,6 +61,51 @@ export default class ModelWrapper extends Component {
     return null
   }
 
+  /**
+   * Builds a Map of schema options combining explicit discriminator mappings and implicit mappings.
+   * 
+   * @returns {Map<string, string[]>} A Map where:
+   *   - key: the schema name (e.g., "Cat", "Dog")
+   *   - value: array of discriminator values that map to this schema
+   * 
+   * Examples:
+   * 1. Explicit mapping only:
+   *    { "Cat": ["kitty", "kitten"], "Dog": ["puppy"] }
+   * 
+   * 2. Implicit mapping only:
+   *    { "Cat": ["Cat"], "Dog": ["Dog"] }
+   * 
+   * 3. Mixed mapping:
+   *    { "Cat": ["kitty", "kitten"], "Dog": ["Dog"] }
+   *    where "Cat" has explicit mappings but "Dog" uses implicit
+   */
+  buildSchemaOptions = (name, discriminator, schemaMap) => {
+    const options = new Map()
+    const mapping = discriminator && discriminator.get("mapping")
+    
+    // First add any explicit mappings
+    if (mapping && mapping.size > 0) {
+      mapping.forEach((schemaRef, key) => {
+        const schemaName = this.getModelName(schemaRef)
+        if (schemaName) {
+          const existing = options.get(schemaName) || []
+          options.set(schemaName, [...existing, key])
+        }
+      })
+    }
+
+    // Then add implicit mappings for any schemas not already mapped
+    const childSchemas = schemaMap[name] || []
+    childSchemas.forEach(childName => {
+      if (!options.has(childName)) {
+        // No explicit mapping for this schema, use implicit
+        options.set(childName, [childName])
+      }
+    })
+
+    return options
+  }
+
   render() {
     let { getComponent, getConfigs, schema, specSelectors } = this.props
     const Model = getComponent("Model")
@@ -71,28 +116,27 @@ export default class ModelWrapper extends Component {
     }
 
     const name = this.getModelName(schema.get("$$ref"))
-
     const schemaMap = specSelectors.getSchemaMap().toJS()
-    const childSchemas = schemaMap[name] || []
-
-    const hasDiscriminator = schema.get("discriminator")
-    const showDropdown = hasDiscriminator && childSchemas.length > 0
-
+    const discriminator = schema.get("discriminator")
+    
+    const options = this.buildSchemaOptions(name, discriminator, schemaMap)
+    const showDropdown = !!discriminator && options.size > 0
+  
     // Use selected schema or original base schema
     const effectiveSchema = this.state.selectedSchema 
       ? specSelectors.findDefinition(this.state.selectedSchema)
       : schema
 
-
     return (
       <div className="model-box">
-        {/* TODO: consider names from discriminator "mapping" */}
         {showDropdown && (
           <div className="model-box-control">
             <select onChange={this.onSchemaSelect} value={this.state.selectedSchema || ""}>
               <option value="">Base: {name}</option>
-              {childSchemas.map(childName => (
-                <option key={childName} value={childName}>{childName}</option>
+              {Array.from(options.entries()).map(([schemaName, keys]) => (
+                <option key={schemaName} value={schemaName}>
+                  {keys.length > 1 ? `${keys.join(" | ")} (${schemaName})` : schemaName}
+                </option>
               ))}
             </select>
           </div>