slices / packages / codegen / src / client.ts

import type { SourceFile } from "ts-morph";
import { Scope } from "ts-morph";
import type { Lexicon } from "./mod.ts";
import { nsidToPascalCase, capitalizeFirst, sanitizeIdentifier } from "./mod.ts";

interface NestedStructure {
  [key: string]: NestedStructure | string | QueryProcedureInfo[] | undefined;
  _recordType?: string;
  _collectionPath?: string;
  _queryProcedures?: QueryProcedureInfo[];
}

interface QueryProcedureInfo {
  nsid: string;
  type: "query" | "procedure";
  methodName: string;
  parametersType?: string;
  inputType?: string;
  outputType?: string;
}

interface PropertyInfo {
  name: string;
  type: string;
}

interface MethodInfo {
  name: string;
  parameters: Array<{ name: string; type: string; hasQuestionToken?: boolean }>;
  returnType: string;
}

export function generateClient(
  sourceFile: SourceFile,
  lexicons: Lexicon[]
): void {
  // Create nested structure from lexicons
  const nestedStructure: NestedStructure = {};

  for (const lexicon of lexicons) {
    if (lexicon.definitions && typeof lexicon.definitions === "object") {
      // Check if this lexicon has any records, queries, or procedures
      const hasRecordsOrEndpoints = Object.values(lexicon.definitions).some(
        (defValue) =>
          (defValue.type === "record" && defValue.record) ||
          defValue.type === "query" ||
          defValue.type === "procedure"
      );

      // Only build nested structure for lexicons that have records
      if (hasRecordsOrEndpoints) {
        for (const [_, defValue] of Object.entries(lexicon.definitions)) {
          if (defValue.type === "record" && defValue.record) {
            const parts = lexicon.id.split(".");
            let current = nestedStructure;

            // Build nested structure
            for (const part of parts) {
              if (!current[part]) {
                current[part] = {};
              }
              current = current[part] as NestedStructure;
            }

            // Add the record interface name and store collection path
            current._recordType = nsidToPascalCase(lexicon.id);
            current._collectionPath = lexicon.id;
          }
        }
      }
    }
  }

  // Add query/procedure methods to the appropriate parent clients
  function addQueryProcedureMethods(
    obj: NestedStructure,
    lexicons: Lexicon[]
  ): void {
    // Group query/procedure endpoints by their parent collection path
    const endpointsByParent = new Map<string, QueryProcedureInfo[]>();

    for (const lexicon of lexicons) {
      if (lexicon.definitions && typeof lexicon.definitions === "object") {
        for (const [_, defValue] of Object.entries(lexicon.definitions)) {
          if (defValue.type === "query" || defValue.type === "procedure") {
            const parts = lexicon.id.split(".");
            const methodName = parts[parts.length - 1];

            // Find the parent collection by removing the method name
            // e.g., "network.slices.slice.getJobStatus" -> "network.slices.slice"
            const parentPath = parts.slice(0, -1).join(".");

            const queryProcedureInfo: QueryProcedureInfo = {
              nsid: lexicon.id,
              type: defValue.type as "query" | "procedure",
              methodName,
              parametersType:
                defValue.parameters?.properties &&
                Object.keys(defValue.parameters.properties).length > 0
                  ? `${nsidToPascalCase(lexicon.id)}Params`
                  : undefined,
              inputType: defValue.input
                ? `${nsidToPascalCase(lexicon.id)}Input`
                : undefined,
              outputType: defValue.output
                ? `${nsidToPascalCase(lexicon.id)}Output`
                : undefined,
            };

            if (!endpointsByParent.has(parentPath)) {
              endpointsByParent.set(parentPath, []);
            }
            endpointsByParent.get(parentPath)!.push(queryProcedureInfo);
          }
        }
      }
    }

    // Add endpoints to their respective parent clients
    function addEndpointsToNode(
      current: NestedStructure,
      path: string[],
      fullPath: string
    ): void {
      if (path.length === 0) {
        // We've reached the target node, add the endpoints if this has a record type
        const endpoints = endpointsByParent.get(fullPath);
        if (endpoints && current._recordType) {
          if (!current._queryProcedures) {
            current._queryProcedures = [];
          }
          current._queryProcedures.push(...endpoints);
        }
        return;
      }

      const [head, ...tail] = path;
      if (current[head]) {
        addEndpointsToNode(current[head] as NestedStructure, tail, fullPath);
      }
    }

    // Add endpoints to their parent nodes
    for (const parentPath of endpointsByParent.keys()) {
      const pathParts = parentPath.split(".");
      addEndpointsToNode(obj, pathParts, parentPath);
    }
  }

  // Add query/procedure methods before generating classes
  addQueryProcedureMethods(nestedStructure, lexicons);

  // Generate nested class structure
  function generateNestedClass(
    obj: NestedStructure,
    className = "CollectionNode",
    currentPath: string[] = []
  ): void {
    const properties: PropertyInfo[] = [];
    const methods: MethodInfo[] = [];

    let collectionPath = "";

    for (const [key, value] of Object.entries(obj)) {
      if (key === "_recordType") {
        // Add collection operations for this record type
        const recordName = value as string;
        // Check if we have sortable fields for this record
        const hasSortFields = sourceFile.getTypeAlias(
          `${recordName}SortFields`
        );
        const sortFieldsType = hasSortFields
          ? `${recordName}SortFields`
          : "IndexedRecordFields";
        const whereFieldsType = hasSortFields
          ? `${recordName}SortFields | IndexedRecordFields`
          : "IndexedRecordFields";

        methods.push({
          name: "getRecords",
          parameters: [
            {
              name: "params",
              type: `{ limit?: number; cursor?: string; where?: { [K in ${whereFieldsType}]?: WhereCondition }; orWhere?: { [K in ${whereFieldsType}]?: WhereCondition }; sortBy?: SortField<${sortFieldsType}>[]; }`,
              hasQuestionToken: true,
            },
          ],
          returnType: `Promise<GetRecordsResponse<${value}>>`,
        });
        methods.push({
          name: "getRecord",
          parameters: [{ name: "params", type: "GetRecordParams" }],
          returnType: `Promise<RecordResponse<${value}>>`,
        });
        methods.push({
          name: "countRecords",
          parameters: [
            {
              name: "params",
              type: `{ limit?: number; cursor?: string; where?: { [K in ${whereFieldsType}]?: WhereCondition }; orWhere?: { [K in ${whereFieldsType}]?: WhereCondition }; sortBy?: SortField<${sortFieldsType}>[]; }`,
              hasQuestionToken: true,
            },
          ],
          returnType: "Promise<CountRecordsResponse>",
        });
        // Add create, update, delete methods
        methods.push({
          name: "createRecord",
          parameters: [
            { name: "record", type: value as string },
            { name: "useSelfRkey", type: "boolean", hasQuestionToken: true },
          ],
          returnType: `Promise<{ uri: string; cid: string }>`,
        });
        methods.push({
          name: "updateRecord",
          parameters: [
            { name: "rkey", type: "string" },
            { name: "record", type: value as string },
          ],
          returnType: `Promise<{ uri: string; cid: string }>`,
        });
        methods.push({
          name: "deleteRecord",
          parameters: [{ name: "rkey", type: "string" }],
          returnType: `Promise<void>`,
        });
      } else if (key === "_collectionPath") {
        collectionPath = value as string;
      } else if (key === "_queryProcedures") {
        // Add query and procedure methods
        const queryProcedures = value as QueryProcedureInfo[];
        for (const qp of queryProcedures) {
          if (qp.type === "query") {
            // Generate query method (GET)
            const parameters = [];
            if (qp.parametersType) {
              parameters.push({
                name: "params",
                type: qp.parametersType,
                hasQuestionToken: true,
              });
            }
            methods.push({
              name: qp.methodName,
              parameters,
              returnType: `Promise<${qp.outputType || "void"}>`,
            });
          } else if (qp.type === "procedure") {
            // Generate procedure method (POST)
            const parameters = [];
            if (qp.inputType) {
              parameters.push({
                name: "input",
                type: qp.inputType,
              });
            } else if (qp.parametersType) {
              parameters.push({
                name: "params",
                type: qp.parametersType,
              });
            }
            methods.push({
              name: qp.methodName,
              parameters,
              returnType: `Promise<${qp.outputType || "void"}>`,
            });
          }
        }
      } else if (typeof value === "object" && Object.keys(value).length > 0) {
        // Add nested property with PascalCase class name
        // Sanitize the key for both class name and property name
        const sanitizedKey = sanitizeIdentifier(key);
        const nestedClassName = `${capitalizeFirst(sanitizedKey)}${className}`;
        generateNestedClass(value as NestedStructure, nestedClassName, [
          ...currentPath,
          key,
        ]);
        properties.push({
          name: sanitizedKey,
          type: nestedClassName,
        });
      }
    }

    if (properties.length > 0 || methods.length > 0) {
      // Use proper naming for the main client
      const finalClassName =
        className === "Client" ? "AtProtoClient" : className;

      const classDeclaration = sourceFile.addClass({
        name: finalClassName,
        isExported: className === "Client",
        extends: className === "Client" ? "SlicesClient" : undefined,
        properties: [
          ...properties.map((p) => ({
            name: p.name,
            type: p.type,
            isReadonly: true,
          })),
          // Add OAuth client to the main AtProtoClient
          ...(className === "Client"
            ? [
                {
                  name: "oauth",
                  type: "OAuthClient | AuthProvider",
                  isReadonly: true,
                  hasQuestionToken: true,
                },
              ]
            : [
                // Nested classes need a reference to the client
                {
                  name: "client",
                  type: "SlicesClient",
                  scope: Scope.Private,
                  isReadonly: true,
                },
              ]),
        ],
      });

      // Add constructor
      const ctor = classDeclaration.addConstructor({
        parameters:
          className === "Client"
            ? [
                { name: "baseUrl", type: "string" },
                { name: "sliceUri", type: "string" },
                {
                  name: "oauthClient",
                  type: "OAuthClient | AuthProvider",
                  hasQuestionToken: true,
                },
              ]
            : [{ name: "client", type: "SlicesClient" }],
      });

      if (className === "Client") {
        ctor.addStatements([
          "super(baseUrl, sliceUri, oauthClient);",
          ...properties.map((p) => `this.${p.name} = new ${p.type}(this);`),
          "this.oauth = oauthClient;",
        ]);
      } else {
        // Nested classes store reference to parent client
        ctor.addStatements([
          "this.client = client;",
          ...properties.map((p) => `this.${p.name} = new ${p.type}(client);`),
        ]);
      }

      // Add methods with implementations
      for (const method of methods) {
        const methodDecl = classDeclaration.addMethod({
          name: method.name,
          parameters: method.parameters,
          returnType: method.returnType,
          isAsync: true,
        });

        // Add basic implementation using shared client methods
        // Use this.client for nested classes, this for main client
        const clientRef = className === "Client" ? "this" : "this.client";

        if (method.name === "getRecords") {
          methodDecl.addStatements([
            `return await ${clientRef}.getRecords('${collectionPath}', params);`,
          ]);
        } else if (method.name === "getRecord") {
          methodDecl.addStatements([
            `return await ${clientRef}.getRecord('${collectionPath}', params);`,
          ]);
        } else if (method.name === "createRecord") {
          methodDecl.addStatements([
            `return await ${clientRef}.createRecord('${collectionPath}', record, useSelfRkey);`,
          ]);
        } else if (method.name === "updateRecord") {
          methodDecl.addStatements([
            `return await ${clientRef}.updateRecord('${collectionPath}', rkey, record);`,
          ]);
        } else if (method.name === "deleteRecord") {
          methodDecl.addStatements([
            `return await ${clientRef}.deleteRecord('${collectionPath}', rkey);`,
          ]);
        } else if (method.name === "countRecords") {
          methodDecl.addStatements([
            `return await ${clientRef}.countRecords('${collectionPath}', params);`,
          ]);
        } else {
          // Handle query and procedure methods
          const queryProcedures = obj._queryProcedures || [];
          const matchingQP = queryProcedures.find(
            (qp) => qp.methodName === method.name
          );

          if (matchingQP) {
            if (matchingQP.type === "query") {
              // Query methods use GET with query parameters
              const paramArg = method.parameters.length > 0 ? "params" : "{}";
              methodDecl.addStatements([
                `return await ${clientRef}.makeRequest<${
                  matchingQP.outputType || "void"
                }>('${matchingQP.nsid}', 'GET', ${paramArg});`,
              ]);
            } else if (matchingQP.type === "procedure") {
              // Procedure methods use POST with body
              const paramArg =
                method.parameters.length > 0 ? method.parameters[0].name : "{}";
              methodDecl.addStatements([
                `return await ${clientRef}.makeRequest<${
                  matchingQP.outputType || "void"
                }>('${matchingQP.nsid}', 'POST', ${paramArg});`,
              ]);
            }
          }
        }
      }
    }
  }

  // Generate the main client class
  if (Object.keys(nestedStructure).length > 0) {
    generateNestedClass(nestedStructure, "Client");
  }
}