import { parseAgentEnv } from "@code/env/agent";
import {
  makePiAgentOsConfig,
  makePiHomeFiles,
  decodeWorkAttemptExecutionInput,
  WorkAttemptExecutionError,
  piSessionEnv,
} from "@code/primitives";
import type {
  ExecutionEvent,
  WorkAttemptExecutionResult,
} from "@code/primitives";
import { agentOS, setup } from "@rivet-dev/agentos";
import { createHostDirBackend } from "@rivet-dev/agentos-core";
import { createClient } from "@rivet-dev/agentos/client";
import { Effect } from "effect";

import { HostRepositoryWorkspace } from "./host-repository";

const piConfig = makePiAgentOsConfig();
const workspace = agentOS<undefined, { token: string }>({
  onBeforeConnect: (_context, params) => {
    if (params.token !== process.env.RIVET_WORKSPACE_TOKEN) {
      throw new Error("Unauthorized workspace connection");
    }
  },
  options: {
    actionTimeout: 10 * 60 * 1000,
  },
  permissions: piConfig.permissions,
  software: piConfig.software,
});

export const runtimeRegistry = setup({ use: { workspace } });

const event = (
  sequence: number,
  kind: ExecutionEvent["kind"],
  message: string,
  metadata: Record<string, string> = {}
): ExecutionEvent => ({
  kind,
  message,
  metadata,
  occurredAt: Date.now(),
  sequence,
});

const executionError = (
  message: string,
  reason: WorkAttemptExecutionError["reason"],
  retryable: boolean
) => new WorkAttemptExecutionError({ message, reason, retryable });

const classifyRuntimeFailure = (cause: unknown): WorkAttemptExecutionError => {
  if (cause instanceof WorkAttemptExecutionError) {
    return cause;
  }
  const message = cause instanceof Error ? cause.message : "Execution failed";
  if (/auth|credential|401|403/iu.test(message)) {
    return executionError(message, "Authentication", false);
  }
  if (/clone|checkout|repository|git /iu.test(message)) {
    return executionError(message, "RepositoryFailed", false);
  }
  if (/timeout|timed out/iu.test(message)) {
    return executionError(message, "Timeout", true);
  }
  if (/connect|network|unavailable|502|503|504/iu.test(message)) {
    return executionError(message, "ProviderUnavailable", true);
  }
  return executionError(message, "HarnessFailed", false);
};

