type FlattenContext struct {
    Message      *protogen.Message
    FlattenInfos []*FlattenFieldInfo
}
 
type FlattenFieldInfo struct {
    Field  *protogen.Field
    Prefix string            // flatten_prefix value (may be empty)
}

func (x *MessageName) MarshalJSON() ([]byte, error) {
    if x == nil {
        return []byte("null"), nil
    }
    data, err := protojson.Marshal(x)
    if err != nil {
        return nil, err
    }
    var raw map[string]json.RawMessage
    if err := json.Unmarshal(data, &raw); err != nil {
        return nil, err
    }
 
    // For each flattened field:
    // 1. Remove the nested field from raw
    // 2. Marshal the child message (use json.Marshal to respect child's MarshalJSON)
    // 3. Add each child field to raw with optional prefix
 
    if x.FlattenedFieldGoName != nil {
        delete(raw, "flattenedFieldJsonName")
        childData, err := json.Marshal(x.FlattenedFieldGoName)
        if err != nil {
            return nil, err
        }
        var childRaw map[string]json.RawMessage
        if err := json.Unmarshal(childData, &childRaw); err != nil {
            return nil, err
        }
        for k, v := range childRaw {
            raw["prefix"+k] = v  // prefix may be empty
        }
    }
 
    return json.Marshal(raw)
}

func (x *MessageName) UnmarshalJSON(data []byte) error {
    var raw map[string]json.RawMessage
    if err := json.Unmarshal(data, &raw); err != nil {
        return err
    }
 
    // For each flattened field:
    // 1. Extract child fields from raw into a new map (with prefix stripped)
    // 2. Marshal the child map and unmarshal into the child message
    // 3. Delete the extracted fields from raw
 
    childRaw := make(map[string]json.RawMessage)
    // For each known child field of the flattened message:
    // Check if "prefix" + childJsonName exists in raw
    if v, ok := raw["prefixChildField1"]; ok {
        childRaw["childField1"] = v
        delete(raw, "prefixChildField1")
    }
    if v, ok := raw["prefixChildField2"]; ok {
        childRaw["childField2"] = v
        delete(raw, "prefixChildField2")
    }
    // ... for each child field
 
    if len(childRaw) > 0 {
        childData, err := json.Marshal(childRaw)
        if err != nil {
            return err
        }
        x.FlattenedFieldGoName = &ChildMessageType{}
        if err := json.Unmarshal(childData, x.FlattenedFieldGoName); err != nil {
            return err
        }
    }
 
    // Re-marshal remaining and unmarshal via protojson for non-flattened fields
    remaining, err := json.Marshal(raw)
    if err != nil {
        return err
    }
    return protojson.Unmarshal(remaining, x)
}

// Without flatten:
export interface Order {
  id: string;
  billing: Address;
  shipping: Address;
}
 
// With flatten + prefix:
export interface Order {
  id: string;
  billing_street: string;
  billing_city: string;
  billing_zip: string;
  shipping_street: string;
  shipping_city: string;
  shipping_zip: string;
}

Order:
  allOf:
    - type: object
      properties:
        id:
          type: string
    - type: object
      description: "billing (flattened with prefix billing_)"
      properties:
        billing_street:
          type: string
        billing_city:
          type: string
        billing_zip:
          type: string
    - type: object
      description: "shipping (flattened with prefix shipping_)"
      properties:
        shipping_street:
          type: string
        shipping_city:
          type: string
        shipping_zip:
          type: string

Order:
  type: object
  properties:
    id:
      type: string
    billing_street:
      type: string
    billing_city:
      type: string
    billing_zip:
      type: string
    shipping_street:
      type: string
    shipping_city:
      type: string
    shipping_zip:
      type: string

syntax = "proto3";
package flatten.test;
 
import "sebuf/http/annotations.proto";
 
option go_package = "github.com/SebastienMelki/sebuf/internal/httpgen/testdata/flatten";
 
// Child message used for flattening
message Address {
  string street = 1;
  string city = 2;
  string zip = 3;
}
 
message ContactInfo {
  string email = 1;
  string phone = 2;
}
 
// Basic flatten without prefix
message SimpleFlatten {
  string id = 1;
  Address address = 2 [(sebuf.http.flatten) = true];
}
 
// Flatten with prefix (two flattened fields of same type)
message DualFlatten {
  string id = 1;
  Address billing = 2 [
    (sebuf.http.flatten) = true,
    (sebuf.http.flatten_prefix) = "billing_"
  ];
  Address shipping = 3 [
    (sebuf.http.flatten) = true,
    (sebuf.http.flatten_prefix) = "shipping_"
  ];
}
 
// Flatten with mixed fields (some flattened, some not)
message MixedFlatten {
  string id = 1;
  Address address = 2 [(sebuf.http.flatten) = true];
  ContactInfo contact = 3;  // Not flattened
  string notes = 4;
}
 
// No flatten annotation (backward compatible)
message PlainNested {
  string id = 1;
  Address address = 2;
}
 
service FlattenService {
  option (sebuf.http.service_config) = {
    base_path: "/api/v1"
  };
 
  rpc TestSimpleFlatten(SimpleFlatten) returns (SimpleFlatten) {
    option (sebuf.http.config) = {
      path: "/flatten/simple"
      method: HTTP_METHOD_POST
    };
  }
 
  rpc TestDualFlatten(DualFlatten) returns (DualFlatten) {
    option (sebuf.http.config) = {
      path: "/flatten/dual"
      method: HTTP_METHOD_POST
    };
  }
 
  rpc TestMixedFlatten(MixedFlatten) returns (MixedFlatten) {
    option (sebuf.http.config) = {
      path: "/flatten/mixed"
      method: HTTP_METHOD_POST
    };
  }
 
  rpc TestPlainNested(PlainNested) returns (PlainNested) {
    option (sebuf.http.config) = {
      path: "/flatten/plain"
      method: HTTP_METHOD_POST
    };
  }
}

UPDATE_GOLDEN=1 go test -run TestExhaustiveGoldenFiles ./internal/httpgen/...
UPDATE_GOLDEN=1 go test -run TestGoldenFiles ./internal/clientgen/...
UPDATE_GOLDEN=1 go test -run TestGoldenFiles ./internal/tsclientgen/...
UPDATE_GOLDEN=1 go test -run TestExhaustiveGoldenFiles ./internal/openapiv3/...