import type {
  UdlAction,
  UdlDocument,
  UdlInstrument,
  UdlValue,
} from "./schema.js";
import { targetIds } from "./primitives.js";
import type { Stranded } from "./diagnostics.js";

export interface FinanceIssue {
  path: string;
  message: string;
  stranded?: Stranded;
}
type Sum = Map<string, bigint>;
type Balance = Sum | null;
interface Hold {
  from: string;
  to: string;
  amount: Balance;
}
interface State {
  status: string;
  actions: string[];
  paths: Map<string, string[]>;
  balances: Map<string, Balance>;
  holds: Map<string, Hold>;
}
const add = (a: Balance, b: Balance, sign = 1n): Balance => {
  if (a === null || b === null) return null;
  const result = new Map(a);
  for (const [key, value] of b) {
    const next = (result.get(key) ?? 0n) + value * sign;
    if (next) result.set(key, next);
    else result.delete(key);
  }
  return result;
};
const show = (sum: Balance): string =>
  sum === null
    ? "unknown"
    : JSON.stringify(
        [...sum]
          .sort(([a], [b]) => (a < b ? -1 : a > b ? 1 : 0))
          .map(([key, value]) => [key, String(value)]),
      );
const same = (a: Balance, b: Balance) => show(a) === show(b);

/** Join loop balances conservatively; a balance equality proves an unknown balance empty. */
export function analyzeInstrumentFinance(
  instrument: UdlInstrument,
  document?: UdlDocument,
): FinanceIssue[] {
  const owned = new Set(
    instrument.fields
      .filter((f) => f.type === "account" && f.owner === "self")
      .map((f) => `self.${f.name}`),
  );
  /** The owned account a `.balance` read names; this pass tracks no other account. */
  const ownedAccount = (field: string): string | undefined => {
    if (!field.endsWith(".balance")) return undefined;
    const account = field.slice(0, -".balance".length);
    return owned.has(account) ? account : undefined;
  };
  const shared = new Set<string>();
  for (const owner of document?.instruments ?? []) {
    if (owner.id === instrument.id) continue;
    for (const action of Object.values(owner.actions))
      for (const move of action.moves)
        if ("amount" in move)
          for (const endpoint of [move.from, move.to]) {
            const parts = endpoint.split(".");
            if (parts.shift() !== "self") continue;
            let targets = [owner];
            while (parts.length > 1) {
              const key = parts.shift();
              targets = targets.flatMap((target) => {
                const reference = target.fields.find(
                  (field) => field.name === key,
                );
                if (reference?.type !== "ref") return [];
                const ids = targetIds(reference.target);
                return (
                  document?.instruments.filter((candidate) =>
                    ids.includes(candidate.id),
                  ) ?? []
                );
              });
            }
            if (
              parts.length === 1 &&
              targets.some((target) => target.id === instrument.id) &&
              owned.has(`self.${parts[0]}`)
            )
              shared.add(`self.${parts[0]}`);
          }
  }
  const problems: FinanceIssue[] = [];
  const report = (message: string) =>
    problems.push({ path: ".actions", message });
  // A field some action sets is a free symbol: its default no longer holds.
  const assigned = new Set(
    Object.values(instrument.actions).flatMap((action) =>
      Object.keys(action.set ?? {}).map((key) => `self.${key}`),
    ),
  );
  const reads = (calc: object): string[] =>
    JSON.stringify(calc)
      .match(/"field":"[^"]+"/g)
      ?.map((match) => match.slice(9, -1)) ?? [];
  // Calculated fields downstream of the seeds, following opaque and linear calculations.
  const downstream = (seeds: (field: string) => boolean): Set<string> => {
    const found = new Set<string>();
    for (let grew = true; grew;) {
      grew = false;
      for (const calc of instrument.calculate) {
        const target = `self.${calc.target}`;
        if (found.has(target)) continue;
        if (reads(calc).some((field) => seeds(field) || found.has(field))) {
          found.add(target);
          grew = true;
        }
      }
    }
    return found;
  };
  const balanceBound = downstream(
    (field) => field.endsWith(".balance") || field.endsWith(".reserved"),
  );
  const equations = new Map<string, Sum>();
  for (const field of instrument.fields)
    if (
      field.type === "money" &&
      field.value !== undefined &&
      !assigned.has(`self.${field.name}`)
    )
      equations.set(
        `self.${field.name}`,
        new Map([["#", BigInt(field.value)]]),
      );
  for (const calc of instrument.calculate) {
    if (calc.op === "sum" || calc.op === "subtract") {
      const terms =
        calc.op === "sum"
          ? calc.values.map((v) => [v, 1n] as const)
          : [
              [calc.base, 1n] as const,
              ...calc.subtract.map((v) => [v, -1n] as const),
            ];
      let result: Balance = new Map();
      for (const [value, sign] of terms)
        result = add(result, raw(value, "create"), sign);
      equations.set(`self.${calc.target}`, result!);
    }
  }
  function raw(value: UdlValue, invocation: string): Sum {
    return "field" in value
      ? new Map([
          [
            value.field.startsWith("input.")
              ? `${invocation}:${value.field}`
              : value.field,
            1n,
          ],
        ])
      : new Map([["#", BigInt(String(value.literal))]]);
  }
  function expand(sum: Sum, seen = new Set<string>()): Sum {
    let result: Balance = new Map();
    for (const [key, coefficient] of sum) {
      const equation = equations.get(key);
      result = add(
        result,
        equation && !seen.has(key)
          ? expand(equation, new Set([...seen, key]))
          : new Map([[key, 1n]]),
        coefficient,
      );
    }
    return result!;
  }
  // A top-level create admits the snapshot it reads; an invoked child create
  // reads its parent's live subject but keeps the parent's snapshot.
  const invokedCreate = (document?.instruments ?? []).some((owner) =>
    Object.values(owner.actions).some((candidate) =>
      (candidate.invoke ?? []).some(
        (call) =>
          "instrument" in call &&
          call.instrument === instrument.id &&
          call.action === "create",
      ),
    ),
  );
  let invocation = 0;
  const apply = (state: State, action: UdlAction, name: string): void => {
    const scope = `${name}#${invocation++}`;
    // An unknown balance becomes one opening symbol, so a whole-balance move cancels it.
    const opening = `${scope}:opening:`;
    for (const [account, balance] of state.balances)
      if (balance === null)
        state.balances.set(account, new Map([[`${opening}${account}`, 1n]]));
    const start = new Map(state.balances);
    // Moves read set fields after the set, so they get this invocation's symbol.
    const changed = new Set(
      Object.keys(action.set ?? {}).map((key) => `self.${key}`),
    );
    const stale = downstream((field) => changed.has(field));
    let fresh = new Set<string>();
    // An action that declares subject fields reads the live object, not the agreement.
    const declared = new Set(
      name === "create" && !invokedCreate
        ? []
        : (action.subject?.requirements ?? []).map(
            (entry) => `subject.${entry.field.name}`,
          ),
    );
    const live = new Set([
      ...declared,
      ...downstream((field) => declared.has(field)),
    ]);
    let reading = 0;
    // The pass keeps no reserved total per account, so a reserved read is external even on an owned account.
    const external = (field: string) =>
      field.endsWith(".reserved") ||
      (field.endsWith(".balance") && ownedAccount(field) === undefined);
    const resolve = (sum: Sum): Balance => {
      let result: Balance = new Map();
      for (const [key, coefficient] of sum) {
        const account = ownedAccount(key);
        const scoped =
          balanceBound.has(key) ||
          fresh.has(key) ||
          live.has(key) ||
          external(key);
        const term: Balance =
          account !== undefined
            ? (start.get(account) ?? null)
            : new Map([[scoped ? `${scope}:${key}` : key, 1n]]);
        result = add(result, term, coefficient);
      }
      return result;
    };
    const calculated = new Map<string, Balance>();
    const local = new Map(
      (action.calculate ?? []).map((c) => [`self.${c.target}`, c]),
    );
    const amount = (value: UdlValue, seen = new Set<string>()): Balance => {
      if ("field" in value) {
        if (calculated.has(value.field)) return calculated.get(value.field)!;
        const account = ownedAccount(value.field);
        if (account !== undefined) return state.balances.get(account) ?? null;
        // Another account's balance moves between reads, so each read is its own symbol.
        if (external(value.field))
          return new Map([[`${scope}:${reading++}:${value.field}`, 1n]]);
        const calculation = local.get(value.field);
        if (calculation && !seen.has(value.field)) {
          const next = new Set([...seen, value.field]);
          if (calculation.op === "sum")
            return calculation.values.reduce<Balance>(
              (sum, value) => add(sum, amount(value, next)),
              new Map(),
            );
          if (calculation.op === "subtract")
            return calculation.subtract.reduce<Balance>(
              (sum, value) => add(sum, amount(value, next), -1n),
              amount(calculation.base, next),
            );
          return new Map([[`${scope}:${value.field}`, 1n]]);
        }
      }
      return resolve(expand(raw(value, scope)));
    };
    for (const requirement of action.requires)
      if (requirement.kind === "compare" && requirement.operator === "==") {
        for (const [left, right] of [
          [requirement.left, requirement.right],
          [requirement.right, requirement.left],
        ] as const) {
          const account =
            "field" in left ? ownedAccount(left.field) : undefined;
          if (account !== undefined) state.balances.set(account, amount(right));
        }
      }
    for (const key of local.keys()) calculated.set(key, amount({ field: key }));
    const debit = (account: string, value: Balance) => {
      if (owned.has(account))
        state.balances.set(
          account,
          add(
            state.balances.has(account)
              ? state.balances.get(account)!
              : new Map(),
            value,
            -1n,
          ),
        );
    };
    const credit = (account: string, value: Balance) => {
      if (owned.has(account))
        state.balances.set(
          account,
          add(
            state.balances.has(account)
              ? state.balances.get(account)!
              : new Map(),
            value,
          ),
        );
    };
    fresh = new Set([...changed, ...stale]);
    for (const move of action.moves) {
      if ("amount" in move) {
        const value = amount(move.amount);
        debit(move.from, value);
        if (move.operation === "internal_transfer.reserve") {
          const key = `self.${move.capture}`;
          if (state.holds.has(key))
            report(`${move.capture} overwrites a live reservation`);
          state.holds.set(key, { from: move.from, to: move.to, amount: value });
        } else credit(move.to, value);
      } else {
        const hold = state.holds.get(move.transfer);
        if (!hold) {
          report(`${move.key} consumes a missing reservation ${move.transfer}`);
          continue;
        }
        credit(
          move.operation === "internal_transfer.post" ? hold.to : hold.from,
          hold.amount,
        );
        state.holds.delete(move.transfer);
      }
    }
    // Old values of set fields are gone; this invocation's values become current.
    const settle = (sum: Balance): Balance =>
      sum === null ||
      [...sum.keys()].some((key) => key.startsWith(opening) || fresh.has(key))
        ? null
        : new Map(
            [...sum].map(([key, value]) => {
              const field = key.slice(scope.length + 1);
              return key.startsWith(`${scope}:`) &&
                fresh.has(field) &&
                !balanceBound.has(field)
                ? [field, value]
                : [key, value];
            }),
          );
    for (const [account, balance] of state.balances)
      state.balances.set(account, settle(balance));
    for (const hold of state.holds.values()) hold.amount = settle(hold.amount);
  };
  const initial: State = {
    status: instrument.lifecycle.initial,
    actions: [],
    paths: new Map(),
    balances: new Map([...owned].map((key) => [key, new Map()])),
    holds: new Map(),
  };
  if (instrument.actions.create) {
    apply(initial, instrument.actions.create, "create");
    for (const [account, balance] of initial.balances)
      if (balance === null || balance.size)
        initial.paths.set(account, ["create"]);
  }
  const pending = [initial];
  const reached = new Map<string, State>();
  const copy = (state: State): State => ({
    status: state.status,
    actions: [...state.actions],
    paths: new Map(state.paths),
    balances: new Map(state.balances),
    holds: new Map([...state.holds].map(([key, hold]) => [key, { ...hold }])),
  });
  let work = 0;
  while (pending.length && !problems.length) {
    if (++work > 4096) {
      report("owned proof exceeds its finite work budget");
      break;
    }
    let state = pending.pop()!;
    const key = JSON.stringify([state.status, [...state.holds.keys()].sort()]);
    const previous = reached.get(key);
    if (previous) {
      const joined = copy(previous);
      let changed = false;
      for (const account of owned) {
        const before = previous.balances.has(account)
          ? previous.balances.get(account)!
          : new Map();
        const after = state.balances.has(account)
          ? state.balances.get(account)!
          : new Map();
        if (before !== null && !same(before, after)) {
          joined.balances.set(account, null);
          if (!before.size && state.paths.has(account))
            joined.paths.set(account, state.paths.get(account)!);
          changed = true;
        }
      }
      for (const [name, hold] of joined.holds) {
        const next = state.holds.get(name)!;
        if (hold.from !== next.from || hold.to !== next.to) {
          report(`${name} changes reservation endpoints across paths`);
          break;
        }
        if (hold.amount !== null && !same(hold.amount, next.amount)) {
          hold.amount = null;
          changed = true;
        }
      }
      if (!changed) continue;
      state = joined;
    }
    if (reached.size >= 256 && !previous) {
      report("owned proof exceeds 256 state and reservation combinations");
      break;
    }
    reached.set(key, copy(state));
    const edges = Object.entries(instrument.lifecycle.transitions).filter(
      ([, edge]) => edge.from.includes(state.status),
    );
    if (
      !edges.length &&
      ([...state.balances].some(
        ([account, balance]) =>
          !shared.has(account) && (balance === null || balance.size > 0),
      ) ||
        state.holds.size)
    )
      problems.push({
        path: ".actions",
        message: `terminal state ${state.status} may strand owned money or an open reservation`,
        stranded: {
          state: state.status,
          accounts: [...state.balances]
            .filter(
              ([account, balance]) =>
                !shared.has(account) && (balance === null || balance.size > 0),
            )
            .map(([account]) => account.slice(5)),
          actions: state.actions,
          paths: Object.fromEntries(
            [...state.paths].map(([account, path]) => [account.slice(5), path]),
          ),
        },
      });
    for (const [name, edge] of edges) {
      const action = instrument.actions[name];
      if (!action) continue;
      const next = copy(state);
      if (edge.to !== "preserve") next.status = edge.to;
      apply(next, action, name);
      next.actions.push(name);
      for (const [account, balance] of next.balances) {
        if (balance !== null && !balance.size) next.paths.delete(account);
        else
          next.paths.set(account, [
            ...(state.paths.get(account) ?? state.actions),
            name,
          ]);
      }
      pending.push(next);
    }
  }
  return problems;
}