// memoize.ts (Bun-compatible, no Node Buffers)
import { mkdir } from "node:fs/promises";
import path from "node:path";

type MemoOpts<V> = {
  ttlMs?: number; // time-to-live for entries
  maxEntries?: number; // cap; oldest (LRU) evicted
  persistDir?: string; // set to enable disk cache (e.g. ".cache/memo")
  keyFn?: (...args: any[]) => string; // custom key if you need it
  cacheErrors?: boolean; // default false
};

type Entry<V> = {
  v: V;
  exp: number | null;
  at: number; // last hit (LRU)
};

const enc = new TextEncoder();
const dec = new TextDecoder();

const stableStringify = (x: any): string => {
  const seen = new WeakSet();
  const S = (v: any): any => {
    if (v && typeof v === "object") {
      if (seen.has(v)) return "[Circular]";
      seen.add(v);
      if (Array.isArray(v)) return v.map(S);
      return Object.fromEntries(
        Object.keys(v)
          .sort()
          .map((k) => [k, S(v[k])]),
      );
    }
    if (typeof v === "function") return `[Function:${v.name || "anon"}]`;
    if (typeof v === "undefined") return "__undefined__";
    return v;
  };
  return JSON.stringify(S(x));
};

async function sha256Hex(s: string) {
  const h = await crypto.subtle.digest("SHA-256", enc.encode(s));
  return Array.from(new Uint8Array(h))
    .map((b) => b.toString(16).padStart(2, "0"))
    .join("");
}

function now() {
  return Date.now();
}

export function memoize<
  F extends (...args: any[]) => any,
  V = Awaited<ReturnType<F>>,
>(fn: F, opts: MemoOpts<V> = {}): F {
  const ttl = opts.ttlMs ?? 0;
  const max = opts.maxEntries ?? 0;
  const dir = opts.persistDir ? path.resolve(opts.persistDir) : null;

  const mem = new Map<string, Entry<V>>();
  const inflight = new Map<string, Promise<V>>();

  async function keyOf(args: any[]): Promise<string> {
    const base = opts.keyFn ? opts.keyFn(...args) : stableStringify(args);
    return dir ? await sha256Hex(base) : base; // hash when persisting (safe filename)
  }

  async function readDisk(k: string): Promise<Entry<V> | null> {
    if (!dir) throw new Error("no dir!");
    const f = Bun.file(path.join(dir, `${k}.json`));
    if (!(await f.exists())) return null;
    try {
      const obj = JSON.parse(await f.text());
      return obj as Entry<V>;
    } catch {
      return null;
    }
  }

  async function writeDisk(k: string, e: Entry<V>) {
    if (!dir) throw new Error("no dir!");
    await Bun.write(path.join(dir, `${k}.json`), JSON.stringify(e));
  }

  function evictLRU() {
    if (!max || mem.size <= max) return;
    const arr = [...mem.entries()].sort((a, b) => a[1].at - b[1].at);
    for (let i = 0; i < mem.size - max; i++) mem.delete(arr[i][0]);
  }

  async function getOrCall(args: any[]): Promise<V> {
    const k = await keyOf(args);
    const t = now();

    // in-flight coalescing
    if (inflight.has(k)) return inflight.get(k)!;

    // memory hit
    const m = mem.get(k);
    if (m && (!m.exp || t < m.exp)) {
      m.at = t;
      return m.v;
    }

    // disk hit
    const d = await readDisk(k);
    if (d && (!d.exp || t < d.exp)) {
      d.at = t;
      mem.set(k, d);
      evictLRU();
      return d.v;
    }

    // miss → call underlying
    const call = (async () => {
      try {
        const r = fn.apply(undefined, args);
        const v: V = r instanceof Promise ? await r : (r as V);
        const e: Entry<V> = { v, exp: ttl ? t + ttl : null, at: t };
        mem.set(k, e);
        evictLRU();
        await writeDisk(k, e);
        return v;
      } catch (err) {
        if (opts.cacheErrors) {
          const e: Entry<any> = { v: err, exp: ttl ? t + ttl : null, at: t };
          mem.set(k, e);
          await writeDisk(k, e as Entry<V>);
        }
        throw err;
      } finally {
        inflight.delete(k);
      }
    })();

    inflight.set(k, call);
    return call;
  }

  // Wrap preserving arity & `this` for methods
  const wrapped = function (this: any, ...args: any[]) {
    const maybe = getOrCall(args).then((v) => v);
    // If original fn is sync (per your signature), unwrap to sync only when it's truly sync.
    // We can detect by calling without awaiting once—dangerous—so be conservative:
    return maybe as unknown as ReturnType<F>;
  } as any as F;

  return wrapped;
}