import type { Schema } from "../schema.ts";
import { isObject } from "./data.ts";
import { unwrap } from "./fields.ts";

type Path = readonly (string | number)[];
type Expand = (value: unknown, path: Path) => unknown;

const expanders = new WeakMap<Schema, Expand | undefined>();

function fail(path: Path, message: string): never {
  throw new Error(`${message} at ${[...path, "$type"].join(".")}`);
}

function children(schema: Schema): readonly Schema[] {
  switch (schema.type) {
    case "array":
      return [schema.items];
    case "object":
      return Object.values(schema.properties).map((field) => unwrap(field).schema);
    case "ref":
      return [schema.target];
    case "record":
      return [schema.record];
    case "union":
      return Object.values(schema.variants);
    default:
      return [];
  }
}

// Every schema reachable from the root, and those that rewrite values or contain one that
// does. Anything else is shared with the input rather than copied.
function survey(root: Schema): { visited: Set<Schema>; required: Set<Schema> } {
  const parents = new Map<Schema, Set<Schema>>([[root, new Set()]]);
  const required = new Set<Schema>();
  for (const schema of parents.keys()) {
    if (expanders.has(schema)) {
      if (expanders.get(schema)) required.add(schema);
      continue;
    }
    if (schema.type === "record" || schema.type === "union") required.add(schema);
    for (const child of children(schema)) {
      const known = parents.get(child);
      if (known) known.add(schema);
      else parents.set(child, new Set([schema]));
    }
  }
  // Propagate through the whole cycle before deciding which branches can be shared.
  for (const schema of required) {
    for (const parent of parents.get(schema) ?? []) required.add(parent);
  }
  return { visited: new Set(parents.keys()), required };
}

function expanderFor(schema: Schema, build: (schema: Schema) => Expand | undefined): Expand {
  const passthrough: Expand = (value) => value;
  switch (schema.type) {
    case "array": {
      const item = build(schema.items) ?? passthrough;
      return (value, path) =>
        Array.isArray(value) ? value.map((element, i) => item(element, [...path, i])) : value;
    }
    case "object": {
      const fields = new Map<string, Expand>();
      for (const [key, field] of Object.entries(schema.properties)) {
        const expand = build(unwrap(field).schema);
        if (expand) fields.set(key, expand);
      }
      return (value, path) => {
        if (!isObject(value)) return value;
        return Object.fromEntries(
          Object.entries(value).map(([key, item]) => {
            const field = fields.get(key);
            return [key, field ? field(item, [...path, key]) : item];
          }),
        );
      };
    }
    case "ref":
      return build(schema.target) ?? passthrough;
    case "record": {
      const body = build(schema.record) ?? passthrough;
      return (value, path) => {
        if (!isObject(value)) return value;
        if (Object.hasOwn(value, "$type") && value.$type !== schema.id) {
          fail(path, `Expected ${schema.id}`);
        }
        return { ...(body(value, path) as object), $type: schema.id };
      };
    }
    case "union": {
      const { namespace } = schema;
      const variants = new Map(
        Object.entries(schema.variants).map(([id, variant]) => [id, build(variant)]),
      );
      return (value, path) => {
        if (!isObject(value) || typeof value.$type !== "string") return value;
        let $type = value.$type;
        if ($type.startsWith("#")) {
          if (namespace === undefined) fail(path, "Only namespaced unions accept short variants");
          $type = namespace + $type;
          if (!variants.has($type)) fail(path, `Unknown union variant ${value.$type}`);
        }
        const resolved = $type === value.$type ? value : { ...value, $type };
        const variant = variants.get($type);
        return variant ? variant(resolved, path) : resolved;
      };
    }
    default:
      return passthrough;
  }
}

// A value with its $type discriminators filled in and expanded, as far as the schema needs.
export function expand(schema: Schema, value: unknown): unknown {
  return expander(schema)?.(value, []) ?? value;
}

export function expander(root: Schema): Expand | undefined {
  if (expanders.has(root)) return expanders.get(root);
  const { visited, required } = survey(root);
  const compiled = new Map<Schema, Expand | undefined>();
  // Recursive schemas get a forwarder before their own expander exists.
  function build(schema: Schema): Expand | undefined {
    if (expanders.has(schema)) return expanders.get(schema);
    if (compiled.has(schema)) return compiled.get(schema);
    if (!required.has(schema)) return undefined;
    let expand: Expand;
    compiled.set(schema, (value, path) => expand(value, path));
    return expand = expanderFor(schema, build);
  }
  const result = build(root);
  for (const schema of visited) {
    if (!expanders.has(schema)) expanders.set(schema, compiled.get(schema));
  }
  return result;
}