#!/usr/bin/env ts-node

import { jsonToLex, LexValue, xrpc } from "@atproto/lex";
import fs from "fs";
import path from "path";
import { glob } from "glob";

interface LexiconPath {
  segments: string[];
  dotPath(): string;
  filePath(): string;
}

class LexiconPathImpl implements LexiconPath {
  constructor(public segments: string[]) {}

  dotPath(): string {
    return this.segments.join(".");
  }

  filePath(): string {
    return this.segments.join("/");
  }

  static fromString(pathStr: string): LexiconPath {
    return new LexiconPathImpl(pathStr.split("."));
  }
}

// Constants
const BASE_DIR = path.dirname(path.resolve(import.meta.dirname));
const LEXICON_DIR = path.join(BASE_DIR, "lexicons");

const TYPE_MAPPING: Record<string, string> = {
  string: "String",
  integer: "Int",
  number: "Float",
  boolean: "Boolean",
  bytes: "String",
};

const FORMAT_MAPPING: Record<string, string> = {
  did: "ID",
  handle: "ID",
  uri: "String",
  cid: "String",
  datetime: "String",
  date: "String",
  time: "String",
  "at-identifier": "ID",
};

const SKIP_FIELDS = new Set(["debug"]); // Fields to skip in GraphQL output

function resolveLexiconPath(lexiconPath: LexiconPath): string {
  const dirParts = lexiconPath.segments;
  const fileName = `${dirParts[dirParts.length - 1]}.json`;
  const filePath = path.join(LEXICON_DIR, ...dirParts.slice(0, -1), fileName);

  if (!fs.existsSync(filePath)) {
    throw new Error(`Lexicon file not found: ${filePath}`);
  }

  return filePath;
}

function normalizeReference(base: string, ref: string): string {
  if (base.includes("#") && ref[0] === "#") {
    return base.split("#")[0] + ref;
  } else if (!base.includes("#") && ref[0] === "#") {
    return base + ref;
  } else {
    return ref;
  }
}

function resolveReference(ref: string): LexValue | null {
  try {
    if (ref.includes("#")) {
      const [lexiconId, typeName] = ref.split("#", 2);
      return resolveReferenceType(lexiconId, typeName || "main");
    } else {
      return resolveReferenceType(ref, "main");
    }
  } catch (e) {
    console.error(`Warning: Could not resolve reference ${ref}: ${e}`, e);
    return null;
  }
}

function resolveReferenceType(
  lexiconId: string,
  typeName: string,
): LexValue | null {
  try {
    const lexiconPath = resolveLexiconPath(
      LexiconPathImpl.fromString(lexiconId),
    );
    const lexiconData = JSON.parse(fs.readFileSync(lexiconPath, "utf-8"));
    const parsed = jsonToLex(lexiconData);

    if (!parsed || !parsed.defs) return null;

    const defs = parsed.defs;
    if (defs[typeName]) {
      return defs[typeName];
    }

    return null;
  } catch (e) {
    console.error(`Error resolving reference ${lexiconId}#${typeName}: ${e}`);
    return null;
  }
}

function lexTypeToGraphQLType(
  base: string,
  lexType: any,
  isRequired: boolean = false,
): string {
  let output;
  if (lexType.type === "array") {
    const itemsType = lexType.items;
    const itemType = lexTypeToGraphQLType(base, itemsType, true);
    output = `[${itemType}]`;
  } else if (lexType.type === "ref") {
    const ref = normalizeReference(base, lexType.ref) || "";
    const typeName = `Lexicon_${ref.replace(/#/g, ".").replace(/\./g, "_")}`;
    output = typeName;
  } else if (lexType.type === "string") {
    output =
      lexType.format && FORMAT_MAPPING[lexType.format]
        ? FORMAT_MAPPING[lexType.format]
        : "String";
  } else if (TYPE_MAPPING[lexType.type]) {
    output = TYPE_MAPPING[lexType.type];
  } else {
    output = "Unknown";
  }

  const nullable_indicator = isRequired ? "!" : "";
  return `${output}${nullable_indicator}`;
}

function toUpperSnakeCase(s: string): string {
  return s.replace(/[^a-zA-Z0-9]+/g, "_").toUpperCase();
}

