import { inspect } from "node:util";
import { eqType, getSymbol } from "../primitives/type.js";
import { zip } from "../util/util.js";
import { pretty, prettyT } from '../util/pretty.js';
import { isTypeVar, TYPE_VARS } from "../primitives/typevars.js";
import { inspectType } from "../meta/type_constructor.js";
import { symbolFunction } from "../structures/type_constructors.js";

// helper for creating generic types
//   for instance, the type:
//     a -> a -> Bool
//   is created by
//     makeGeneric(a => fnType(() => a)(() => fnType(() => a)(() => Bool)))
export const makeGeneric = callback => {
  // type variables to make available:
  const type = callback(...TYPE_VARS);
  return type;
};

const _occurring = stack => type => {
  if (isTypeVar(type)) {
    return new Set([getSymbol(type)]);
  }
  const tag = stack.length;
  const newStack = [...stack, tag];
  return new Set(type.params.flatMap(p => {
    const innerType = p(tag);
    if (newStack.includes(innerType)) {
      return []; // no endless recursion!
    }
    return [..._occurring(newStack)(innerType)]; 
  }));
};

// Get set of type variables in type.
export const occurring = _occurring([]);

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

const checkConflict = (m1, m2) => {
  for (const [symbol1, typ1] of m1) {
    if (m2.has(symbol1)) {
      const other = m2.get(symbol1);
      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 remaining = 2;
  while (remaining > 0) {
    // notice we swap m2 and m1, so the rewriting can happen both ways:
    let stable;
    [stable, m2, m1] = mergeOneWay(m1, m2);
    remaining -= stable;
  }
  const result = new Map([...m1, ...m2]);
  // console.log("mergeTwoWay result =", result);
  return result;
};

export class UnifyError extends Error {}
export class NotAFunctionError extends Error {}

// 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.
export const __unify = (fType, aType, fStack=[], aStack=[]) => {
  // console.log("__unify", {typeVars, fType: prettyT(fType), aType: prettyT(aType), fStack, aStack});
  if (isTypeVar(fType)) {
    // simplest case: formalType is a type paramater
    //   => substitute with actualType
    // console.log(`assign ${prettyT(aType)} to ${prettyT(fType)}`);
    return {
      substitutions: new Map([[getSymbol(fType), aType]]),
      type: aType,
    };
  }
  if (isTypeVar(aType)) {
    // same as above, but in the other direction
    // console.log(`assign ${prettyT(fType)} to ${prettyT(aType)}`);
    return {
      substitutions: new Map([[getSymbol(aType), fType]]),
      type: fType,
    };
  }

  // recursively unify
  if (fType.symbol !== aType.symbol) {
    throw new UnifyError(`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 UnifyError("cannot unify: types differ in their recursion");
      }
      if (fStack[fParam] !== undefined) {
        const type = fStack[fParam];
        return () => ({
          substitutions: new Map(),
          type,
        });
      }
      return parent => __unify(fParam, aParam,
        [...fStack, parent],
        [...aStack, parent]);
    });

  const unifiedParams = unifications.map(getParam => {
    return parent => getParam(parent).type;
  });

  const unifiedSubstitutions = unifications.reduce((acc, getParam) => {
    const self = Symbol(); // dirty, just need something unique
    const paramSubstitutions = getParam(self).substitutions;
    const substitutions = mergeTwoWay(acc, paramSubstitutions);
    return substitutions;
  }, new Map());

  return {
    substitutions: unifiedSubstitutions,
    type: {
      symbol: fType.symbol,
      params: unifiedParams,
      [inspect.custom]: inspectType,
    },
  };
};

export const unify = (fType, aType) => {
  [fType, aType] = recomputeTypeVars([fType, aType]);
  const {type, substitutions} = __unify(fType, aType);
  // console.log('unification complete! substitutions:', substitutions);
  return recomputeTypeVars([type])[0];
};

export const substitute = (type, substitutions, stack=[]) => {
  // console.log('substitute...', {type, substitutions, stack});
  const found = substitutions.get(getSymbol(type));
  if (found) {
    return found;
  }
  return {
    symbol: getSymbol(type),
    params: type.params.map(getParam => parent => {
      const param = getParam(parent);
      if (stack.includes(param)) {
        // param points back up - that's ok - means we don't have to recurse
        return param;
      }
      return substitute(param, substitutions, [...stack, parent]);
    }),
    [inspect.custom]: inspectType,
  };
};

export const assignFn = (funType, paramType) => {
  const [inType, inSubst, outType, outSubst] = assignFnSubstitutions(funType, paramType);
  // return recomputeTypeVars([outType])[0];
  return outType;
};

// same as above, but also returns the substitutions that took place
export const assignFnSubstitutions = (funType, paramType, skip=0) => {
  if (getSymbol(funType) !== symbolFunction) {
    throw new NotAFunctionError(`${prettyT(funType)} is not a function type!`);
  }
  const [[refunType, funS], [reparamType, paramS]] = recomputeTypeVarsSubstitutions([funType, paramType], skip);
  const [inType, outType] = refunType.params.map(p => p(refunType));
  const {type: newInType, substitutions} = __unify(inType, reparamType);
  const totalParamSubstitutions = mergeTwoWay(substitutions, paramS);
  const newOutType = substitute(outType, substitutions);
  const [[finalOutType, outsubst]] = recomputeTypeVarsSubstitutions([newOutType], skip);
  const totalOutSubstitutions = mergeTwoWay(funS, outsubst);
  return [newInType, totalParamSubstitutions, finalOutType, totalOutSubstitutions];
};

// Ensures that no type variables overlap
export const recomputeTypeVars = types => {
  return recomputeTypeVarsSubstitutions(types)
    .map(([newType, _subst]) => newType);
};

export const recomputeTypeVarsSubstitutions = (types, skip=0) => {
  let nextIdx = skip;
  return types.map(type => {
    const substitutions = new Map();
    const typeVars = occurring(type);
    for (const typeVar of typeVars) {
      substitutions.set(typeVar, TYPE_VARS[nextIdx++]);
    }
    return [substitute(type, substitutions), substitutions];
  });
};