inline-schema/src/csv-loader/reference-resolver.ts

import * as fs from "fs";
import * as path from "path";
import { parseValue } from "../index.js";
import type {
  Schema,
  ReferenceSchema,
  ReverseReferenceSchema,
} from "../types.js";
import type { ReferenceFieldInfo } from "./types.js";
import { parseCsv } from "./loader.js";

/** Cache for loaded referenced tables */
const referenceTableCache = new Map<string, Record<string, unknown>[]>();

/**
 * Cache for reverse-reference lookups: filePath -> (foreignKey -> (fkValue -> rows[])).
 * Built once per referenced table + foreign key, mirroring what module-gen emits,
 * so resolveReverseReference doesn't re-filter the whole table per row.
 */
const reverseLookupCache = new Map<
  string,
  Map<string, Map<string, Record<string, unknown>[]>>
>();

/** Set of file paths currently being loaded (to detect circular references) */
const loadingFiles = new Set<string>();

export function hasNestedReferences(schema: Schema): boolean {
  switch (schema.type) {
    case "reference":
    case "reverseReference":
      return true;
    case "tuple":
      return schema.elements.some((el) => hasNestedReferences(el.schema));
    case "array":
      return hasNestedReferences(schema.element);
    case "union":
      return schema.members.some((m) => hasNestedReferences(m));
    default:
      return false;
  }
}

function loadReferenceTableData(
  schema: ReferenceSchema | ReverseReferenceSchema,
  refBaseDir: string | undefined,
  currentFilePath: string | undefined,
): {
  refTable: Record<string, unknown>[];
  refFilePath: string;
} {
  const baseDir =
    refBaseDir ||
    (currentFilePath ? path.dirname(currentFilePath) : process.cwd());
  const fileName = `${schema.tableName}.csv`;
  const refFilePath = path.isAbsolute(fileName)
    ? fileName
    : path.join(baseDir, fileName);

  let refTable: Record<string, unknown>[];
  if (referenceTableCache.has(refFilePath)) {
    refTable = referenceTableCache.get(refFilePath)!;
  } else {
    if (loadingFiles.has(refFilePath)) {
      throw new Error(
        `Circular reference detected: table "${schema.tableName}" (${refFilePath}) is already being loaded`,
      );
    }
    loadingFiles.add(refFilePath);
    try {
      const refContent = fs.readFileSync(refFilePath, "utf-8");
      const refResult = parseCsv(refContent, {
        currentFilePath: refFilePath,
        emitTypes: false,
      });
      refTable = refResult.data;
      referenceTableCache.set(refFilePath, refTable);
    } catch (error) {
      throw new Error(
        `Failed to load referenced table "${schema.tableName}" from ${refFilePath}: ${error instanceof Error ? error.message : String(error)}`,
      );
    } finally {
      loadingFiles.delete(refFilePath);
    }
  }

  return { refTable, refFilePath };
}

/**
 * Load a referenced table and build a forward lookup keyed by primary key.
 * Only call this when forward resolution is actually needed; the reverse
 * path uses `loadReferenceTableData` directly to avoid building a throwaway map.
 */
export function loadReferenceTable(
  schema: ReferenceSchema | ReverseReferenceSchema,
  refBaseDir: string | undefined,
  defaultPrimaryKey: string,
  currentFilePath: string | undefined,
): {
  lookup: Map<string, Record<string, unknown>>;
  refTable: Record<string, unknown>[];
  refFilePath: string;
} {
  const { refTable, refFilePath } = loadReferenceTableData(
    schema,
    refBaseDir,
    currentFilePath,
  );

  const lookup = new Map<string, Record<string, unknown>>();
  refTable.forEach((row) => {
    const pkValue = row[defaultPrimaryKey];
    if (pkValue !== undefined) {
      lookup.set(String(pkValue), row);
    }
  });

  return { lookup, refTable, refFilePath };
}

/**
 * Build (and cache) a reverse lookup for a referenced table + foreign key:
 * a map from the foreign-key value to the rows that reference it.
 */
function getReverseLookup(
  schema: ReverseReferenceSchema,
  refBaseDir: string | undefined,
  defaultPrimaryKey: string,
  currentFilePath: string | undefined,
): Map<string, Record<string, unknown>[]> {
  const { refTable, refFilePath } = loadReferenceTableData(
    schema,
    refBaseDir,
    currentFilePath,
  );

  let byForeignKey = reverseLookupCache.get(refFilePath);
  if (!byForeignKey) {
    byForeignKey = new Map();
    reverseLookupCache.set(refFilePath, byForeignKey);
  }

  const cached = byForeignKey.get(schema.foreignKey);
  if (cached) return cached;

  const lookup = new Map<string, Record<string, unknown>[]>();
  for (const row of refTable) {
    const fkValue = row[schema.foreignKey];
    const fkStr =
      fkValue !== null && fkValue !== undefined && typeof fkValue === "object"
        ? String((fkValue as Record<string, unknown>)[defaultPrimaryKey])
        : String(fkValue);
    const bucket = lookup.get(fkStr);
    if (bucket) {
      bucket.push(row);
    } else {
      lookup.set(fkStr, [row]);
    }
  }
  byForeignKey.set(schema.foreignKey, lookup);
  return lookup;
}