function lexObjectDefinitionToGraphQL(typeName: string, typeDef: any): string {
  const graphqlLines: string[] = [];

  if (typeDef.type === "string" && typeDef.knownValues) {
    // Enums
    graphqlLines.push(`# ${typeName}`);
    graphqlLines.push(
      `enum Lexicon_${typeName.replace(/#/g, "_").replace(/\./g, "_")} {`,
    );

    for (const fieldName of typeDef.knownValues) {
      graphqlLines.push(`  ${toUpperSnakeCase(fieldName)}`);
    }
  } else if (typeDef.type === "object") {
    // Struct/object definitions
    graphqlLines.push(`# ${typeName}`);
    graphqlLines.push(
      `type Lexicon_${typeName.replace(/#/g, "_").replace(/\./g, "_")} {`,
    );

    const properties = typeDef.properties || {};
    const requiredFields = new Set(typeDef.required || []);

    for (const [fieldName, fieldDef] of Object.entries(properties)) {
      // Skip debug field as per requirements
      if (SKIP_FIELDS.has(fieldName)) {
        continue;
      }

      // Skip ref unions
      if (fieldDef.type === "refUnion") {
        continue;
      }

      // Skip "unknown" fields
      if (fieldDef.type === "unknown") {
        continue;
      }

      const isRequired = requiredFields.has(fieldName);
      const graphqlType = lexTypeToGraphQLType(typeName, fieldDef, isRequired);
      // TODO better resolve "Unknown" types
      if (!graphqlType.startsWith("Unknown")) {
        graphqlLines.push(`  ${fieldName}: ${graphqlType}`);
      }
    }
  }

  graphqlLines.push("}");
  return graphqlLines.join("\n");
}

export function outputLexiconNamespaces(
  namespaces: Record<string, string>,
  root: LexiconPath,
): string[] {
  const lines: string[] = [];
  const chunks: string[] = [];
  const rootLen = root.segments.length;

  while (Object.keys(namespaces).length > 0) {
    const lexiconPath = LexiconPathImpl.fromString(Object.keys(namespaces)[0]);
    if (rootLen === lexiconPath.segments.length - 1) {
      // Done, output the type
      const path = lexiconPath.dotPath();
      lines.push(namespaces[path]);
      delete namespaces[path];
    } else {
      const nextSegment = lexiconPath.segments[rootLen];
      const nextPath =
        rootLen > 0 ? `${root.dotPath()}.${nextSegment}` : nextSegment;

      const group: Record<string, string> = {};
      for (const [path, def] of Object.entries(namespaces)) {
        if (path.startsWith(`${nextPath}.`)) {
          group[path] = def;
          delete namespaces[path];
        }
      }

      const nestedChunks = outputLexiconNamespaces(
        group,
        LexiconPathImpl.fromString(nextPath),
      );
      chunks.push(...nestedChunks);

      const typeName = `Lexicon_${LexiconPathImpl.fromString(nextPath).segments.join("_")}`;
      lines.push(`${nextSegment}: ${typeName}!`);
    }
  }

  const content = lines.map((line) => `  ${line}`).join("\n");
  const typeName = ["Lexicon"].concat(root.segments).join("_");
  chunks.push(`type ${typeName} {\n${content}\n}\n`);
  return chunks;
}

export function generateLexiconStructure(lexiconPath: LexiconPath): {
  methodDef: string;
  objectDefinitions: Record<string, any>;
  methodCall: Function;
} | null {
  const methodName = lexiconPath.segments[lexiconPath.segments.length - 1];

  try {
    const filePath = resolveLexiconPath(lexiconPath);
    const lexiconData = JSON.parse(fs.readFileSync(filePath, "utf-8"));
    const parsed = jsonToLex(lexiconData);

    if (!parsed || typeof parsed !== "object" || !("defs" in parsed)) {
      return null;
    }

    const defs = parsed.defs;
    const mainDef = defs.main;

    if (!mainDef || mainDef.type !== "query") {
      return null;
    }

    const referencedTypes = new Set<string>();
    const objectDefinitions: Record<string, any> = {};

    // Add parameters as fields
    const methodFields: string[] = [];
    if (mainDef.parameters) {
      const params = mainDef.parameters;
      const requiredFields = new Set(params.required || []);

      for (const [fieldName, fieldDef] of Object.entries(
        params.properties || {},
      )) {
        const isRequired = requiredFields.has(fieldName);
        const graphqlType = lexTypeToGraphQLType(
          lexiconPath.dotPath(),
          fieldDef,
          isRequired,
        );
        methodFields.push(`${fieldName}: ${graphqlType}`);

        // Sweep up refs
        if (fieldDef.type === "ref") {
          referencedTypes.add(fieldDef.ref);
        }
      }
    }

    const methodParams =
      methodFields.length > 0 ? `(${methodFields.join(", ")})` : "";

    let returnType = null;

    if (mainDef.output) {
      const outputSchema = mainDef.output.schema;

      if (outputSchema.type === "ref") {
        // Use the output type as return type
        returnType = lexTypeToGraphQLType(lexiconPath.dotPath(), outputSchema);

        // Sweep up refs
        referencedTypes.add(outputSchema.ref);
      } else if (outputSchema.type === "object") {
        // Handle inline object definitions for return type
        const outputType = `Lexicon_${lexiconPath.segments.join("_")}Output`;
        returnType = outputType;
        objectDefinitions[`${lexiconPath.dotPath()}Output`] = outputSchema;
      }
    }

    // Resolve references
    for (const ref of referencedTypes) {
      const normalizedRef = normalizeReference(lexiconPath.dotPath(), ref);
      const refDef = resolveReference(normalizedRef);
      if (refDef) {
        objectDefinitions[ref] = refDef;
      }
    }

    const methodDef = returnType
      ? `${methodName}${methodParams}: ${returnType}`
      : `${methodName}${methodParams}`;

    const methodCall = async function (_, params: any) {
      const result = await xrpc(
        "https://api.bsky.app",
        await import("./lexicons/" + lexiconPath.segments.join("/")),
        {
          params,
        },
      );
      if (!result.success) {
        throw new Error("Error in response: " + JSON.stringify(result.payload));
      }
      return result.payload.body;
    };

    return { methodDef, objectDefinitions, methodCall };
  } catch (e) {
    console.error(
      `Error processing lexicon ${lexiconPath.segments.join(".")}: ${e}`,
    );
    return null;
  }
}

