import { randomUUID } from "node:crypto";

import type {
  Extension,
  ExtensionRuntimeContext,
} from "../../../kernel";
import { Event, ExtensionId, Hook } from "../../catalog";
import {
  type ChatMessage,
  type WorkspaceService,
} from "../workspace";

export interface ModelProvider {
  id: string;
  chat(messages: ChatMessage[], signal?: AbortSignal): AsyncIterable<string>;
}

export interface AgentService {
  chat(input: string, signal?: AbortSignal): AsyncIterable<string>;
}

export function createAgentExtension(): Extension {
  let runtime: ExtensionRuntimeContext | undefined;
  let workspace: WorkspaceService | undefined;
  let model: ModelProvider | undefined;

  const agent: AgentService = {
    async *chat(input, signal) {
      if (!runtime || !workspace || !model) {
        throw new Error("Agent is not running.");
      }

      const activeRuntime = runtime;
      const activeWorkspace = workspace;
      const activeModel = model;
      const runId = randomUUID();

      const userMessage = await activeWorkspace.append("user", input);
      await activeRuntime.emit(Event.MessageAdded, {
        workspaceId: activeWorkspace.workspaceId,
        conversationId: activeWorkspace.conversationId,
        message: userMessage,
      });
      await activeRuntime.emit(Event.RunStarted, { runId, input });

      let answer = "";

      try {
        const messages = await activeWorkspace.messages();

        for await (const chunk of activeModel.chat(messages, signal)) {
          answer += chunk;
          yield chunk;
        }

        const assistantMessage = await activeWorkspace.append("assistant", answer);
        await activeRuntime.emit(Event.MessageAdded, {
          workspaceId: activeWorkspace.workspaceId,
          conversationId: activeWorkspace.conversationId,
          message: assistantMessage,
        });
        await activeRuntime.emit(Event.RunCompleted, { runId });
      } catch (error) {
        await activeRuntime.emit(Event.RunFailed, { runId, error: String(error) });
        throw error;
      }
    },
  };

  return {
    setup(context) {
      context.add(Hook.Agent, agent);
    },

    start(context) {
      const providers = context.all<ModelProvider>(Hook.ModelProviders);

      if (providers.length === 0) {
        throw new Error("Agent needs at least one model provider.");
      }

      runtime = context;
      workspace = context.get<WorkspaceService>(Hook.Workspace);
      model = providers[0];
    },

    stop() {
      runtime = undefined;
      workspace = undefined;
      model = undefined;
    },
  };
}

createAgentExtension.id = ExtensionId.Agent;