import type { Env } from "../../../src/core/eval/env";
import { envEmpty, envSet } from "../../../src/core/eval/env";
import type { Store } from "../../../src/core/eval/store";
import type { State, Frame, StepOutcome } from "../../../src/core/eval/machine";
import type { Val } from "../../../src/core/eval/values";
import { VUnit, VTrue, VFalse } from "../../../src/core/eval/values";
import { registerConditionPrims } from "../../../src/core/conditions/prims";

function deepEqual(a: Val, b: Val): boolean {
  if (a.tag !== b.tag) return false;
  switch (a.tag) {
    case "Unit":
      return true;
    case "Num":
      return (b as { n: number }).n === a.n;
    case "Bool":
      return (b as { b: boolean }).b === a.b;
    case "Str":
      return (b as { s: string }).s === a.s;
    case "Sym":
      return (b as { name: string }).name === a.name;
    case "Vector": {
      const other = b as { items: Val[] };
      if (other.items.length !== a.items.length) return false;
      for (let i = 0; i < a.items.length; i++) {
        if (!deepEqual(a.items[i], other.items[i])) return false;
      }
      return true;
    }
    default:
      return JSON.stringify(a) === JSON.stringify(b);
  }
}

function listToArray(value: Val): Val[] {
  const items: Val[] = [];
  let current = value;
  while (current.tag === "Vector" && current.items.length >= 2) {
    items.push(current.items[0]);
    current = current.items[1];
  }
  return items;
}

function arrayToList(items: Val[]): Val {
  let result: Val = VUnit;
  for (let i = items.length - 1; i >= 0; i--) {
    result = { tag: "Vector", items: [items[i], result] };
  }
  return result;
}

