bufpiv/src/app/proto-definition-selector/proto-definition.service.ts

import { Injectable } from '@angular/core';
import { FieldBase, MapField, ReflectionObject, parse } from 'protobufjs';
import {
  BooleanMessage,
  EnumMessage,
  ListMessage,
  MapMessage,
  MessageConfiguration,
  MessageTypeEnum,
  NumericMessage,
  ObjectMessage,
  ProtoBase,
  ProtoEnum,
  ProtoMessage,
  StringMessage,
  UnknownProto,
} from '../model/proto-message.model';
import { readTextFile } from '@tauri-apps/api/fs';

@Injectable({
  providedIn: 'root',
})
export class ProtoDefinitionService {
  async parseAllDefinitions(allProtoFiles: File[]): Promise<ProtoMessage[]> {
    const messages: (ProtoBase | undefined)[] = [];
    for (const protoFile of allProtoFiles) {
      const definition = await parse(await protoFile.text());
      // Try resolve imports, if they fail (due to invalid path or running app in browser)
      // they'll be skipped
      if (definition.imports) {
        for (const importPath of definition.imports) {
          try {
            const fileName = importPath.substring(
              importPath.lastIndexOf('/') + 1
            );
            allProtoFiles.push(
              new File(
                [
                  new Blob([await readTextFile(importPath)], {
                    type: 'text/plain',
                  }),
                ],
                fileName
              )
            );
          } catch (err) {
            console.error(
              `Failed to import file ${importPath}. Skipping...`,
              err
            );
          }
        }
      }
      const messageDefinitions = definition.root;
      const parsedMessages = this.parseAllNestedMessages(
        messageDefinitions,
        protoFile.name,
        definition.package
      );
      messages.push(...parsedMessages);
    }

    if (messages.length === 0) {
      return [];
    }

    // Now check for duplicates in case we auto-imported files
    for (let i = 0; i < messages.length - 1; i++) {
      for (let j = i + 1; j < messages.length; j++) {
        if (messages[i]?.fullName === messages[j]?.fullName) {
          messages[j] = undefined;
        }
      }
    }
    const distinctMessages = messages.filter(
      (message) => message
    ) as ProtoBase[];

    // Then populate all object/enum messages for the final passes
    const standardMessages = distinctMessages
      .filter((messageObject) => 'values' in messageObject)
      .map((messageObject) => messageObject as ProtoMessage);
    for (const messageObject of standardMessages) {
      for (const value of messageObject.values) {
        let configuration: MessageConfiguration | null = null;
        if (value.configuration.type === MessageTypeEnum.Object) {
          configuration = value.configuration;
        } else if (value.configuration.type === MessageTypeEnum.List) {
          const listConfiguration = value.configuration as ListMessage;
          if (
            listConfiguration.subConfiguration.type === MessageTypeEnum.Object
          ) {
            configuration = listConfiguration.subConfiguration;
          }
        } else if (value.configuration.type === MessageTypeEnum.Map) {
          const mapConfiguration = value.configuration as MapMessage;
          if (
            mapConfiguration.valueConfiguration.type === MessageTypeEnum.Object
          ) {
            configuration = mapConfiguration.valueConfiguration;
          }
        }
        if (configuration) {
          this.populateNestedObject(
            configuration as ObjectMessage,
            messageObject.packageName!,
            standardMessages
          );
        }
      }
    }
    const enumMessages = distinctMessages
      .filter((messageObject) => 'options' in messageObject)
      .map((messageObject) => messageObject as ProtoEnum);
    for (const messageObject of standardMessages) {
      for (const value of messageObject.values) {
        if (value.configuration.type === MessageTypeEnum.Object) {
          value.configuration = this.populateEnumMessages(
            value.configuration as ObjectMessage,
            enumMessages
          );
        } else if (value.configuration.type === MessageTypeEnum.List) {
          const listConfiguration = value.configuration as ListMessage;
          if (
            listConfiguration.subConfiguration.type === MessageTypeEnum.Object
          ) {
            listConfiguration.subConfiguration = this.populateEnumMessages(
              listConfiguration.subConfiguration as ObjectMessage,
              enumMessages
            );
          }
        } else if (value.configuration.type === MessageTypeEnum.Map) {
          const mapConfiguration = value.configuration as MapMessage;
          if (
            mapConfiguration.valueConfiguration.type === MessageTypeEnum.Object
          ) {
            mapConfiguration.valueConfiguration = this.populateEnumMessages(
              mapConfiguration.valueConfiguration as ObjectMessage,
              enumMessages
            );
          }
        }
      }
    }

    return standardMessages;
  }

