/**
 * Lexicon-to-Zod codegen script.
 *
 * Fetches AT Protocol lexicon JSON files from the official atproto repo,
 * resolves cross-references, and generates Zod schemas.
 *
 * Usage: pnpm codegen
 */

import { writeFileSync, mkdirSync } from "node:fs";
import { resolve, dirname } from "node:path";
import { fileURLToPath } from "node:url";

// ---------------------------------------------------------------------------
// Configuration
// ---------------------------------------------------------------------------

const GITHUB_BASE =
  "https://raw.githubusercontent.com/bluesky-social/atproto/main/lexicons";

/** Record NSIDs that appear as Jetstream `collection` values. */
const ROOT_RECORDS: string[] = [
  "app.bsky.feed.post",
  "app.bsky.feed.like",
  "app.bsky.feed.repost",
  "app.bsky.feed.threadgate",
  "app.bsky.feed.postgate",
  "app.bsky.graph.follow",
  "app.bsky.graph.block",
  "app.bsky.graph.list",
  "app.bsky.graph.listitem",
  "app.bsky.graph.listblock",
  "app.bsky.actor.profile",
];

const __filename = fileURLToPath(import.meta.url);
const __dirname = dirname(__filename);

const OUTPUT_PATH = resolve(__dirname, "../src/generated/lexicons.ts");

// ---------------------------------------------------------------------------
// Lexicon JSON types (internal to codegen)
// ---------------------------------------------------------------------------

interface LexiconDoc {
  lexicon: number;
  id: string;
  defs: Record<string, LexiconDef>;
}

type LexiconDef =
  | RecordDef
  | ObjectDef
  | StringDef
  | IntegerDef
  | BooleanDef
  | ArrayDef
  | UnionDef
  | RefDef
  | TokenDef
  | BlobDef
  | BytesDef
  | UnknownDef
  | CidLinkDef;

interface RecordDef {
  type: "record";
  key?: string;
  description?: string;
  record: ObjectDef;
}
interface ObjectDef {
  type: "object";
  required?: string[];
  nullable?: string[];
  properties: Record<string, LexiconDef>;
  description?: string;
}
interface StringDef {
  type: "string";
  format?: string;
  maxLength?: number;
  minLength?: number;
  maxGraphemes?: number;
  knownValues?: string[];
  enum?: string[];
  const?: string;
  default?: string;
  description?: string;
}
interface IntegerDef {
  type: "integer";
  minimum?: number;
  maximum?: number;
  const?: number;
  default?: number;
  enum?: number[];
  description?: string;
}
interface BooleanDef {
  type: "boolean";
  const?: boolean;
  default?: boolean;
  description?: string;
}
interface ArrayDef {
  type: "array";
  items: LexiconDef;
  maxLength?: number;
  minLength?: number;
  description?: string;
}
interface UnionDef {
  type: "union";
  refs: string[];
  closed?: boolean;
  description?: string;
}
interface RefDef {
  type: "ref";
  ref: string;
  description?: string;
}
interface TokenDef {
  type: "token";
  description?: string;
}
interface BlobDef {
  type: "blob";
  accept?: string[];
  maxSize?: number;
  description?: string;
}
interface BytesDef {
  type: "bytes";
  description?: string;
}
interface UnknownDef {
  type: "unknown";
  description?: string;
}
interface CidLinkDef {
  type: "cid-link";
  description?: string;
}

// ---------------------------------------------------------------------------
// Phase 2: Fetch lexicons
// ---------------------------------------------------------------------------

const lexiconCache = new Map<string, LexiconDoc>();

function nsidToPath(nsid: string): string {
  return nsid.replace(/\./g, "/");
}

async function fetchLexicon(nsid: string): Promise<LexiconDoc> {
  const cached = lexiconCache.get(nsid);
  if (cached) return cached;

  const url = `${GITHUB_BASE}/${nsidToPath(nsid)}.json`;
  console.log(`  fetch ${nsid}`);
  const res = await fetch(url);
  if (!res.ok) {
    throw new Error(`Failed to fetch ${url}: ${res.status} ${res.statusText}`);
  }
  const doc = (await res.json()) as LexiconDoc;
  lexiconCache.set(nsid, doc);
  return doc;
}

