import ts, { StringLiteral } from "typescript";
import { resolve, relative, dirname, join } from "path";
import { Glob } from "bun";

const RPC_PKG = resolve(import.meta.dirname!, "../packages/rpc");
const SRC_DIR = join(RPC_PKG, "src");
const INDEX_FILE = join(RPC_PKG, "index.ts");

const glob = new Glob("**/*_pb.ts");
const files = Array.from(glob.scanSync({ cwd: SRC_DIR, absolute: true })).sort();

if (files.length === 0) {
  console.error("No _pb.ts files found in", SRC_DIR);
  process.exit(1);
}

// Read the rpc package's tsconfig so module resolution finds @bufbuild/protobuf
const tsconfigPath = join(RPC_PKG, "tsconfig.json");
const { config } = ts.readConfigFile(tsconfigPath, ts.sys.readFile);
const parsed = ts.parseJsonConfigFileContent(config, ts.sys, RPC_PKG);

const program = ts.createProgram(files, {
  ...parsed.options,
  noEmit: true,
});

const checker = program.getTypeChecker();
const factory = ts.factory;

function hasTypeNamed(type: ts.Type, name: string): boolean {
  if (type.aliasSymbol?.name === name) return true;
  if (type.symbol?.name === name) return true;
  const target = (type as any).target as ts.Type | undefined;
  if (target?.symbol?.name === name) return true;
  if (type.isIntersection()) {
    return type.types.some((t) => hasTypeNamed(t, name));
  }
  if (type.isUnion()) {
    return type.types.some((t) => hasTypeNamed(t, name));
  }
  return false;
}

const isStringLiteral = (s: string | StringLiteral | undefined): s is StringLiteral => {
  if (s === undefined) {
    return false;
  }
  return (s as StringLiteral).forEachChild != undefined;
}

const serviceNameToRouterPropertyName = (n:string) : string => {
  return n.toLowerCase().replace('service','') + 's'
}

const exportDecl = (exp: Map<string, boolean>, typesOnly:boolean = false, from?: string | StringLiteral) => {
  return factory.createExportDeclaration(
    undefined,
    typesOnly,
    factory.createNamedExports(
      Array.from(exp.entries()).map((exp) => {
        return factory.createExportSpecifier(
          exp[1],
          undefined,
          exp[0]
        )
      })
    ),
    (
      isStringLiteral(from) ?
        from :
        (
          from != undefined ?
            factory.createStringLiteral(from) :
            from
        )
    )
  )
}

const importDecl = (imp: Map<string, boolean>, from: string | StringLiteral, typesOnly: boolean = false) => {
  return factory.createImportDeclaration(
    undefined,
    factory.createImportClause(
      typesOnly,
      undefined,
      factory.createNamedImports(
        Array.from(imp.entries()).map((k) => {
          return factory.createImportSpecifier(
            k[1],
            undefined,
            factory.createIdentifier(k[0])
          )
        })
      )
    ),
    (isStringLiteral(from) ? from : factory.createStringLiteral(from))
  )
}

const statements: ts.Statement[] = [];
const serviceNames: string[] = [];
const typeNames: string[] = [];

for (const filePath of files) {
  const sourceFile = program.getSourceFile(filePath);
  if (!sourceFile) continue;

  const symbol = checker.getSymbolAtLocation(sourceFile);
  if (!symbol) continue;

  const exports = checker.getExportsOfModule(symbol);

  const fileServiceNames: string[] = [];
  const fileTypeNames: string[] = [];

  for (const exp of exports) {
    const declarations = exp.getDeclarations();
    if (!declarations || declarations.length === 0) continue;

    const decl = declarations[0];

    if (ts.isTypeAliasDeclaration(decl) || ts.isInterfaceDeclaration(decl)) {
      const type = checker.getDeclaredTypeOfSymbol(exp);
      if (hasTypeNamed(type, "Message")) {
        fileTypeNames.push(exp.getName());
      }
    } else {
      const type = checker.getTypeOfSymbolAtLocation(exp, decl);
      if (hasTypeNamed(type, "GenService")) {
        fileServiceNames.push(exp.getName());
      }
    }
  }

  let modulePath = "./" + relative(dirname(INDEX_FILE), filePath);
  modulePath = modulePath.replace(/\.ts$/, "");

  const moduleSpecifier = factory.createStringLiteral(modulePath);
  statements.push(
    importDecl(
      new Map<string, boolean>([
        ["Client", true],
        ["createClient", false]
      ]),
      "@connectrpc/connect"
    )
  )
  statements.push(
    importDecl(
      new Map<string, boolean>([
        ["createConnectTransport", false]
      ]),
      "@connectrpc/connect-web"
    )
  )
  if (fileServiceNames.length > 0) {
    serviceNames.push(...fileServiceNames)
    statements.push(importDecl(
      new Map<string, boolean>(
        fileServiceNames.map((v) => {
          return [v, false] as const;
        })
      ),
      moduleSpecifier
    ))
  }

  if (fileTypeNames.length > 0) {
    typeNames.push(...fileTypeNames)
    statements.push(importDecl(
      new Map<string, boolean>(
        fileTypeNames.map((v) => {
          return [v, false] as const;
        })
      ),
      moduleSpecifier,
      true
    ))
  }
}