  private populateEnumMessages(
    objectMessage: ObjectMessage,
    enumMessages: ProtoEnum[]
  ): MessageConfiguration {
    let useFullName = false;
    let objectMessageName = objectMessage.messageDefinition.name;
    if (objectMessageName.includes('.')) {
      // First do a check for current package.
      for (const enumMessage of enumMessages) {
        if (
          enumMessage.packageName ===
            objectMessage.messageDefinition.packageName &&
          enumMessage.fullName ===
            `${objectMessage.messageDefinition.packageName}.${objectMessageName}`
        ) {
          return EnumMessage(enumMessage.options);
        }
      }

      // If there's no match, then we'll be matching on full name, assuming a separate package
      useFullName = true;
    }

    // None matching messages in this package found, check external package
    for (const enumMessage of enumMessages) {
      if (
        (useFullName && enumMessage.fullName === objectMessageName) ||
        (!useFullName && enumMessage.name === objectMessageName)
      ) {
        return EnumMessage(enumMessage.options);
      }
    }
    objectMessage.messageDefinition.values =
      objectMessage.messageDefinition.values.map((obj) => {
        if (obj.configuration.type === MessageTypeEnum.Object) {
          return {
            ...obj,
            configuration: this.populateEnumMessages(
              obj.configuration as ObjectMessage,
              enumMessages
            ),
          };
        } else if (obj.configuration.type === MessageTypeEnum.List) {
          const listConfiguration = obj.configuration as ListMessage;
          if (
            listConfiguration.subConfiguration.type === MessageTypeEnum.Object
          ) {
            return {
              ...obj,
              subConfiguration: this.populateEnumMessages(
                listConfiguration.subConfiguration as ObjectMessage,
                enumMessages
              ),
            };
          }
        } else if (obj.configuration.type === MessageTypeEnum.Map) {
          const mapConfiguration = obj.configuration as MapMessage;
          if (
            mapConfiguration.valueConfiguration.type === MessageTypeEnum.Object
          ) {
            return {
              ...obj,
              valueConfiguration: this.populateEnumMessages(
                mapConfiguration.valueConfiguration as ObjectMessage,
                enumMessages
              ),
            };
          }
        }
        return obj;
      });
    return objectMessage;
  }

  private parseAllNestedMessages(
    obj: ReflectionObject,
    fileName: string,
    packageName?: string
  ) {
    const messageObjects: ProtoBase[] = [];
    // Nested messages/enums
    if ('nestedArray' in obj && obj.nestedArray) {
      const nestedArray = obj.nestedArray as ReflectionObject[];
      for (const nestedObj of nestedArray) {
        messageObjects.push(
          ...this.parseAllNestedMessages(nestedObj, fileName, packageName)
        );
      }
    }
    // Message
    if ('fieldsArray' in obj) {
      messageObjects.push(
        this.convertProtoDefinitionToModel(obj, fileName, packageName)
      );
    }
    // Enum
    if ('values' in obj) {
      messageObjects.push({
        name: obj.name,
        fullName: obj.fullName,
        packageName,
        fileName,
        options: Object.entries(obj.values as Record<string, string>)
          .sort((a, b) => Number(a[1]) - Number(b[1]))
          .map(([key, _]) => key),
      } as ProtoEnum);
    }
    return messageObjects;
  }

