import type { Id } from "@code/backend/convex/_generated/dataModel";
import { env } from "@code/env/server";
import { AgentOs, createAgentOsActorEffect } from "@code/primitives/agent-os";
import { Effect, Schema } from "effect";
import { setup } from "rivetkit";
import { createClient } from "rivetkit/client";

import type { DaemonCommand } from "./convex";
import { ConvexControlPlane } from "./convex";

export class AgentOsCommandError extends Schema.TaggedErrorClass<AgentOsCommandError>()(
  "AgentOsCommandError",
  {
    method: Schema.String,
    cause: Schema.Defect(),
  }
) {}

const asArgs = (value: unknown): unknown[] => {
  if (!Array.isArray(value)) {
    throw new TypeError("agentOS command args must be an array");
  }
  return value;
};

export const makeAgentOsRuntime = Effect.gen(function* () {
  const actor = yield* createAgentOsActorEffect();
  const registry = setup({ use: { agentOs: actor } });
  registry.start();
  const client = createClient<typeof registry>(env.RIVET_ENDPOINT);

  const execute = Effect.fn("AgentOsRuntime.execute")(function* (
    command: DaemonCommand
  ) {
    const controlPlane = yield* ConvexControlPlane;
    const handle = client.agentOs.getOrCreate([...command.actorKey]);
    return yield* Effect.tryPromise({
      try: () =>
        handle.action({
          name: command.method,
          args: asArgs(command.args),
        }),
      catch: (cause) =>
        new AgentOsCommandError({ method: command.method, cause }),
    }).pipe(
      Effect.map((result) =>
        result instanceof Uint8Array ? result.slice().buffer : result
      ),
      Effect.tap((result) =>
        controlPlane.recordEvent({
          commandId: command._id,
          kind: "agentos.action.succeeded",
          data: { method: command.method, result },
        })
      ),
      Effect.tapError((error) =>
        controlPlane.recordEvent({
          commandId: command._id,
          kind: "agentos.action.failed",
          data: { method: command.method, error: String(error.cause) },
        })
      )
    );
  });

  const recordLifecycleEvent = (
    commandId: Id<"daemonCommands"> | undefined,
    kind: string,
    data: unknown
  ) =>
    Effect.gen(function* () {
      const controlPlane = yield* ConvexControlPlane;
      yield* controlPlane.recordEvent({ commandId, kind, data });
    });

  return { execute, recordLifecycleEvent } as const;
}).pipe(Effect.provide(AgentOs.layer));