export function generateDefinitions(lexiconIds: string[]): {
  chunks: string[];
  lexiconCalls: Record<string, Function>;
} {
  const chunks: string[] = [];
  const lexiconFields: Record<string, string> = {};
  let lexiconCalls: Record<string, Function> = {};
  let objectDefinitions: Record<string, LexValue> = {};

  for (const lexiconId of lexiconIds) {
    const lexiconPath = LexiconPathImpl.fromString(lexiconId);

    // Generate main lexicon structure
    const result = generateLexiconStructure(lexiconPath);
    if (result) {
      const {
        methodDef,
        objectDefinitions: collectedObjectDefinitions,
        methodCall,
      } = result;
      lexiconFields[lexiconPath.dotPath()] = methodDef;
      lexiconCalls[lexiconPath.dotPath()] = methodCall;
      objectDefinitions = {
        ...objectDefinitions,
        ...collectedObjectDefinitions,
      };
    }
  }

  // Walk output types for more refs
  const finalDefinitions: Record<string, LexValue> = {};
  let newDefinitions: Record<string, LexValue> = {};

  while (Object.keys(objectDefinitions).length > 0) {
    for (const [objName, objDef] of Object.entries(objectDefinitions)) {
      if (objDef.type === "object") {
        for (const [fieldName, fieldDef] of Object.entries(
          objDef.properties || {},
        )) {
          // Export refs
          if (fieldDef.type === "ref") {
            const ref = normalizeReference(objName, fieldDef.ref);
            const refDef = resolveReference(ref);
            if (refDef) {
              newDefinitions[ref] = refDef;
            }
            continue;
          } else if (
            fieldDef.type === "array" &&
            fieldDef.items?.type === "ref"
          ) {
            const ref = normalizeReference(objName, fieldDef.items.ref);
            const refDef = resolveReference(ref);
            if (
              refDef &&
              (refDef.type === "object" || refDef.type === "string")
            ) {
              newDefinitions[ref] = refDef;
            }
            continue;
          }
        }
      }

      finalDefinitions[objName] = objDef;
    }

    // Reset objectDefinitions to only be ones we haven't seen yet.
    objectDefinitions = {};
    for (const [objKey, objDef] of Object.entries(newDefinitions)) {
      if (!(objKey in finalDefinitions)) {
        objectDefinitions[objKey] = objDef;
      }
    }
  }

  objectDefinitions = finalDefinitions;

  // Generate types for all output types
  for (const [name, typeDef] of Object.entries(objectDefinitions)) {
    chunks.push(lexObjectDefinitionToGraphQL(name, typeDef) + "\n\n");
  }

  // Generate Lexicon structs iteratively.
  chunks.push(
    ...outputLexiconNamespaces(lexiconFields, new LexiconPathImpl([])),
  );

  return { chunks, lexiconCalls };
}

export async function readSchemaFiles(
  schemaFiles: string[] | null,
): Promise<string> {
  if (!schemaFiles || schemaFiles.length === 0) {
    return "";
  }

  let output = "";

  for (const schemaFile of schemaFiles) {
    const stats = fs.statSync(schemaFile);

    if (stats.isFile()) {
      // If it's a single file, read it
      const content = fs.readFileSync(schemaFile, "utf-8");
      output += `\n${content}\n`;
    } else if (stats.isDirectory()) {
      // If it's a directory, recursively read all .graphql files
      const files = await glob("**/*.graphql", { cwd: schemaFile });
      for (const file of files) {
        const fullPath = path.join(schemaFile, file);
        const content = fs.readFileSync(fullPath, "utf-8");
        output += `\n${content}\n`;
      }
    }
  }

  return output;
}