import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { buildSessionContext, convertToLlm } from "@earendil-works/pi-coding-agent";
import type { Message } from "@earendil-works/pi-ai/compat";

export type Grounding = { getGroundingPrefix(ctx: ExtensionContext): Message[] };

type PrefixOwner = { sessionId: string; leafId: string | null };
type CachedPrefix = PrefixOwner & { prefix: Message[] };

function ownerFor(ctx: ExtensionContext): PrefixOwner {
  return {
    sessionId: ctx.sessionManager.getSessionId(),
    leafId: ctx.sessionManager.getLeafId(),
  };
}

function sameOwner(left: PrefixOwner, right: PrefixOwner): boolean {
  return left.sessionId === right.sessionId && left.leafId === right.leafId;
}

/**
 * Register a read-only `context` listener that caches the converted provider
 * prefix only for its owning session and tree leaf. A prefix must never cross
 * a session or branch boundary, even while Pi is resuming a session.
 */
export function registerGrounding(pi: ExtensionAPI): Grounding {
  let cached: CachedPrefix | undefined;
  const clear = () => { cached = undefined; };

  pi.on("context", (event, ctx) => {
    // A failed conversion or identity lookup must discard, rather than retain,
    // a prefix from a previous context event.
    clear();
    try {
      const owner = ownerFor(ctx);
      cached = { ...owner, prefix: convertToLlm(event.messages) };
    } catch {
      clear();
    }
  });
  pi.on("session_start", clear);
  pi.on("session_tree", (event) => {
    try {
      if (event.oldLeafId !== event.newLeafId) clear();
    } catch {
      clear();
    }
  });
  pi.on("session_shutdown", clear);

  return {
    getGroundingPrefix(ctx) {
      try {
        const owner = ownerFor(ctx);
        if (cached && sameOwner(cached, owner)) return cached.prefix;

        const prefix = convertToLlm(buildSessionContext(ctx.sessionManager.getEntries(), owner.leafId).messages);
        cached = { ...owner, prefix };
        return prefix;
      } catch {
        clear();
        return [];
      }
    },
  };
}