// Semantic types produced by the checker.

export type Type =
  | { k: "prim"; name: PrimName }
  | { k: "optional"; inner: Type }
  | { k: "result"; inner: Type }
  | { k: "list"; elem: Type }
  | { k: "map"; key: Type; value: Type }
  | { k: "set"; elem: Type }
  | { k: "tuple"; items: Type[] }
  | { k: "range" }
  | { k: "fun"; params: Type[]; ret: Type; async: boolean }
  | { k: "struct"; name: string; fields: FieldInfo[]; kind: "struct" | "schema" | "class" | "contract"; typeArgs: Type[] }
  | { k: "enum"; name: string; variants: string[] }
  | { k: "enumval"; of: string; name: string }
  | { k: "union"; name: string; variants: VariantInfo[] }
  | { k: "variant"; of: string; name: string; fields: FieldInfo[] }
  | { k: "typeparam"; name: string; id: number }
  | { k: "bytes" }
  | { k: "any" }
  | { k: "void" }
  | { k: "err" }
  | { k: "expect" }
  | { k: "module"; moduleId: string }
  | { k: "unknown" };

export type PrimName = "Int" | "Float" | "Bool" | "String" | "Char";

export interface FieldInfo {
  name: string;
  type: Type;
}

export interface VariantInfo {
  name: string;
  fields: FieldInfo[];
  span?: { file: string; start: number; end: number };
}

export const INT: Type = { k: "prim", name: "Int" };
export const FLOAT: Type = { k: "prim", name: "Float" };
export const BOOL: Type = { k: "prim", name: "Bool" };
export const STRING: Type = { k: "prim", name: "String" };
export const CHAR: Type = { k: "prim", name: "Char" };
export const VOID: Type = { k: "void" };
export const ANY: Type = { k: "any" };
export const ERR: Type = { k: "err" };
export const RANGE: Type = { k: "range" };
export const BYTES: Type = { k: "bytes" };
export const UNKNOWN: Type = { k: "unknown" };

export const prim = (name: PrimName): Type => ({ k: "prim", name });
export const optional = (inner: Type): Type => ({ k: "optional", inner });
export const result = (inner: Type): Type => ({ k: "result", inner });
export const list = (elem: Type): Type => ({ k: "list", elem });
export const mapType = (key: Type, value: Type): Type => ({ k: "map", key, value });
export const set = (elem: Type): Type => ({ k: "set", elem });
export const tuple = (items: Type[]): Type => ({ k: "tuple", items });
export const fun = (params: Type[], ret: Type, async = false): Type => ({ k: "fun", params, ret, async });
export const typeparam = (name: string, id: number): Type => ({ k: "typeparam", name, id });

export function typeToString(t: Type, seen = new Set<Type>()): string {
  if (seen.has(t)) return "...";
  seen.add(t);
  switch (t.k) {
    case "prim":
      return t.name;
    case "optional":
      return typeToString(t.inner, seen) + "?";
    case "result":
      return typeToString(t.inner, seen) + "!";
    case "list":
      return `List<${typeToString(t.elem, seen)}>`;
    case "map":
      return `Map<${typeToString(t.key, seen)}, ${typeToString(t.value, seen)}>`;
    case "set":
      return `Set<${typeToString(t.elem, seen)}>`;
    case "tuple":
      return `(${t.items.map((i) => typeToString(i, seen)).join(", ")})`;
    case "range":
      return "Range";
    case "fun": {
      const ps = t.params.map((p) => typeToString(p, seen)).join(", ");
      const r = typeToString(t.ret, seen);
      return t.async ? `async (${ps}) -> ${r}` : `(${ps}) -> ${r}`;
    }
    case "struct":
      return t.typeArgs.length ? `${t.name}<${t.typeArgs.map((a) => typeToString(a, seen)).join(", ")}>` : t.name;
    case "enum":
      return t.name;
    case "enumval":
      return `${t.of}.${t.name}`;
    case "union":
      return t.variants.map((v) => v.fields.length ? `${v.name}(${v.fields.map((f) => f.name).join(", ")})` : v.name).join(" | ");
    case "variant":
      return `${t.of}.${t.name}`;
    case "typeparam":
      return t.name;
    case "bytes":
      return "Bytes";
    case "any":
      return "Any";
    case "void":
      return "Void";
    case "err":
      return "Err";
    case "expect":
      return "Expect";
    case "module":
      return t.moduleId;
    case "unknown":
      return "?";
  }
}

