import { eqType } from "../primitives/type.js";
import { zip } from "../util/util.js";
import { pretty } from '../util/pretty.js';
import { 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,
});

// From the given set of type variables, return only those that occur in the given type.
export const occurring = (type, typeVars) => {
  // console.log("occurring", type);
  if (typeVars.has(type)) {
    // type IS a type variable:
    return new Set([type]);
  }
  return new Set(type.params.flatMap(p => [...occurring(p, typeVars)]));
};

// 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 deleted = 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);
    deleted = deleted.union(d);
  }
  const result = {
    substitutions: new Map([...m1, ...m2]),
    deleted, // 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
const unifyInternal = (typeVars, fType, aType) => {
  // console.log("unify",  pretty({typeVars, fType, aType}));
  
  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)}`);
  }
  else {
    // console.log("symbols match - unify recursively", formal.symbol);
    const unifiedParams =
      zip(fType.params, aType.params)
      .map(([fParam, aParam]) => unifyInternal(typeVars, fParam, aParam));
    const {substitutions, deleted} =
      unifiedParams.reduce(({substitutions: s, deleted: d}, cur) => {
        // console.log('merging', s, cur.substitutions);
        const {substitutions, deleted} = mergeTwoWay(s, cur.substitutions);
        return {
          substitutions,
          deleted: deleted.union(d),
        };
      }, { substitutions: new Map(), deleted: new Set() });
    // console.log(pretty({unifiedParams}));
    return {
      substitutions,
      genericType: {
        typeVars: typeVars.difference(substitutions).difference(deleted),
        type: {
          symbol: fType.symbol,
          params: unifiedParams.map(p => p.genericType.type),
        },
      },
    };
  }
};

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

export const substitute = (type, substitutions) => {
  // console.log("substitute", {type, substitutions})
  if (substitutions.has(type)) {
    return substitutions.get(type);
  }
  if (typeof type === "symbol") {
    return type; // nothing to substitute here
  }
  return {
    symbol: type.symbol,
    params: type.params.map(p => substitute(p, substitutions)),
  };
};

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);
  // console.log({genFnType: prettyT(genFnType), paramType: prettyT(paramType), substitutions})
  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];
}