statements.push(
  exportDecl(
    new Map<string, boolean>(serviceNames.map((service) => {
      return [service, false] as const;
    }))
  )
)

statements.push(
  exportDecl(
    new Map<string, boolean>(typeNames.map((type) => {
      return [type, false] as const;
    })),
    true
  )
)

statements.push(
  factory.createInterfaceDeclaration(
    undefined,
    factory.createIdentifier("Router"),
    undefined,
    undefined,
    serviceNames.map((service) => {
      return factory.createPropertySignature(
        undefined,
        serviceNameToRouterPropertyName(service),
        undefined,
        factory.createTypeReferenceNode(
          factory.createIdentifier("Client"),
          [
            factory.createTypeQueryNode(
              factory.createIdentifier(service)
            )
          ]
        )
      )
    })
  )
)

statements.push(exportDecl(
  new Map<string, boolean>([
    ['Router', true]
  ])
))

statements.push(factory.createFunctionDeclaration(
  undefined,
  undefined,
  'createRouter',
  undefined,
  [
    factory.createParameterDeclaration(
      undefined,
      undefined,
      'baseUrl',
      undefined,
      factory.createKeywordTypeNode(ts.SyntaxKind.StringKeyword),
      undefined
    )
  ],
  factory.createTypeReferenceNode(factory.createIdentifier("Router"), undefined),
  factory.createBlock(
    [
      factory.createVariableStatement(
        undefined,
        factory.createVariableDeclarationList(
          [
            factory.createVariableDeclaration(
              'transport',
              undefined,
              undefined,
              factory.createCallExpression(
                factory.createIdentifier("createConnectTransport"),
                undefined,
                [
                  factory.createObjectLiteralExpression([
                    factory.createPropertyAssignment(
                      'baseUrl',
                      factory.createIdentifier('baseUrl')
                    )
                  ])
                ]
              )
            ),
          ],
          ts.NodeFlags.Const
        ),
      ),
      factory.createVariableStatement(
        undefined,
        factory.createVariableDeclarationList(
          [
            factory.createVariableDeclaration(
              'router',
              undefined,
              undefined,
              factory.createObjectLiteralExpression(
                [
                  ...serviceNames.map((service) => {
                    return factory.createPropertyAssignment(
                      serviceNameToRouterPropertyName(service),

                      factory.createCallExpression(
                        factory.createIdentifier('createClient'),
                        undefined,
                        [
                          factory.createIdentifier(service),
                          factory.createIdentifier('transport')
                        ]
                      )
                    )
                  })
                ]
              )
            )
          ],
          ts.NodeFlags.Const
        )
      ),
      factory.createReturnStatement(
        factory.createIdentifier('router')
      )
    ],
    true
  )
))

statements.push(exportDecl(new Map<string,boolean>([
  ['createRouter', false]
])))

const outputFile = ts.createSourceFile(INDEX_FILE, "", ts.ScriptTarget.ESNext, false, ts.ScriptKind.TS);
const printer = ts.createPrinter({ newLine: ts.NewLineKind.LineFeed });

const header = "// @generated by scripts/gen-rpc-index.ts — do not edit manually\n\n";
const body = statements
  .map((s) => printer.printNode(ts.EmitHint.Unspecified, s, outputFile))
  .join("\n");

await Bun.write(INDEX_FILE, header + body + "\n");
console.log(`Wrote ${INDEX_FILE}`);