// ---------------------------------------------------------------------------
// Phase 3: Resolve references — collect all reachable definitions
// ---------------------------------------------------------------------------

/** Fully-qualified def id, e.g. "app.bsky.feed.post#main" */
type QualifiedId = string;

/** Resolved definition with metadata */
interface ResolvedDef {
  qualifiedId: QualifiedId;
  nsid: string;
  defName: string;
  def: LexiconDef;
}

const resolvedDefs = new Map<QualifiedId, ResolvedDef>();
const visited = new Set<QualifiedId>();

function qualifiedId(nsid: string, defName: string): QualifiedId {
  return `${nsid}#${defName}`;
}

/**
 * Resolve a ref string relative to a context NSID.
 * Returns [nsid, defName].
 */
function parseRef(ref: string, contextNsid: string): [string, string] {
  if (ref.startsWith("#")) {
    return [contextNsid, ref.slice(1)];
  }
  const hashIdx = ref.indexOf("#");
  if (hashIdx !== -1) {
    return [ref.slice(0, hashIdx), ref.slice(hashIdx + 1)];
  }
  return [ref, "main"];
}

/**
 * Recursively resolve a definition and all its transitive dependencies.
 * Skips defs whose name contains "view" (case-insensitive) at the top level
 * to avoid pulling in API response types that have circular refs.
 */
async function resolveDef(
  nsid: string,
  defName: string,
  isRoot: boolean = false,
): Promise<void> {
  const qid = qualifiedId(nsid, defName);
  if (visited.has(qid)) return;
  visited.add(qid);

  // Skip view defs unless they are explicitly part of a root record
  if (
    !isRoot &&
    defName !== "main" &&
    /view/i.test(defName)
  ) {
    return;
  }

  const doc = await fetchLexicon(nsid);
  const def = doc.defs[defName];
  if (!def) {
    console.warn(`  warning: definition ${qid} not found, skipping`);
    return;
  }

  resolvedDefs.set(qid, { qualifiedId: qid, nsid, defName, def });

  // Walk refs within this definition
  const refs = collectRefs(def, nsid);
  for (const [refNsid, refDefName] of refs) {
    await resolveDef(refNsid, refDefName);
  }
}

/** Collect all ref targets from a definition. */
function collectRefs(
  def: LexiconDef,
  contextNsid: string,
): Array<[string, string]> {
  const refs: Array<[string, string]> = [];

  function walk(d: LexiconDef): void {
    if (!d || typeof d !== "object") return;

    switch (d.type) {
      case "ref":
        refs.push(parseRef(d.ref, contextNsid));
        break;
      case "union":
        for (const r of d.refs) {
          refs.push(parseRef(r, contextNsid));
        }
        break;
      case "array":
        walk(d.items);
        break;
      case "object":
        if (d.properties) {
          for (const prop of Object.values(d.properties)) {
            walk(prop);
          }
        }
        break;
      case "record":
        walk(d.record);
        break;
    }
  }

  walk(def);
  return refs;
}

// ---------------------------------------------------------------------------
// Phase 4: Track union members
// ---------------------------------------------------------------------------

function findUnionMembers(): Set<QualifiedId> {
  const members = new Set<QualifiedId>();

  function walkForUnions(def: LexiconDef, contextNsid: string): void {
    if (!def || typeof def !== "object") return;

    switch (def.type) {
      case "union":
        for (const r of def.refs) {
          const [rNsid, rDef] = parseRef(r, contextNsid);
          members.add(qualifiedId(rNsid, rDef));
        }
        break;
      case "array":
        walkForUnions(def.items, contextNsid);
        break;
      case "object":
        if (def.properties) {
          for (const prop of Object.values(def.properties)) {
            walkForUnions(prop, contextNsid);
          }
        }
        break;
      case "record":
        walkForUnions(def.record, contextNsid);
        break;
    }
  }

  for (const [, resolved] of resolvedDefs) {
    walkForUnions(resolved.def, resolved.nsid);
  }

  return members;
}

// ---------------------------------------------------------------------------
// Phase 5: Topological sort
// ---------------------------------------------------------------------------

