import type { Context } from "@earendil-works/chord";
import { awaitWithContext } from "@earendil-works/chord/context";
import type {
  Agent,
  AnyTask,
  ConversationId,
  RegistryReader,
  Task,
} from "@earendil-works/pi-durable";

/** Adapt phase agents until pi-durable supports conversation-scoped tool registration. */
export function conversationRegistry(
  registry: RegistryReader,
  resolve: (agent: Agent, id: ConversationId, context: Context) => Promise<Agent>,
): RegistryReader {
  return {
    subscribe: (listener) => registry.subscribe(listener),
    snapshot() {
      const snapshot = registry.snapshot();
      const tasks = new Map<string, AnyTask>();
      type Handler = Task<unknown, { phase: string }, unknown, object>["definition"]["abort"];
      const wrap = (handler: Handler): Handler => (record, runtime, context) => {
        let agent: Promise<Agent> | undefined;
        const adapted = new Proxy(runtime, {
          get(target, key) {
            if (key === "agent") {
              return async (during: Context) => {
                const base = await target.agent(during);
                agent ??= resolve(base, target.conversationId, during);
                return await awaitWithContext(agent, during);
              };
            }
            return Reflect.get(target, key);
          },
        });
        return handler(record, adapted, context);
      };
      for (const original of snapshot.tasks()) {
        const task = original as Task<unknown, { phase: string }, unknown, object>;
        const phases = Object.fromEntries(
          Object.entries(task.definition.phases).map(([name, phase]) => [name, wrap(phase)]),
        );
        tasks.set(task.definition.name, {
          ...task,
          definition: { ...task.definition, phases, abort: wrap(task.definition.abort) },
        });
      }
      return {
        installed: () => snapshot.installed(),
        extension: (name) => snapshot.extension(name),
        tools: () => snapshot.tools(),
        sections: () => snapshot.sections(),
        tasks: () => [...tasks.values()],
        task: (name) => tasks.get(name),
      };
    },
  };
}