import { compareTypes } from "../compare/type.js";
import { getHumanReadableName } from "../primitives/symbol.js";
import { eqType, getSymbol } from "../primitives/type.js";
import { isTypeVar } from "../primitives/typevars.js";
import { emptySet, add, has } from "../structures/set.js";
import { prettyT } from "../util/pretty.js";
import { zip } from "../util/util.js";
import { UnifyError } from "./generics.js";

const emptyTypeSet = emptySet(compareTypes);

// Low-level unify
// Assumes that if types variables in typeA and typeB are overlapping, they are the same, so it may be necessary to re-compute type variables before calling this function.
export const unifyLL = (typeA, typeB, stackA = [], stackB = []) => {
  if (eqType(typeA)(typeB)) {
    return new Map();
  }

  if (isTypeVar(typeA) || isTypeVar(typeB)) {
    const unifA = isTypeVar(typeA)
      ? new Map([[getSymbol(typeA), add(emptyTypeSet)(typeB)]])
      : new Map();

    const unifB = isTypeVar(typeB)
      ? new Map([[getSymbol(typeB), add(emptyTypeSet)(typeA)]])
      : new Map();

    return mergeUnifications(unifA, unifB);
  }

  // recursively unify
  if (typeA.symbol !== typeB.symbol) {
    throw new UnifyError(`cannot unify ${prettyT(typeA)} and ${prettyT(typeB)}`);
  }

  const tagA = stackA.length;
  const tagB = stackB.length;

  const unifParams = zip(typeA.params, typeB.params)
    .map(([getParamA, getParamB]) => {
      const paramA = getParamA(tagA);
      const paramB = getParamB(tagB);

      // type recursively points to an enclosing type that we've already seen
      if (stackA[paramA] !== stackB[paramB]) {
        // note that both are also allowed not to be mapped (undefined)
        throw new UnifyError("cannot unify: types differ in their recursion");
      }
      if (stackA[paramA] !== undefined) {
        // we've already seen this type, don't endlessly recurse:
        return new Map();
      }
      return unifyLL(paramA, paramB,
        [...stackA, tagA],
        [...stackB, tagB]);
    });

  return unifParams.reduce(
    (acc, cur) => mergeUnifications(acc, cur),
    new Map());
};

// Given two unifications, try to merge them (may throw UnifyError).
// Useful when the same type variable occurs in multiple places, to see if there are conflicts.
export const mergeUnifications = (unifA, unifB) => {
  const allSymbols = new Set([...unifA.keys(), ...unifB.keys()]);
  const result = new Map();
  for (const symbol of allSymbols) {
    const setOfTypesA = unifA.get(symbol) || emptyTypeSet;
    const setOfTypesB = unifB.get(symbol) || emptyTypeSet;
    let union = setOfTypesA;
    for (const typeB of setOfTypesB.keys()) {
      union = addIfSafe(union, typeB);
    }
    result.set(symbol, union);
  }
  // console.log(`mergeUnifications(${prettyU(unifA)}, ${prettyU(unifB)}) = ${prettyU(result)}`);
  // return transitivelyMerge(result);
  return result;
};

const transitivelyGrow = (unif) => {
  let stable = true;
  const result = new Map();
  for (const [symbol, setOfTypes] of unif) {
    let newSetOfTypes = setOfTypes;
    for (const type of setOfTypes.keys()) {
      if (isTypeVar(type)) {
        const haveTypes = unif.get(getSymbol(type));
        if (haveTypes) {
          for (const transitiveType of haveTypes.keys()) {
            if (!has(newSetOfTypes)(transitiveType)) {
              newSetOfTypes = addIfSafe(newSetOfTypes, transitiveType);
              stable = false;  
            }
          }
        }
      }
    }
    result.set(symbol, newSetOfTypes);
  }
  // repeat until stable
  return stable ? result : transitivelyGrow(result);
};

const addIfSafe = (setOfTypes, typeToAdd) => {
  for (const alreadyHaveType of setOfTypes.keys()) {
    if (!has(setOfTypes)(typeToAdd)) {
      if (isTypeVar(alreadyHaveType) && isTypeVar(typeToAdd)) {
        continue; // not a problem
      }
      // console.log('can unify', prettyT(typeToAdd), 'and', prettyT(alreadyHaveType), '?');
      unifyLL(alreadyHaveType, typeToAdd); // may throw
    }
  }
  return add(setOfTypes)(typeToAdd);
};

// Given a non-conflicting, non-empty set of types, reduce it to a single type
export const reduce = (setOfTypes) => {
  for (const type of setOfTypes.keys()) {
    if (!isTypeVar(type)) {
      // console.log('reduce', prettyST(setOfTypes), 'to', prettyT(type));
      return type;
    }
  }
  // console.log('reduce', prettyST(setOfTypes), 'to', prettyT(setOfTypes.keys()[0]));
  return setOfTypes.keys()[0];
};

// Reduce a unification to a mapping: {symbol => Type}
// this mapping can then be used for substituting the typevars (=symbols) in a type by concrete types
export const reduceUnification = (unif) => {
  // console.log('b4 grown:', prettyU(unif));
  const grown = transitivelyGrow(unif);
  // console.log('grown:', prettyU(grown));
  const result = new Map([...grown]
    .map(([symbol, types]) =>
      [symbol, reduce(types)]));
  // console.log('reduce', prettyU(grown), 'to', result);
  return result;
};

// For debugging
const prettyU = unif => {
  return `{${[...unif].map(([symbol, types]) => `${getHumanReadableName(symbol)} => ${prettyST(types)}`).join(', ')}}`;
};
const prettyST = st => {
  return `(${st.keys().map(prettyT).join(',')})`;
};