export const executeAgentOsAttempt = async (
  rawInput: unknown
): Promise<WorkAttemptExecutionResult> => {
  try {
    const input = await Effect.runPromise(
      decodeWorkAttemptExecutionInput(rawInput)
    );
    const env = parseAgentEnv(process.env);
    const client = createClient<typeof runtimeRegistry>({
      disableMetadataLookup: true,
      endpoint: env.RIVET_PUBLIC_ENDPOINT ?? env.RIVET_ENDPOINT,
    });
    const vm = client.workspace.getOrCreate([input.workspaceKey], {
      params: { token: env.RIVET_WORKSPACE_TOKEN },
    });
    const events: ExecutionEvent[] = [
      event(0, "runtime.preparing", "AgentOS workspace selected", {
        workspaceKey: input.workspaceKey,
      }),
    ];
    const hostRepository = new HostRepositoryWorkspace();
    const prepared = await hostRepository.prepare({
      attemptId: input.attemptId,
      piHomeFiles: makePiHomeFiles({
        api: env.AGENT_MODEL_API,
        apiKeyEnvironmentVariable: "AGENT_MODEL_API_KEY",
        baseUrl: env.AGENT_MODEL_BASE_URL,
        contextWindow: env.AGENT_MODEL_CONTEXT_WINDOW,
        maxTokens: env.AGENT_MODEL_MAX_TOKENS,
        model: env.AGENT_MODEL_NAME,
        provider: env.AGENT_MODEL_PROVIDER,
      }),
    });
    events.push(
      event(
        1,
        "runtime.preparing",
        prepared.created
          ? "Isolated Zopu worktree created on the execution host"
          : "Isolated Zopu worktree recreated"
      )
    );
    const mounts = [
      {
        hostPath: prepared.checkoutPath,
        path: "/workspace/repository",
        readOnly: false,
      },
      {
        hostPath: prepared.sourceRepositoryPath,
        path: prepared.sourceRepositoryPath,
        readOnly: false,
      },
      {
        hostPath: prepared.piHomePath,
        path: "/home/zopu",
        readOnly: false,
      },
      {
        hostPath: prepared.toolsPath,
        path: "/opt/zopu-tools",
        readOnly: true,
      },
    ] as const;
    const mountNext = async (index: number): Promise<void> => {
      const mount = mounts[index];
      if (!mount) {
        return;
      }
      try {
        await vm.mountFs({
          path: mount.path,
          plugin: createHostDirBackend({
            hostPath: mount.hostPath,
            readOnly: mount.readOnly,
          }),
          readOnly: mount.readOnly,
        });
      } catch (error) {
        const message = error instanceof Error ? error.message : String(error);
        throw executionError(
          `Failed to mount ${mount.hostPath} at ${mount.path}: ${message}`,
          "HarnessFailed",
          false
        );
      }
      await mountNext(index + 1);
    };
    await mountNext(0);
    const { baseRevision } = prepared;
    events.push(
      event(2, "repository.ready", "Repository checkout is ready", {
        baseRevision,
      })
    );

    const sessionId = `pi-${input.attemptId}`;
    await vm.openSession({
      additionalDirectories: [`${prepared.sourceRepositoryPath}/.git`],
      additionalInstructions:
        "Work only inside /workspace/repository. This is an isolated worktree of the Zopu product repository. Read AGENTS.md and the relevant product specifications before changing code. Never reveal credentials, modify the read-only base checkout, push, or open a pull request. Implement the requested change and run focused verification.",
      agent: "pi",
      cwd: "/workspace/repository",
      env: piSessionEnv(env.AGENT_MODEL_API_KEY),
      permissionPolicy: "allow_all",
      sessionId,
    });
    events.push(
      event(3, "harness.started", "Pi implementation session started")
    );
    const promptResult = await vm.prompt({
      content: [{ text: input.prompt, type: "text" }],
      idempotencyKey: input.attemptId,
      sessionId,
    });
    if (promptResult.stopReason !== "end_turn") {
      const reason =
        promptResult.stopReason === "cancelled" ? "Cancelled" : "HarnessFailed";
      throw executionError(
        `Pi stopped with ${promptResult.stopReason}`,
        reason,
        promptResult.stopReason === "max_tokens" ||
          promptResult.stopReason === "max_turn_requests"
      );
    }
    events.push(
      event(4, "harness.progress", "Pi implementation turn completed", {
        stopReason: promptResult.stopReason,
      })
    );

    const collected = await HostRepositoryWorkspace.collect({
      attemptId: input.attemptId,
      baseRevision,
      checkoutPath: prepared.checkoutPath,
    });
    const { candidateRevision, changedFiles, diff } = collected;
    events.push(
      event(
        5,
        "repository.changed",
        `${changedFiles.length} changed file(s) collected`,
        {
          candidateRevision,
        }
      ),
      event(6, "runtime.completed", "AgentOS execution completed")
    );

    return {
      baseRevision,
      candidateRevision,
      changedFiles,
      diff,
      environmentId: input.workspaceKey,
      events,
      summary:
        changedFiles.length > 0
          ? `Pi changed ${changedFiles.length} file(s)`
          : "Pi completed without repository changes",
    };
  } catch (error) {
    throw classifyRuntimeFailure(error);
  }
};

export const cancelAgentOsAttempt = async (
  workspaceKey: string,
  attemptId: string
) => {
  const env = parseAgentEnv(process.env);
  const client = createClient<typeof runtimeRegistry>({
    disableMetadataLookup: true,
    endpoint: env.RIVET_PUBLIC_ENDPOINT ?? env.RIVET_ENDPOINT,
  });
  await client.workspace
    .getOrCreate([workspaceKey], {
      params: { token: env.RIVET_WORKSPACE_TOKEN },
    })
    .cancelPrompt({ sessionId: `pi-${attemptId}` });
};
