import { Bool, Int } from "./primitives/symbols.js";
import { fnType, lsType } from "./type_registry.js";
import { deepEqual } from "./util.js";

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,
  };
};

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();
}


import { inspect } from 'node:util';

function pretty(obj) {
  return inspect(obj, {colors: true});
}

export const matchGeneric = (
  {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 (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 = matchGeneric({typeVars: formalTypeVars, type: formalType.in}, {typeVars: actualTypeVars, type: actualType.in});
      const outType = matchGeneric({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 = new Map([
        ...inType.substitutions,
        ...outType.substitutions,
      ]);
      const newTypeVars = new Set([
        ...actualTypeVars,
        ...formalTypeVars].filter(a => !newSubstitutions.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 = matchGeneric(
        {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),
      };
    }
  }
  throw new Error("i don't know what to do :(")
};

export const matchConcrete = ({typeVars, type: formalType}, actualType) => {
  return matchGeneric({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 = matchGeneric({
    typeVars: genFnType.typeVars,
    type: genFnType.type.in,
  }, paramType);
  const substitutedOutType = substitute(genFnType.type.out, matchedInType.substitutions);
  return {
    typeVars: matchedInType.typeVars,
    type: substitutedOutType,
  };
}

// a -> Int
const a_to_Int = makeGeneric(a => fnType({in: a, out: Int}));
const Bool_to_Int = fnType({in: lsType(Bool), out: Int});
console.log(matchConcrete(a_to_Int, Bool_to_Int));

// (a -> a) -> b
const fnType2 = makeGeneric((a,b) => fnType({in: fnType({in: a, out: a}), out: b}));
// (Bool -> Bool) -> a
const fnType3 = makeGeneric(a => fnType({in: fnType({in: Bool, out: Bool}), out: a}));
console.log(matchGeneric(fnType2, fnType3));

// (a -> b) -> [a] -> [b]
const mapFnType = makeGeneric((a,b) =>
  fnType({
    in: fnType({in: a, out: b}),
    out: fnType({in: lsType(a), out: lsType(b)}),
  }));
// a -> a
const idFnType = makeGeneric(a =>
  fnType({in: a, out: a}));

// should be: [a] -> [a]
console.log(assign(mapFnType, idFnType));