// SPDX-FileCopyrightText: Amolith <amolith@secluded.site>
//
// SPDX-License-Identifier: GPL-3.0-or-later

import { Agent, type AgentEvent, type AgentTool } from "@mariozechner/pi-agent-core";
import { getEnvApiKey, type AssistantMessage } from "@mariozechner/pi-ai";
import type { RumiloConfig } from "../config/schema.js";
import { resolveModel } from "./model-resolver.js";
import { AgentError } from "../util/errors.js";
import { resolveConfigValue } from "../util/env.js";

export interface AgentRunOptions {
  model: string;
  systemPrompt: string;
  tools: AgentTool[];
  onEvent?: (event: AgentEvent) => void;
  config: RumiloConfig;
}

export interface AgentRunResult {
  message: string;
  usage?: unknown;
  requestCount: number;
}

/**
 * Build a getApiKey callback for the Agent.
 *
 * Resolution order:
 * 1. Custom model config — if a custom model for this provider defines an
 *    `apiKey` field, resolve it via `resolveConfigValue` (supports env var
 *    names, `$VAR` references, and `!shell` commands).
 * 2. pi-ai’s built-in env-var lookup (`ANTHROPIC_API_KEY`, `OPENAI_API_KEY`, etc.).
 */
export function buildGetApiKey(config: RumiloConfig): (provider: string) => string | undefined {
  return (provider: string) => {
    if (config.custom_models) {
      for (const model of Object.values(config.custom_models)) {
        if (model.provider === provider && model.api_key) {
          return resolveConfigValue(model.api_key);
        }
      }
    }

    return getEnvApiKey(provider);
  };
}

export async function runAgent(query: string, options: AgentRunOptions): Promise<AgentRunResult> {
  const agent = new Agent({
    initialState: {
      systemPrompt: options.systemPrompt,
      model: resolveModel(options.model, options.config),
      tools: options.tools,
    },
    getApiKey: buildGetApiKey(options.config),
  });

  if (options.onEvent) {
    agent.subscribe(options.onEvent);
  }

  await agent.prompt(query);

  // Check for errors in agent state
  if (agent.state.error) {
    throw new AgentError(agent.state.error);
  }

  const last = agent.state.messages
    .slice()
    .reverse()
    .find((msg): msg is AssistantMessage => msg.role === "assistant");

  // Check if the last assistant message indicates an error
  if (last?.stopReason === "error") {
    throw new AgentError(last.errorMessage ?? "Agent stopped with an unknown error");
  }

  const text = last?.content
    ?.filter((content): content is Extract<typeof content, { type: "text" }> => content.type === "text")
    .map((content) => content.text)
    .join("")
    .trim();

  if (text === undefined || text === "") {
    throw new AgentError("Agent completed without producing a text response");
  }

  const requestCount = agent.state.messages.filter((msg) => msg.role === "assistant").length;

  return {
    message: text,
    usage: last?.usage,
    requestCount,
  };
}