function topologicalSort(): QualifiedId[] {
  const adjacency = new Map<QualifiedId, Set<QualifiedId>>();
  const inDegree = new Map<QualifiedId, number>();

  for (const qid of resolvedDefs.keys()) {
    adjacency.set(qid, new Set());
    inDegree.set(qid, 0);
  }

  // Build adjacency edges (Set deduplicates)
  for (const [qid, resolved] of resolvedDefs) {
    const refs = collectRefs(resolved.def, resolved.nsid);
    for (const [rNsid, rDefName] of refs) {
      const depQid = qualifiedId(rNsid, rDefName);
      if (resolvedDefs.has(depQid) && depQid !== qid) {
        adjacency.get(depQid)!.add(qid);
      }
    }
  }

  // Compute in-degree from adjacency (avoids duplicate-ref inflation)
  for (const neighbors of adjacency.values()) {
    for (const neighbor of neighbors) {
      inDegree.set(neighbor, (inDegree.get(neighbor) ?? 0) + 1);
    }
  }

  // Kahn's algorithm
  const queue: QualifiedId[] = [];
  for (const [qid, deg] of inDegree) {
    if (deg === 0) queue.push(qid);
  }

  const sorted: QualifiedId[] = [];
  while (queue.length > 0) {
    const node = queue.shift()!;
    sorted.push(node);
    for (const neighbor of adjacency.get(node) ?? []) {
      const newDeg = (inDegree.get(neighbor) ?? 1) - 1;
      inDegree.set(neighbor, newDeg);
      if (newDeg === 0) queue.push(neighbor);
    }
  }

  // If some nodes are left (cycle), append them with a warning
  for (const qid of resolvedDefs.keys()) {
    if (!sorted.includes(qid)) {
      console.warn(`  warning: cycle detected involving ${qid}, using z.lazy()`);
      sorted.push(qid);
    }
  }

  return sorted;
}

// ---------------------------------------------------------------------------
// Phase 6: Code generation
// ---------------------------------------------------------------------------

/** Convert a qualified ID to a camelCase variable name. */
function qidToVarName(qid: QualifiedId): string {
  const [nsid, defName] = qid.split("#");
  const parts = nsid.split(".");
  let name = parts
    .map((p, i) => (i === 0 ? p : p.charAt(0).toUpperCase() + p.slice(1)))
    .join("");
  if (defName && defName !== "main") {
    name += defName.charAt(0).toUpperCase() + defName.slice(1);
  }
  return name;
}

/** Convert a qualified ID to a PascalCase type name. */
function qidToTypeName(qid: QualifiedId): string {
  const varName = qidToVarName(qid);
  return varName.charAt(0).toUpperCase() + varName.slice(1);
}

/**
 * Get the $type string for a definition used in unions.
 * For "main" defs: just the NSID. For named defs: "nsid#name".
 */
function typeTag(nsid: string, defName: string): string {
  return defName === "main" ? nsid : `${nsid}#${defName}`;
}

function generateCode(
  sortedIds: QualifiedId[],
  unionMembers: Set<QualifiedId>,
): string {
  const lines: string[] = [];

  lines.push("// AUTO-GENERATED FILE — DO NOT EDIT");
  lines.push(
    "// Generated by scripts/codegen.ts from AT Protocol lexicon definitions",
  );
  lines.push(
    "// Source: https://github.com/bluesky-social/atproto/tree/main/lexicons",
  );
  lines.push("");
  lines.push('import { z } from "zod";');
  lines.push("");

  // Shared blob schema
  lines.push("// Blob wire format (as it appears in Jetstream JSON)");
  lines.push("const blobSchema = z.object({");
  lines.push('  $type: z.literal("blob"),');
  lines.push("  ref: z.object({ $link: z.string() }),");
  lines.push("  mimeType: z.string(),");
  lines.push("  size: z.number(),");
  lines.push("});");
  lines.push("");

  // Generate each schema
  for (const qid of sortedIds) {
    const resolved = resolvedDefs.get(qid)!;
    const varName = qidToVarName(qid);
    const code = generateDefCode(resolved, unionMembers);
    lines.push(`export const ${varName} = ${code};`);
    lines.push("");
  }

  // Type exports
  lines.push("// Type exports");
  for (const qid of sortedIds) {
    const varName = qidToVarName(qid);
    const typeName = qidToTypeName(qid);
    lines.push(
      `export type ${typeName} = z.infer<typeof ${varName}>;`,
    );
  }
  lines.push("");

  // Collection schema mapping
  lines.push("// Collection NSID → record schema mapping");
  lines.push("export const collectionSchemas = {");
  for (const nsid of ROOT_RECORDS) {
    const qid = qualifiedId(nsid, "main");
    const varName = qidToVarName(qid);
    lines.push(`  "${nsid}": ${varName},`);
  }
  lines.push("} as const;");
  lines.push("");
  lines.push(
    "export type CollectionNsid = keyof typeof collectionSchemas;",
  );
  lines.push("");

  return lines.join("\n");
}

