import { eqType } from "../primitives/type.js";
import { zip } from "../util/util.js";
import { pretty, prettyT } from '../util/pretty.js';

// constructor for generic types
//   for instance, the type:
//     ∀a: a -> a -> Bool
//   is created by
//     makeGeneric(a => fnType(() => a)(() => fnType(() => a)(() => Bool)))
export const makeGeneric = callback => {
  // type variables to make available:
  const typeVars = ['a', 'b', 'c', 'd', 'e'].map(Symbol);
  const type = callback(...typeVars);
  return onlyOccurring(type, new Set(typeVars));
};

export const onlyOccurring = (type, typeVars) => ({
  typeVars: occurring(type, typeVars),
  type,
});

const __occurring = state => typeVars => type => {
  if (typeVars.has(type)) {
    return new Set([type]);
  }
  const tag = state.nextTag++;
  state.seen.add(tag);
  return new Set(type.params.flatMap(p => {
    const innerType = p(tag);
    if (state.seen.has(innerType)) {
      return []; // no endless recursion!
    }
    return [...__occurring(state)(typeVars)(innerType)];
  }));
};

// From the given set of type variables, return only those that occur in the given type.
export const occurring = (type, typeVars) => {
  return __occurring({nextTag:0, seen: new Set()})(typeVars)(type);
};

// Merge 2 substitution-mappings, uni-directional.
const mergeOneWay = (m1, m2) => {
  const m1copy = new Map(m1);
  const m2copy = new Map(m2);
  for (const [var1, typ1] of m1copy.entries()) {
    if (m2copy.has(typ1)) {
      // typ1 is a typeVar for which we also have a substitution
      // -> fold substitutions
      m1copy.set(var1, m2.get(typ1));
      m2copy.delete(typ1);
      return [false, m1copy, m2copy, new Set([typ1])];
    }
  }
  return [true, m1copy, m2copy, new Set()]; // stable
};

const checkConflict = (m1, m2) => {
  for (const [var1, typ1] of m1) {
    if (m2.has(var1)) {
      const other = m2.get(var1);
      if (!eqType(typ1, other)) {
        throw new Error(`conflicting substitution: ${pretty(typ1)}vs. ${pretty(other)}`);
      }
    }
  }
};

// Merge 2 substitution-mappings, bi-directional.
export const mergeTwoWay = (m1, m2) => {
  // console.log("mergeTwoWay", {m1, m2});
  checkConflict(m1, m2);
  // checkConflict(m2, m1); // <- don't think this is necessary...
  // actually merge
  let stable = false;
  let deletions = new Set();
  while (!stable) {
    let d;
    // notice we swap m2 and m1, so the rewriting can happen both ways:
    [stable, m2, m1, d] = mergeOneWay(m1, m2);
    deletions = deletions.union(d);
  }
  const result = {
    substitutions: new Map([...m1, ...m2]),
    deletions, // deleted type variables
  };
  // console.log("mergeTwoWay result =", result);
  return result;
};