export function resolveReferenceId(
  id: string,
  lookup: Map<string, Record<string, unknown>>,
  tableName: string,
): Record<string, unknown> {
  const obj = lookup.get(id);
  if (!obj) {
    throw new Error(`Reference to "${tableName}" with id="${id}" not found`);
  }
  return obj;
}

export function parseReferenceIds(
  schema: ReferenceSchema,
  valueString: string,
): unknown {
  const trimmed = valueString.trim();
  if (schema.isOptional && trimmed === "") {
    return null;
  }
  return parseValue(schema, trimmed);
}

export function parseValueWithReferenceIds(
  valueString: string,
  schema: Schema,
): unknown {
  if (!hasNestedReferences(schema)) {
    return parseValue(schema, valueString);
  }

  switch (schema.type) {
    case "reference":
      return parseReferenceIds(schema, valueString);
    case "reverseReference":
      // Reverse references don't store IDs; they're derived at resolution time
      return null;
    case "tuple": {
      const parsed = parseValue(schema, valueString) as unknown[];
      return schema.elements.map((el, i) =>
        hasNestedReferences(el.schema)
          ? extractNestedReferenceIds(parsed[i], el.schema)
          : parsed[i],
      );
    }
    case "array": {
      const parsed = parseValue(schema, valueString) as unknown[];
      return parsed.map((item) =>
        hasNestedReferences(schema.element)
          ? extractNestedReferenceIds(item, schema.element)
          : item,
      );
    }
    case "union": {
      for (const member of schema.members) {
        if (hasNestedReferences(member)) {
          try {
            const parsed = parseValue(member, valueString);
            return extractNestedReferenceIds(parsed, member);
          } catch {}
        }
      }
      return parseValue(schema, valueString);
    }
    default:
      return parseValue(schema, valueString);
  }
}

export function extractNestedReferenceIds(
  value: unknown,
  schema: Schema,
): unknown {
  switch (schema.type) {
    case "reference":
      if (value === null || value === undefined) return value;
      if (schema.isArray) {
        const ids = Array.isArray(value) ? value : [value];
        return ids.map((id) => String(id));
      }
      return String(value);
    case "reverseReference":
      // Reverse references don't store IDs; return null placeholder
      return null;
    case "tuple": {
      if (!Array.isArray(value)) return value;
      return schema.elements.map((el, i) =>
        hasNestedReferences(el.schema)
          ? extractNestedReferenceIds(value[i], el.schema)
          : value[i],
      );
    }
    case "array": {
      if (!Array.isArray(value)) return value;
      return value.map((item) =>
        hasNestedReferences(schema.element)
          ? extractNestedReferenceIds(item, schema.element)
          : item,
      );
    }
    case "union": {
      for (const member of schema.members) {
        if (hasNestedReferences(member)) {
          try {
            return extractNestedReferenceIds(value, member);
          } catch {}
        }
      }
      return value;
    }
    default:
      return value;
  }
}

export function collectReferenceFields(
  schema: Schema,
  name: string,
): ReferenceFieldInfo[] {
  const fields: ReferenceFieldInfo[] = [];
  switch (schema.type) {
    case "reference":
      fields.push({
        name,
        tableName: schema.tableName,
        isArray: schema.isArray,
        schema,
      });
      break;
    case "reverseReference":
      fields.push({
        name,
        tableName: schema.tableName,
        isArray: true,
        foreignKey: schema.foreignKey,
        schema,
      });
      break;
    case "tuple":
      for (const el of schema.elements) {
        fields.push(...collectReferenceFields(el.schema, name));
      }
      break;
    case "array":
      fields.push(...collectReferenceFields(schema.element, name));
      break;
    case "union":
      for (const member of schema.members) {
        fields.push(...collectReferenceFields(member, name));
      }
      break;
  }
  return fields;
}

