import { lsType } from "../structures/list_common.js";
import { fnType } from "../metacircular.js";
import { deepEqual, DefaultMap, pretty } from "../util.js";
import { numDictType } from "../typeclasses/num_type.js";

const genericTypeRegistry = new DefaultMap(underlyingType => ({ generic: underlyingType }));

// type constructor for generic kinds,
//   for instance:
//     a -> a -> Bool
//   is typed by
//     genericType(Function)
export const genericType = underlyingType => genericTypeRegistry.getdefault(underlyingType, true);

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

// From the given set of type variables, return only those that occur in the given type.
export const occurring = (type, typeVars) => {
  if (typeVars.has(type)) {
    return new Set([type]);
  }
  if (type.in !== undefined) {
    // function type
    return new Set([
      ...occurring(type.in, typeVars),
      ...occurring(type.out, typeVars)]);
  }
  if (type.listOf !== undefined) {
    return occurring(type.listOf, typeVars);
  }
  return new Set();
}

export const properUnify = eqDict => (
  {typeVars: formalTypeVars, type: formalType},
  {typeVars: actualTypeVars, type: actualType},
) => {
  if (getEq(eqDict)(formalType)(actualType)) {
    return {
      substitutions: new Map(),
      typeVars: new Set([
        ...actualTypeVars,
        // ...formalTypeVars, // <- i don't think we need these?
      ]),
      type: actualType,
    }
  }

  if (formalTypeVars.has(formalType)) {
    // formalType is type variable -> substitute it by actualType
    return {
      substitutions: new Map([[formalType, actualType]]),
      typeVars: new Set([
        ...actualTypeVars,
        ...formalTypeVars,
      ].filter(a => a !== formalType)),
      type: actualType,
    }
  }
  if (actualTypeVars.has(actualType)) {
    // same as above, but in opposite direction:
    // actualType is type variable -> substitute it by formalType
    return {
      substitutions: new Map([[actualType, formalType]]),
      typeVars: new Set([
        ...actualTypeVars,
        ...formalTypeVars,
      ].filter(a => a !== actualType)),
      type: formalType,
    }
  }
  
  // WIP...
}

const mergeOneWay = (m1, m2) => {
  const m1copy = new Map(m1);
  const m2copy = new Map(m2);
  for (const [key1, val1] of m1copy.entries()) {
    if (m2copy.has(val1)) {
      m1copy.set(key1, m2.get(val1));
      m2copy.delete(val1);
      return [false, m1copy, m2copy, new Set([val1])];
    }
  }
  return [true, m1copy, m2copy, new Set()]; // stable
}

export const mergeSubstitutions = (m1, m2) => {
  let stable = false;
  let deletedTypeVars = 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);
    deletedTypeVars = deletedTypeVars.union(d);
  }
  return [new Map([...m1, ...m2]), deletedTypeVars];
}

// Currently very ad-hoc.

// 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
export const unify = (
  {typeVars: formalTypeVars, type: formalType},
  {typeVars: actualTypeVars, type: actualType},
) => {
  if (deepEqual(formalType, actualType)) {
    return {
      substitutions: new Map(),
      typeVars: new Set([
        ...actualTypeVars,
        ...formalTypeVars]),
      type: actualType,
    }
  }

  if (formalTypeVars.has(formalType)) {
    // simplest case: substitute formal type param
    return {
      substitutions: new Map([[formalType, actualType]]),
      typeVars: new Set([
        ...actualTypeVars,
        ...formalTypeVars].filter(a => a !== formalType)),
      type: actualType,
    };
  }

  if (actualTypeVars.has(actualType)) {
    // same as above, but in the other direction
    return {
      substitutions: new Map([[actualType, formalType]]),
      typeVars: new Set([
        ...actualTypeVars,
        ...formalTypeVars].filter(a => a !== actualType)),
      type: formalType,
    };
  }

  if (formalType.in !== undefined) {
    // function type
    if (actualType.in === undefined) {
      throw new Error(`cannot assign ${pretty(actualType)} to ${pretty(formalType)}`);
    }
    else {
      // both are function type
      const inType = unify({typeVars: formalTypeVars, type: formalType.in}, {typeVars: actualTypeVars, type: actualType.in});
      const outType = unify({typeVars: formalTypeVars, type: formalType.out}, {typeVars: actualTypeVars, type: actualType.out});
      // check for conflicts between 'in' and 'out' subsitutions
      for (const [typeVar, actual] of inType.substitutions) {
        if (outType.substitutions.has(typeVar)) {
          if (!deepEqual(actual, outType.substitutions.get(typeVar))) {
            throw new Error(`conflicting assignment for ${pretty(typeVar)}: ${pretty(a)}`);
          }
        }
      }
      // merge substitutions
      const [newSubstitutions, deletedTypeVars] = mergeSubstitutions(
        inType.substitutions, outType.substitutions);
      // const newSubstitutions = new Map([
      //   ...inType.substitutions,
      //   ...outType.substitutions,
      // ]);
      const newTypeVars = new Set([
        ...actualTypeVars,
        ...formalTypeVars].filter(a => !newSubstitutions.has(a) && !deletedTypeVars.has(a)));
      return {
        substitutions: newSubstitutions,
        typeVars: newTypeVars,
        type: fnType({in: inType.type, out: outType.type}),
      };
    }
  }

  if (formalType.listOf !== undefined) {
    // list type
    if (actualType.listOf === undefined) {
      throw new Error(`cannot assign ${pretty(actualType)} to ${pretty(formalType)}`);
    }
    else {
      // both are list type
      const elementType = unify(
        {typeVars: formalTypeVars, type: formalType.listOf},
        {typeVars: actualTypeVars, type: actualType.listOf});
      return {
        substitutions: elementType.substitutions,
        typeVars: new Set([
          ...actualTypeVars,
          ...formalTypeVars].filter(a => !elementType.substitutions.has(a))),
        type: lsType(elementType.type),
      };
    }
  }

  if (formalType.numDict !== undefined) {
    if (actualType.numDict === undefined) {
      throw new Error(`cannot assign ${pretty(actualType)} to ${pretty(formalType)}`);
    }
    else {
      // both are NumDict type
      const underlyingType = unify(
        {typeVars: formalTypeVars, type: formalType.numDict},
        {typeVars: actualTypeVars, type: actualType.numDict});
      return {
        substitutions: underlyingType.substitutions,
        typeVars: new Set([
          ...actualTypeVars,
          ...formalTypeVars].filter(a => !underlyingType.substitutions.has(a))),
        type: numDictType(underlyingType.type),
      };
    }
  }
  throw new Error("i don't know what to do :(");
};

// export const matchConcrete = ({typeVars, type: formalType}, actualType) => {
//   return unify({typeVars, type: formalType}, {typeVars: new Set(), type: actualType});
// };

export const substitute = (type, substitutions) => {
  if (substitutions.has(type)) {
    return substitutions.get(type);
  }
  if (type.listOf !== undefined) {
    // list type
    return lsType(substitute(type.listOf, substitutions));
  }
  if (type.in !== undefined) {
    // function type
    return fnType({
      in: substitute(type.in, substitutions),
      out: substitute(type.out, substitutions),
    })
  }
  return type;
}

export const assign = (genFnType, paramType) => {
  const matchedInType = unify({
    typeVars: genFnType.typeVars,
    type: genFnType.type.in,
  }, paramType);
  const substitutedOutType = substitute(genFnType.type.out, matchedInType.substitutions);
  return {
    typeVars: matchedInType.typeVars,
    type: substitutedOutType,
  };
};