// Structural assignability. `a` is assignable to `b`.
export function assignable(a: Type, b: Type, typeParams: Map<number, Type> = new Map()): boolean {
  if (a.k === "unknown" || b.k === "unknown") return true;
  if (a.k === "any" || b.k === "any") return true;
  if (a.k === "typeparam") {
    const inst = typeParams.get(a.id);
    if (inst) return assignable(inst, b, typeParams);
    return true;
  }
  if (b.k === "typeparam") {
    const inst = typeParams.get(b.id);
    if (inst) return assignable(a, inst, typeParams);
    return true;
  }
  if (a.k === b.k) {
    switch (a.k) {
      case "prim":
        return a.name === (b as { k: "prim"; name: PrimName }).name || (a.name === "Int" && (b as { k: "prim"; name: PrimName }).name === "Float");
      case "optional":
        return assignable(a.inner, (b as { k: "optional"; inner: Type }).inner, typeParams);
      case "result":
        return assignable(a.inner, (b as { k: "result"; inner: Type }).inner, typeParams);
      case "list":
        return assignable(a.elem, (b as { k: "list"; elem: Type }).elem, typeParams);
      case "set":
        return assignable(a.elem, (b as { k: "set"; elem: Type }).elem, typeParams);
      case "map":
        return assignable(a.key, (b as { k: "map"; key: Type; value: Type }).key, typeParams) &&
          assignable(a.value, (b as { k: "map"; key: Type; value: Type }).value, typeParams);
      case "tuple":
        return a.items.length === (b as { k: "tuple"; items: Type[] }).items.length &&
          a.items.every((x, i) => assignable(x, (b as { k: "tuple"; items: Type[] }).items[i]!, typeParams));
      case "fun":
        return fnAssignable(a, b as { k: "fun"; params: Type[]; ret: Type; async: boolean }, typeParams);
      case "struct": {
        const bs = b as { k: "struct"; name: string; fields: FieldInfo[]; kind: string };
        if (a.kind === "contract" || bs.kind === "contract") {
          // contract target: every method of the contract must be present in a
          const target = bs.kind === "contract" ? bs : a;
          const source = bs.kind === "contract" ? a : bs;
          for (const m of target.fields) {
            const f = source.fields.find((x) => x.name === m.name);
            if (!f) return false;
            if (!fnAssignable(f.type as Type, m.type as Type, typeParams)) return false;
          }
          return true;
        }
        return bs.fields.every((f) => {
          const af = a.fields.find((x) => x.name === f.name);
          return af !== undefined && assignable(af.type, f.type, typeParams);
        });
      }
      case "union": {
        const bu = b as { k: "union"; name: string; variants: VariantInfo[] };
        if (a.name !== bu.name) {
          // a's variants must all be covered by b's variants
          return a.variants.every((v) => bu.variants.some((bv) => bv.name === v.name));
        }
        return true;
      }
      case "enum":
        return a.name === (b as { k: "enum"; name: string }).name;
      case "enumval":
        return a.of === (b as { k: "enumval"; of: string }).of && a.name === (b as { k: "enumval"; name: string }).name;
      case "variant": {
        const bv = b as { k: "variant"; of: string; name: string };
        return a.name === bv.name && (a.of === bv.of || true);
      }
      case "bytes":
      case "err":
      case "expect":
      case "range":
      case "void":
        return true;
      default:
        return true;
    }
  }
  // different kinds
  if (a.k === "variant" && b.k === "union") {
    return assignableVariantToUnion(a, b);
  }
  if (a.k === "enumval" && b.k === "enum") {
    return a.of === b.name;
  }
  if (a.k === "struct" && b.k === "tuple") {
    // positional compatibility: struct → tuple of its field types
    return a.fields.length === b.items.length && a.fields.every((f, i) => assignable(f.type, b.items[i]!, typeParams));
  }
  if (a.k === "optional" && b.k === "prim") {
    // narrowing optional to its inner type is only allowed when it's still optional... 
    return false;
  }
  return false;
}