  private populateNestedObject(
    objectMessage: ObjectMessage,
    packageName: string,
    availableMessages: ProtoMessage[]
  ) {
    let useFullName = false;
    let objectMessageName = objectMessage.messageDefinition.name;
    if (objectMessageName.includes('.')) {
      // First do a check for current package.
      for (const message of availableMessages) {
        if (
          message.packageName === packageName &&
          message.fullName === `${packageName}.${objectMessageName}`
        ) {
          this.updateObjectMessage(objectMessage, message, availableMessages);
          return;
        }
      }

      // If there's no match, then we'll be matching on full name, assuming a separate package
      useFullName = true;
    }

    for (const message of availableMessages) {
      if (
        (useFullName && message.fullName === objectMessageName) ||
        (!useFullName && message.name === objectMessageName)
      ) {
        this.updateObjectMessage(objectMessage, message, availableMessages);
        return;
      }
    }
    objectMessage.messageDefinition = UnknownProto(
      objectMessage.messageDefinition.name,
      objectMessage.messageDefinition.fullName
    );
  }

  private updateObjectMessage(
    objectMessage: ObjectMessage,
    message: ProtoMessage,
    availableMessages: ProtoMessage[]
  ) {
    objectMessage.messageDefinition = {
      name: objectMessage.messageDefinition.name,
      values: message.values.map((value) => {
        let configuration: MessageConfiguration | null = null;
        if (value.configuration.type === MessageTypeEnum.Object) {
          configuration = value.configuration;
        } else if (value.configuration.type === MessageTypeEnum.List) {
          const listConfiguration = value.configuration as ListMessage;
          if (
            listConfiguration.subConfiguration.type === MessageTypeEnum.Object
          ) {
            configuration = listConfiguration.subConfiguration;
          }
        } else if (value.configuration.type === MessageTypeEnum.Map) {
          const mapConfiguration = value.configuration as MapMessage;
          if (
            mapConfiguration.valueConfiguration.type === MessageTypeEnum.Object
          ) {
            configuration = mapConfiguration.valueConfiguration;
          }
        }
        if (configuration) {
          this.populateNestedObject(
            configuration as ObjectMessage,
            message.packageName!,
            availableMessages
          );
        }
        return structuredClone(value);
      }),
    };
  }

  private convertProtoDefinitionToModel(
    messageDefinition: ReflectionObject,
    fileName: string,
    packageName?: string
  ): ProtoBase {
    if (messageDefinition && 'fieldsArray' in messageDefinition) {
      const convertedFields: ProtoMessage = {
        name: messageDefinition.name,
        fullName: messageDefinition.fullName.substring(1), // trim leading .
        packageName,
        fileName,
        values: [],
      };
      const fields = messageDefinition.fieldsArray as FieldBase[];
      for (const field of fields) {
        let type: MessageConfiguration | undefined;
        if ('rule' in field && field.rule === 'repeated') {
          const subType = this.getTypeFromField(field.type);
          if (subType) {
            type = ListMessage(subType);
          }
        } else if (field instanceof MapField) {
          // Map
          const keyConfiguration = this.getTypeFromField(field.keyType);
          const valueConfiguration = this.getTypeFromField(field.type);
          if (keyConfiguration && valueConfiguration) {
            type = MapMessage(keyConfiguration, valueConfiguration);
          }
        } else {
          type = this.getTypeFromField(field.type);
        }
        if (type) {
          convertedFields.values.push({
            name: field.name,
            configuration: type,
          });
        } else {
          console.error(`Failed to find type ${field.type}`);
        }
      }
      return convertedFields;
    }
    throw 'ReflectionObject is not a message';
  }

  private getTypeFromField(
    fieldType: string
  ): MessageConfiguration | undefined {
    switch (fieldType) {
      case 'string':
        return StringMessage();
        break;
      case 'int32':
      case 'int64':
      case 'float':
      case 'double':
        return NumericMessage();
        break;
      case 'bool':
        return BooleanMessage();
        break;
      // TODO: bytes as well, though that's pretty useless (can't really represent/edit it?)
    }
    return ObjectMessage({ name: fieldType, values: [] });
  }
}