// Thanks to Hans for pointing out that this algorithm exactly like "Unification" in Prolog (hence the function name):
//  https://www.dai.ed.ac.uk/groups/ssp/bookpages/quickprolog/node12.html
//
// Parameters:
//   typeVars: all the type variables in both fType and aType
//   fType, aType: generic types to unify
//   fStack, aStack: internal use.
const __unify = (typeVars, fType, aType, fStack=[], aStack=[]) => {
  console.log("__unify", {typeVars, fType, aType, fStack, aStack});
  if (typeVars.has(fType)) {
    // simplest case: formalType is a type paramater
    //   => substitute with actualType
    // console.log("assign actual to formal");
    return {
      substitutions: new Map([[fType, aType]]),
      genericType: {
        typeVars: typeVars.difference(new Set([fType])),
        type: aType,
      },
    };
  }
  if (typeVars.has(aType)) {
    // same as above, but in the other direction
    // console.log("assign formal to actual");
    return {
      substitutions: new Map([[aType, fType]]),
      genericType: {
        typeVars: typeVars.difference(new Set([aType])),
        type: fType,
      },
    };
  }

  // recursively unify
  if (fType.symbol !== aType.symbol) {
    throw new Error(`cannot unify ${prettyT(fType)} and ${prettyT(aType)}`);
  }

  const fTag = fStack.length;
  const aTag = aStack.length;

  const unifications =
    zip(fType.params, aType.params)
    .map(([getFParam, getAParam]) => {
      const fParam = getFParam(fTag);
      const aParam = getAParam(aTag);
      // type recursively points to an enclosing type that we've already seen
      if (fStack[fParam] !== aStack[aParam]) {
        // note that both are also allowed not to be mapped (undefined)
        throw new Error("cannot unify: types differ in their recursion");
      }
      if (fStack[fParam] !== undefined) {
        const type = fStack[fParam];
        return () => ({
          substitutions: new Map(),
          genericType: {
            typeVars,
            type,
          },
        });
      }
      return parent => __unify(typeVars, fParam, aParam,
        [...fStack, parent],
        [...aStack, parent]);
    });

  const unifiedParams = unifications.map(getParam => {
    return parent => getParam(parent).genericType.type;
  });
  const type = {
    symbol: fType.symbol,
    params: unifiedParams,
  };

  const [unifiedSubstitutions, unifiedTypeVars] = unifications.reduce((acc, getParam) => {
    const self = Symbol();
    const {substitutions, deletions} = mergeTwoWay(acc[0], getParam(self).substitutions);
    return [substitutions, acc[1]
      .difference(substitutions)
      .difference(deletions)];
  }, [new Map(), typeVars]);

  return {
    substitutions: unifiedSubstitutions,
    genericType: {
      typeVars: unifiedTypeVars,
      type,
    },
  };
};

export const unify = (fGenericType, aGenericType) => {
  let allTypeVars;
  [allTypeVars, fGenericType, aGenericType] = safeUnionTypeVars(fGenericType, aGenericType);
  const {genericType} = __unify(allTypeVars, fGenericType.type, aGenericType.type);
  return recomputeTypeVars(genericType);
};

export const substitute = (type, substitutions, stack=[]) => {
  console.log('substitute...', {type, substitutions, stack});
  return substitutions.get(type)
  || {
    symbol: type.symbol,
    params: type.params.map(getParam => parent => {
      const param = getParam(stack.length);
      const have = stack[param];
      return (have !== undefined)
        ? have
        : substitute(param, substitutions, [...stack, parent]);
    }),
  };
};

export const assign = (genFnType, paramType) => {
  let allTypeVars;
  [allTypeVars, genFnType, paramType] = safeUnionTypeVars(genFnType, paramType);
  const [inType, outType] = genFnType.type.params;
  const {substitutions} = unifyInternal(allTypeVars, inType, paramType.type);
  const substitutedOutType = substitute(outType, substitutions);
  return recomputeTypeVars(onlyOccurring(substitutedOutType, allTypeVars));
};

export const assignFn = (genFnType, paramType) => {
  let allTypeVars;
  [allTypeVars, genFnType, paramType] = safeUnionTypeVars(genFnType, paramType);
  const [inType] = genFnType.type.params;
  const {substitutions} = unifyInternal(allTypeVars, inType, paramType.type);
  const substitutedFnType = substitute(genFnType.type, substitutions);
  return recomputeTypeVars(onlyOccurring(substitutedFnType, allTypeVars));
};

export const recomputeTypeVars = (genType) => {
  const newTypeVars = ['a', 'b', 'c', 'd', 'e', 'f', 'g'].map(Symbol);
  let nextIdx = 0;
  const subst = new Map();
  for (const typeVarA of genType.typeVars) {
    subst.set(typeVarA, newTypeVars[nextIdx++]);
  }
  const substType = {
    typeVars: new Set(subst.values()),
    type: substitute(genType.type, subst),
  };
  return substType;
};

export const safeUnionTypeVars = (genTypeA, genTypeB) => {
  const substTypeA = recomputeTypeVars(genTypeA);
  const substTypeB = recomputeTypeVars(genTypeB);
  const allTypeVars = substTypeA.typeVars.union(substTypeB.typeVars);
  return [allTypeVars, substTypeA, substTypeB];
};