function generateDefCode(
  resolved: ResolvedDef,
  unionMembers: Set<QualifiedId>,
): string {
  const { def, nsid, defName, qualifiedId: qid } = resolved;

  switch (def.type) {
    case "record":
      return generateObjectCode(
        def.record,
        nsid,
        defName,
        unionMembers,
        true,
      );
    case "object":
      return generateObjectCode(
        def,
        nsid,
        defName,
        unionMembers,
        false,
      );
    case "string":
      return generateStringCode(def);
    case "integer":
      return generateIntegerCode(def);
    case "boolean":
      return generateBooleanCode(def);
    case "array":
      return generateArrayCode(def, nsid);
    case "union":
      return generateUnionCode(def, nsid);
    case "ref":
      return generateRefCode(def, nsid);
    case "token":
      return `z.literal("${typeTag(nsid, defName)}")`;
    case "blob":
      return "blobSchema";
    case "bytes":
      return "z.unknown()";
    case "cid-link":
      return "z.string()";
    case "unknown":
      return "z.unknown()";
    default:
      console.warn(`  warning: unknown type "${(def as any).type}", using z.unknown()`);
      return "z.unknown()";
  }
}

function generateObjectCode(
  def: ObjectDef,
  nsid: string,
  defName: string,
  unionMembers: Set<QualifiedId>,
  isRecord: boolean,
): string {
  const qid = qualifiedId(nsid, defName);
  const requiredSet = new Set(def.required ?? []);
  const props: string[] = [];

  // Add $type for records (always) or union members
  if (isRecord) {
    props.push(`  $type: z.literal("${nsid}")`);
  } else if (unionMembers.has(qid)) {
    props.push(`  $type: z.literal("${typeTag(nsid, defName)}")`);
  }

  if (def.properties) {
    for (const [propName, propDef] of Object.entries(def.properties)) {
      const fieldCode = generateFieldCode(propDef, nsid);
      const isRequired = requiredSet.has(propName);
      const safeName = propName.startsWith("$") ? `"${propName}"` : propName;
      if (isRequired) {
        props.push(`  ${safeName}: ${fieldCode}`);
      } else {
        props.push(`  ${safeName}: ${fieldCode}.optional()`);
      }
    }
  }

  return `z.object({\n${props.join(",\n")},\n})`;
}

function generateFieldCode(def: LexiconDef, contextNsid: string): string {
  switch (def.type) {
    case "string":
      return generateStringCode(def);
    case "integer":
      return generateIntegerCode(def);
    case "boolean":
      return generateBooleanCode(def);
    case "array":
      return generateArrayCode(def, contextNsid);
    case "union":
      return generateUnionCode(def, contextNsid);
    case "ref":
      return generateRefCode(def, contextNsid);
    case "blob":
      return "blobSchema";
    case "bytes":
      return "z.unknown()";
    case "cid-link":
      return "z.string()";
    case "unknown":
      return "z.unknown()";
    case "object":
      // Inline object (rare but possible)
      return generateInlineObjectCode(def, contextNsid);
    default:
      return "z.unknown()";
  }
}

function generateInlineObjectCode(
  def: ObjectDef,
  contextNsid: string,
): string {
  const requiredSet = new Set(def.required ?? []);
  const props: string[] = [];

  for (const [propName, propDef] of Object.entries(def.properties ?? {})) {
    const fieldCode = generateFieldCode(propDef, contextNsid);
    const isRequired = requiredSet.has(propName);
    const safeName = propName.startsWith("$") ? `"${propName}"` : propName;
    if (isRequired) {
      props.push(`  ${safeName}: ${fieldCode}`);
    } else {
      props.push(`  ${safeName}: ${fieldCode}.optional()`);
    }
  }

  return `z.object({\n${props.join(",\n")},\n})`;
}