export function installMembranePrims(store: Store): { env: Env; store: Store } {
  let env: Env = envEmpty();
  let st: Store = store;

  function def(name: string, v: Val) {
    const [nextStore, addr] = st.alloc(v);
    st = nextStore;
    env = envSet(env, name, addr);
  }

  function isCallable(proc: Val): proc is Val {
    return proc.tag === "Native" || proc.tag === "Closure";
  }

  function ensureArity(proc: Val, expected: number, name: string): void {
    if (proc.tag === "Native") {
      if (proc.arity !== "variadic" && proc.arity !== expected) {
        throw new Error(`${name}: expected procedure of arity ${expected}`);
      }
      return;
    }
    if (proc.tag === "Closure") {
      if (proc.params.length !== expected) {
        throw new Error(`${name}: expected procedure of arity ${expected}`);
      }
      return;
    }
    throw new Error(`${name}: expected a procedure`);
  }

  function applyProcedure(proc: Val, procArgs: Val[], state: State): State | StepOutcome {
    if (proc.tag === "Native") {
      return proc.fn(procArgs, state);
    }
    if (proc.tag === "Closure") {
      if (proc.params.length !== procArgs.length) {
        throw new Error(`procedure arity mismatch: expected ${proc.params.length}, got ${procArgs.length}`);
      }
      let nextStore = state.store;
      let nextEnv = proc.env;
      for (let i = 0; i < proc.params.length; i++) {
        const [allocatedStore, addr] = nextStore.alloc(procArgs[i]);
        nextStore = allocatedStore;
        nextEnv = envSet(nextEnv, proc.params[i], addr);
      }
      const kont = state.kont.concat([{ tag: "KCall", savedEnv: state.env } as Frame]);
      return { ...state, control: { tag: "Expr", e: proc.body }, env: nextEnv, store: nextStore, kont };
    }
    throw new Error("expected procedure");
  }

  function applyNative(proc: Val, procArgs: Val[], state: State): Val | null {
    if (proc.tag !== "Native") {
      return null;
    }
    const result = proc.fn(procArgs, state);
    if ("control" in result && result.control.tag === "Val") {
      return result.control.v;
    }
    return null;
  }

  registerConditionPrims(def, { applyProcedure, ensureArity, isCallable });

  def("*uninit*", { tag: "Native", name: "*uninit*", arity: 0, fn: (_args, s) => {
    return { ...s, control: { tag: "Val", v: VUnit } };
  }});

  def("+", { tag: "Native", name: "+", arity: "variadic", fn: (args, s) => {
    return { ...s, control: { tag: "Val", v: { tag: "Num", n: args.reduce((acc, value) => acc + (value as { n: number }).n, 0) } } };
  }});

  def("-", { tag: "Native", name: "-", arity: "variadic", fn: (args, s) => {
    const nums = args.map((value) => (value as { n: number }).n);
    const n = nums.length === 1 ? -nums[0] : nums.slice(1).reduce((acc, value) => acc - value, nums[0]);
    return { ...s, control: { tag: "Val", v: { tag: "Num", n } } };
  }});

  def("*", { tag: "Native", name: "*", arity: "variadic", fn: (args, s) => {
    return { ...s, control: { tag: "Val", v: { tag: "Num", n: args.reduce((acc, value) => acc * (value as { n: number }).n, 1) } } };
  }});

  def("=", { tag: "Native", name: "=", arity: 2, fn: (args, s) => {
    return { ...s, control: { tag: "Val", v: (args[0] as { n: number }).n === (args[1] as { n: number }).n ? VTrue : VFalse } };
  }});

  def("<", { tag: "Native", name: "<", arity: 2, fn: (args, s) => {
    return { ...s, control: { tag: "Val", v: (args[0] as { n: number }).n < (args[1] as { n: number }).n ? VTrue : VFalse } };
  }});

  def(">", { tag: "Native", name: ">", arity: 2, fn: (args, s) => {
    return { ...s, control: { tag: "Val", v: (args[0] as { n: number }).n > (args[1] as { n: number }).n ? VTrue : VFalse } };
  }});

  def("<=", { tag: "Native", name: "<=", arity: 2, fn: (args, s) => {
    return { ...s, control: { tag: "Val", v: (args[0] as { n: number }).n <= (args[1] as { n: number }).n ? VTrue : VFalse } };
  }});

  def(">=", { tag: "Native", name: ">=", arity: 2, fn: (args, s) => {
    return { ...s, control: { tag: "Val", v: (args[0] as { n: number }).n >= (args[1] as { n: number }).n ? VTrue : VFalse } };
  }});

  def("/", { tag: "Native", name: "/", arity: 2, fn: (args, s) => {
    return { ...s, control: { tag: "Val", v: { tag: "Num", n: (args[0] as { n: number }).n / (args[1] as { n: number }).n } } };
  }});

  def("modulo", { tag: "Native", name: "modulo", arity: 2, fn: (args, s) => {
    return { ...s, control: { tag: "Val", v: { tag: "Num", n: (args[0] as { n: number }).n % (args[1] as { n: number }).n } } };
  }});

  def("not", { tag: "Native", name: "not", arity: 1, fn: (args, s) => {
    const value = args[0];
    const isFalsey = (value.tag === "Bool" && !value.b) || value.tag === "Unit";
    return { ...s, control: { tag: "Val", v: isFalsey ? VTrue : VFalse } };
  }});

  def("or", { tag: "Native", name: "or", arity: "variadic", fn: (args, s) => {
    for (const arg of args) {
      const truthy = !(arg.tag === "Bool" && !arg.b) && arg.tag !== "Unit";
      if (truthy) {
        return { ...s, control: { tag: "Val", v: arg } };
      }
    }
    return { ...s, control: { tag: "Val", v: VFalse } };
  }});

  def("and", { tag: "Native", name: "and", arity: "variadic", fn: (args, s) => {
    let last: Val = VTrue;
    for (const arg of args) {
      const truthy = !(arg.tag === "Bool" && !arg.b) && arg.tag !== "Unit";
      if (!truthy) {
        return { ...s, control: { tag: "Val", v: arg } };
      }
      last = arg;
    }
    return { ...s, control: { tag: "Val", v: last } };
  }});

  def("cons", { tag: "Native", name: "cons", arity: 2, fn: (args, s) => {
    return { ...s, control: { tag: "Val", v: { tag: "Vector", items: [args[0], args[1]] } } };
  }});

  def("car", { tag: "Native", name: "car", arity: 1, fn: (args, s) => {
    const value = args[0] as { tag: string; items?: Val[] };
    if (value.tag === "Vector" && value.items && value.items.length >= 1) {
      return { ...s, control: { tag: "Val", v: value.items[0] } };
    }
    throw new Error("car: expected pair");
  }});

  def("cdr", { tag: "Native", name: "cdr", arity: 1, fn: (args, s) => {
    const value = args[0] as { tag: string; items?: Val[] };
    if (value.tag === "Vector" && value.items && value.items.length >= 2) {
      return { ...s, control: { tag: "Val", v: value.items[1] } };
    }
    throw new Error("cdr: expected pair");
  }});

  def("null?", { tag: "Native", name: "null?", arity: 1, fn: (args, s) => {
    const value = args[0] as { tag: string; items?: Val[] };
    const isNull = value.tag === "Unit" || (value.tag === "Vector" && value.items?.length === 0);
    return { ...s, control: { tag: "Val", v: isNull ? VTrue : VFalse } };
  }});

  def("pair?", { tag: "Native", name: "pair?", arity: 1, fn: (args, s) => {
    const value = args[0] as { tag: string; items?: Val[] };
    const isPair = value.tag === "Vector" && (value.items?.length ?? 0) >= 2;
    return { ...s, control: { tag: "Val", v: isPair ? VTrue : VFalse } };
  }});

  def("list?", { tag: "Native", name: "list?", arity: 1, fn: (args, s) => {
    const value = args[0] as { tag: string; items?: Val[] };
    const isNull = value.tag === "Unit" || (value.tag === "Vector" && value.items?.length === 0);
    const isPair = value.tag === "Vector" && (value.items?.length ?? 0) >= 2;
    return { ...s, control: { tag: "Val", v: isNull || isPair ? VTrue : VFalse } };
  }});

  def("list", { tag: "Native", name: "list", arity: "variadic", fn: (args, s) => {
    return { ...s, control: { tag: "Val", v: arrayToList(args) } };
  }});

  def("reverse", { tag: "Native", name: "reverse", arity: 1, fn: (args, s) => {
    const items = listToArray(args[0]);
    let result: Val = VUnit;
    for (const item of items) {
      result = { tag: "Vector", items: [item, result] };
    }
    return { ...s, control: { tag: "Val", v: result } };
  }});

  def("string-length", { tag: "Native", name: "string-length", arity: 1, fn: (args, s) => {
    const value = args[0] as { tag: string; s?: string };
    if (value.tag !== "Str") throw new Error("string-length: expected string");
    return { ...s, control: { tag: "Val", v: { tag: "Num", n: value.s!.length } } };
  }});

  def("string-append", { tag: "Native", name: "string-append", arity: "variadic", fn: (args, s) => {
    let result = "";
    for (const arg of args) {
      if (arg.tag !== "Str") throw new Error("string-append: expected strings");
      result += arg.s;
    }
    return { ...s, control: { tag: "Val", v: { tag: "Str", s: result } } };
  }});

  def("substring", { tag: "Native", name: "substring", arity: 3, fn: (args, s) => {
    const str = args[0];
    const start = args[1];
    const end = args[2];
    if (str.tag !== "Str") throw new Error("substring: expected string");
    if (start.tag !== "Num" || end.tag !== "Num") throw new Error("substring: expected numbers for indices");
    return { ...s, control: { tag: "Val", v: { tag: "Str", s: str.s.substring(start.n, end.n) } } };
  }});

  def("get", { tag: "Native", name: "get", arity: "variadic", fn: (args, s) => {
    if (args.length < 2 || args.length > 3) throw new Error("get: expected (get coll key) or (get coll key default)");
    const coll = args[0];
    const key = args[1];
    const defaultVal: Val = args.length === 3 ? args[2] : VUnit;

    if (coll.tag === "Map") {
      for (const [entryKey, entryValue] of coll.entries) {
        if (deepEqual(entryKey, key)) {
          return { ...s, control: { tag: "Val", v: entryValue } };
        }
      }
      return { ...s, control: { tag: "Val", v: defaultVal } };
    }

    if (coll.tag === "Vector" && key.tag === "Num") {
      if (key.n >= 0 && key.n < coll.items.length) {
        return { ...s, control: { tag: "Val", v: coll.items[key.n] } };
      }
      return { ...s, control: { tag: "Val", v: defaultVal } };
    }

    return { ...s, control: { tag: "Val", v: defaultVal } };
  }});

  def("equal?", { tag: "Native", name: "equal?", arity: 2, fn: (args, s) => {
    return { ...s, control: { tag: "Val", v: deepEqual(args[0], args[1]) ? VTrue : VFalse } };
  }});

  def("map", { tag: "Native", name: "map", arity: 2, fn: (args, s) => {
    const proc = args[0];
    const items = listToArray(args[1]);

    if (items.length === 0) {
      return { ...s, control: { tag: "Val", v: VUnit } };
    }

    if (proc.tag === "Native") {
      const results: Val[] = [];
      for (const item of items) {
        const result = applyNative(proc, [item], s);
        if (result === null) throw new Error("map: Native function did not return a value");
        results.push(result);
      }
      return { ...s, control: { tag: "Val", v: arrayToList(results) } };
    }

    if (proc.tag === "Closure") {
      const [first, ...rest] = items;
      const frame: Frame = { tag: "KMapRest", fn: proc, remaining: rest, acc: [], env: s.env };
      return applyProcedure(proc, [first], { ...s, kont: [...s.kont, frame] });
    }

    throw new Error("map: expected procedure");
  }});

  def("filter", { tag: "Native", name: "filter", arity: 2, fn: (args, s) => {
    const proc = args[0];
    const items = listToArray(args[1]);

    if (items.length === 0) {
      return { ...s, control: { tag: "Val", v: VUnit } };
    }

    if (proc.tag === "Native") {
      const results: Val[] = [];
      for (const item of items) {
        const result = applyNative(proc, [item], s);
        if (result === null) throw new Error("filter: Native function did not return a value");
        if ((result.tag === "Bool" && result.b) || (result.tag !== "Bool" && result.tag !== "Unit")) {
          results.push(item);
        }
      }
      return { ...s, control: { tag: "Val", v: arrayToList(results) } };
    }

    if (proc.tag === "Closure") {
      const [first, ...rest] = items;
      const frame: Frame = { tag: "KFilterRest", fn: proc, remaining: rest, currentItem: first, acc: [], env: s.env };
      return applyProcedure(proc, [first], { ...s, kont: [...s.kont, frame] });
    }

    throw new Error("filter: expected procedure");
  }});

  return { env, store: st };
}