export function parseValueWithReferences(
  valueString: string,
  schema: Schema,
  refBaseDir: string | undefined,
  defaultPrimaryKey: string,
  currentFilePath: string | undefined,
  currentRowPk?: unknown,
): unknown {
  if (!hasNestedReferences(schema)) {
    return parseValue(schema, valueString);
  }

  switch (schema.type) {
    case "reference":
      return parseReferenceValue(
        schema,
        valueString,
        refBaseDir,
        defaultPrimaryKey,
        currentFilePath,
      );
    case "reverseReference": {
      if (currentRowPk === undefined) return [];
      return resolveReverseReference(
        schema,
        currentRowPk,
        refBaseDir,
        defaultPrimaryKey,
        currentFilePath,
      );
    }
    case "tuple": {
      const parsed = parseValue(schema, valueString) as unknown[];
      return schema.elements.map((el, i) =>
        resolveNestedReferences(
          parsed[i],
          el.schema,
          refBaseDir,
          defaultPrimaryKey,
          currentFilePath,
          currentRowPk,
        ),
      );
    }
    case "array": {
      const parsed = parseValue(schema, valueString) as unknown[];
      return parsed.map((item) =>
        resolveNestedReferences(
          item,
          schema.element,
          refBaseDir,
          defaultPrimaryKey,
          currentFilePath,
          currentRowPk,
        ),
      );
    }
    case "union": {
      // Try reference members first (more specific), then non-reference members.
      // This makes the fallback structural rather than error-message-driven.
      const refMembers = schema.members.filter(hasNestedReferences);
      const nonRefMembers = schema.members.filter((m) => !hasNestedReferences(m));
      const errors: Error[] = [];
      for (const member of [...refMembers, ...nonRefMembers]) {
        try {
          const parsed = parseValue(member, valueString);
          return resolveNestedReferences(
            parsed,
            member,
            refBaseDir,
            defaultPrimaryKey,
            currentFilePath,
            currentRowPk,
          );
        } catch (e) {
          errors.push(e instanceof Error ? e : new Error(String(e)));
        }
      }
      throw errors[0] ?? new Error("Value does not match any union member");
    }
    default:
      return parseValue(schema, valueString);
  }
}

export function resolveReverseReference(
  schema: ReverseReferenceSchema,
  pkValue: unknown,
  refBaseDir: string | undefined,
  defaultPrimaryKey: string,
  currentFilePath: string | undefined,
): Record<string, unknown>[] {
  const lookup = getReverseLookup(
    schema,
    refBaseDir,
    defaultPrimaryKey,
    currentFilePath,
  );
  return lookup.get(String(pkValue)) ?? [];
}

export function resolveNestedReferences(
  value: unknown,
  schema: Schema,
  refBaseDir: string | undefined,
  defaultPrimaryKey: string,
  currentFilePath: string | undefined,
  currentRowPk?: unknown,
): unknown {
  switch (schema.type) {
    case "reference": {
      if (value === null || value === undefined) return value;
      const { lookup } = loadReferenceTable(
        schema,
        refBaseDir,
        defaultPrimaryKey,
        currentFilePath,
      );
      if (schema.isArray) {
        const ids = Array.isArray(value) ? value : [value];
        return ids.map((id) =>
          resolveReferenceId(String(id), lookup, schema.tableName),
        );
      }
      return resolveReferenceId(String(value), lookup, schema.tableName);
    }
    case "reverseReference": {
      if (currentRowPk === undefined) return [];
      const results = resolveReverseReference(
        schema,
        currentRowPk,
        refBaseDir,
        defaultPrimaryKey,
        currentFilePath,
      );
      return results;
    }
    case "tuple": {
      if (!Array.isArray(value)) return value;
      return schema.elements.map((el, i) =>
        resolveNestedReferences(
          value[i],
          el.schema,
          refBaseDir,
          defaultPrimaryKey,
          currentFilePath,
          currentRowPk,
        ),
      );
    }
    case "array": {
      if (!Array.isArray(value)) return value;
      return value.map((item) =>
        resolveNestedReferences(
          item,
          schema.element,
          refBaseDir,
          defaultPrimaryKey,
          currentFilePath,
          currentRowPk,
        ),
      );
    }
    case "union": {
      // Try reference members first (more specific), then non-reference members.
      const refMembers = schema.members.filter(hasNestedReferences);
      const nonRefMembers = schema.members.filter((m) => !hasNestedReferences(m));
      const errors: Error[] = [];
      for (const member of [...refMembers, ...nonRefMembers]) {
        try {
          return resolveNestedReferences(
            value,
            member,
            refBaseDir,
            defaultPrimaryKey,
            currentFilePath,
            currentRowPk,
          );
        } catch (e) {
          errors.push(e instanceof Error ? e : new Error(String(e)));
        }
      }
      if (errors.length > 0) {
        throw errors[0];
      }
      return value;
    }
    default:
      return value;
  }
}

export function parseReferenceValue(
  schema: ReferenceSchema,
  valueString: string,
  refBaseDir: string | undefined,
  defaultPrimaryKey: string,
  currentFilePath: string | undefined,
): unknown {
  const trimmed = valueString.trim();
  if (schema.isOptional && trimmed === "") {
    return null;
  }

  const { lookup } = loadReferenceTable(
    schema,
    refBaseDir,
    defaultPrimaryKey,
    currentFilePath,
  );

  const ids = parseValue(schema, trimmed) as string | string[] | null;
  if (ids === null) return null;

  if (schema.isArray && Array.isArray(ids)) {
    return ids.map((id) => resolveReferenceId(id, lookup, schema.tableName));
  }

  return resolveReferenceId(ids as string, lookup, schema.tableName);
}