function generateStringCode(def: StringDef): string {
  if (def.const !== undefined) {
    return `z.literal("${def.const}")`;
  }
  if (def.enum) {
    const values = def.enum.map((v) => `"${v}"`).join(", ");
    return `z.enum([${values}])`;
  }
  let code = "z.string()";
  if (def.maxLength !== undefined) {
    code += `.max(${def.maxLength})`;
  }
  if (def.minLength !== undefined) {
    code += `.min(${def.minLength})`;
  }
  if (def.default !== undefined) {
    code += `.default("${def.default}")`;
  }
  return code;
}

function generateIntegerCode(def: IntegerDef): string {
  if (def.const !== undefined) {
    return `z.literal(${def.const})`;
  }
  if (def.enum) {
    const values = def.enum.map((v) => `${v}`).join(", ");
    return `z.enum([${values}])`;
  }
  let code = "z.number().int()";
  if (def.minimum !== undefined) {
    code += `.min(${def.minimum})`;
  }
  if (def.maximum !== undefined) {
    code += `.max(${def.maximum})`;
  }
  if (def.default !== undefined) {
    code += `.default(${def.default})`;
  }
  return code;
}

function generateBooleanCode(def: BooleanDef): string {
  if (def.const !== undefined) {
    return `z.literal(${def.const})`;
  }
  let code = "z.boolean()";
  if (def.default !== undefined) {
    code += `.default(${def.default})`;
  }
  return code;
}

function generateArrayCode(def: ArrayDef, contextNsid: string): string {
  const itemCode = generateFieldCode(def.items, contextNsid);
  let code = `z.array(${itemCode})`;
  if (def.maxLength !== undefined) {
    code += `.max(${def.maxLength})`;
  }
  if (def.minLength !== undefined) {
    code += `.min(${def.minLength})`;
  }
  return code;
}

function generateUnionCode(def: UnionDef, contextNsid: string): string {
  const memberRefs = def.refs.map((r) => {
    const [rNsid, rDefName] = parseRef(r, contextNsid);
    const qid = qualifiedId(rNsid, rDefName);
    return qidToVarName(qid);
  });

  // Filter to only members we actually resolved (skip view types etc.)
  const availableMembers = memberRefs.filter((varName) => {
    // Check if this variable exists in our resolved defs
    for (const [, resolved] of resolvedDefs) {
      if (qidToVarName(resolved.qualifiedId) === varName) return true;
    }
    return false;
  });

  if (availableMembers.length === 0) {
    return "z.unknown()";
  }
  if (availableMembers.length === 1) {
    return availableMembers[0];
  }

  return `z.discriminatedUnion("$type", [${availableMembers.join(", ")}])`;
}

function generateRefCode(def: RefDef, contextNsid: string): string {
  const [rNsid, rDefName] = parseRef(def.ref, contextNsid);
  const qid = qualifiedId(rNsid, rDefName);

  if (resolvedDefs.has(qid)) {
    return qidToVarName(qid);
  }

  // Unresolved ref — fall back to z.unknown()
  console.warn(`  warning: unresolved ref ${qid}, using z.unknown()`);
  return "z.unknown()";
}

// ---------------------------------------------------------------------------
// Main
// ---------------------------------------------------------------------------

async function main(): Promise<void> {
  console.log("Fetching lexicons...");
  for (const nsid of ROOT_RECORDS) {
    await resolveDef(nsid, "main", true);
  }

  console.log(`\nResolved ${resolvedDefs.size} definitions`);

  const unionMembers = findUnionMembers();
  console.log(`Found ${unionMembers.size} union member definitions`);

  const sorted = topologicalSort();
  console.log(`Topologically sorted ${sorted.length} definitions`);

  console.log("\nGenerating code...");
  const code = generateCode(sorted, unionMembers);

  mkdirSync(dirname(OUTPUT_PATH), { recursive: true });
  writeFileSync(OUTPUT_PATH, code, "utf-8");
  console.log(`\nWritten to ${OUTPUT_PATH}`);
}

main().catch((err) => {
  console.error("Codegen failed:", err);
  process.exit(1);
});