function assignableVariantToUnion(v: { k: "variant"; of: string; name: string }, u: { k: "union"; variants: VariantInfo[] }): boolean {
  const target = u.variants.find((x) => x.name === v.name);
  return target !== undefined;
}

export function fnAssignable(a: Type, b: Type, typeParams: Map<number, Type> = new Map()): boolean {
  if (a.k !== "fun" || b.k !== "fun") return assignable(a, b, typeParams);
  if (a.params.length !== b.params.length) return false;
  if (a.async && !b.async) {
    // async fn can be used where a sync fn is expected? No — awaiting mismatch. Sync where async expected is fine.
  }
  for (let i = 0; i < a.params.length; i++) {
    if (!assignable(b.params[i]!, a.params[i]!, typeParams)) return false;
  }
  return assignable(a.ret, b.ret, typeParams);
}

// Unify a type parameter with an actual type, recording the binding.
export function unify(param: Type, actual: Type, typeParams: Map<number, Type>): void {
  if (param.k === "typeparam") {
    const existing = typeParams.get(param.id);
    if (existing) {
      if (!assignable(actual, existing, typeParams) && !assignable(existing, actual, typeParams)) {
        // incompatible — keep first binding
      }
      return;
    }
    typeParams.set(param.id, actual);
    return;
  }
  if (param.k === "list" && actual.k === "list") unify(param.elem, actual.elem, typeParams);
  else if (param.k === "set" && actual.k === "set") unify(param.elem, actual.elem, typeParams);
  else if (param.k === "map" && actual.k === "map") {
    unify(param.key, actual.key, typeParams);
    unify(param.value, actual.value, typeParams);
  } else if (param.k === "tuple" && actual.k === "tuple") {
    param.items.forEach((p, i) => actual.items[i] && unify(p, actual.items[i]!, typeParams));
  } else if (param.k === "fun" && actual.k === "fun") {
    param.params.forEach((p, i) => actual.params[i] && unify(p, actual.params[i]!, typeParams));
    unify(param.ret, actual.ret, typeParams);
  } else if (param.k === "optional" && actual.k === "optional") {
    unify(param.inner, actual.inner, typeParams);
  } else if (param.k === "result" && actual.k === "result") {
    unify(param.inner, actual.inner, typeParams);
  }
}

// Instantiate a generic type with bindings.
export function instantiate(t: Type, typeParams: Map<number, Type>): Type {
  switch (t.k) {
    case "typeparam": {
      const inst = typeParams.get(t.id);
      return inst ?? t;
    }
    case "optional":
      return optional(instantiate(t.inner, typeParams));
    case "result":
      return result(instantiate(t.inner, typeParams));
    case "list":
      return list(instantiate(t.elem, typeParams));
    case "set":
      return set(instantiate(t.elem, typeParams));
    case "map":
      return mapType(instantiate(t.key, typeParams), instantiate(t.value, typeParams));
    case "tuple":
      return tuple(t.items.map((i) => instantiate(i, typeParams)));
    case "fun":
      return fun(t.params.map((p) => instantiate(p, typeParams)), instantiate(t.ret, typeParams), t.async);
    case "struct":
      return { ...t, fields: t.fields.map((f) => ({ name: f.name, type: instantiate(f.type, typeParams) })), typeArgs: t.typeArgs.map((a) => instantiate(a, typeParams)) };
    case "union":
      return { ...t, variants: t.variants.map((v) => ({ ...v, fields: v.fields.map((f) => ({ name: f.name, type: instantiate(f.type, typeParams) })) })) };
    case "variant":
      return { ...t, fields: t.fields.map((f) => ({ name: f.name, type: instantiate(f.type, typeParams) })) };
    default:
      return t;
  }
}

export function freshTypeParams(params: string[]): { typeParams: Type[]; map: Map<number, Type> } {
  const typeParams = params.map((n, i) => typeparam(n, i));
  return { typeParams, map: new Map() };
}