import { env } from "cloudflare:workers";
import type { ClefRequest, ClefResponse } from "./decide";
import { mockProvider } from "./mock";
import type { ChatMessage, MergeFixProvider, TokenUsage } from "./plan";
import { ModelCallError } from "./plan";
import type { AiEndpoint } from "../ai/endpoint";
import { assertOutboundUrl, isAllowedLoopback } from "../integrations/outbound";
import { assertResolvesPublic } from "../net/resolve-guard";

const MAX_OUTPUT_TOKENS = 2048;

type ChatReply = {
  response?: unknown;
  choices?: Array<{ message?: { content?: unknown } }>;
  usage?: { prompt_tokens?: number; completion_tokens?: number; input_tokens?: number; output_tokens?: number };
};

function readReply(reply: ChatReply): { text: string; usage: TokenUsage } {
  const raw = reply.choices?.[0]?.message?.content ?? reply.response;
  if (typeof raw !== "string") throw new ModelCallError("The model returned no text");
  return {
    text: raw,
    usage: {
      inputTokens: reply.usage?.prompt_tokens ?? reply.usage?.input_tokens ?? 0,
      outputTokens: reply.usage?.completion_tokens ?? reply.usage?.output_tokens ?? 0,
    },
  };
}

type Unchecked = { run(model: string, inputs: unknown, options?: AiOptions): Promise<unknown> };

export type GatewayMetadata = { account_id: string; feature: string } & Record<string, string>;

export type ChatTarget = {
  /** Workers AI model used when the account has no bring-your-own endpoint. */
  model: string;
  /** From resolveAiEndpoint: the repo's or account's own endpoint, which pays for the call. */
  byok: AiEndpoint | null;
  metadata: GatewayMetadata;
  maxTokens?: number;
};

/** The model actually billed for a chat call: the account's BYOK model, else the platform one. */
export function chatModel(target: Pick<ChatTarget, "model" | "byok">): string {
  return target.byok?.model ?? target.model;
}

const BYOK_TIMEOUT_MS = 120_000;
const CF_GATEWAY_ORIGIN = "https://gateway.ai.cloudflare.com";

/** The chat/completions URL for a custom endpoint. Custom base URLs are re-validated (host rules and DNS) on every call. */
export async function endpointUrl(endpoint: AiEndpoint): Promise<string> {
  if (endpoint.kind === "cf-gateway") {
    return `${CF_GATEWAY_ORIGIN}/v1/${endpoint.cfAccountId}/${endpoint.cfGatewayId}/compat/chat/completions`;
  }
  const url = assertOutboundUrl(`${endpoint.baseUrl ?? ""}/chat/completions`);
  if (!isAllowedLoopback(url)) await assertResolvesPublic(url);
  return url.toString();
}

async function customCompletion(endpoint: AiEndpoint, body: unknown, metadata: GatewayMetadata) {
  const headers: Record<string, string> = {
    authorization: `Bearer ${endpoint.apiKey}`,
    "content-type": "application/json",
  };
  if (endpoint.kind === "cf-gateway") headers["cf-aig-metadata"] = JSON.stringify(metadata);
  const response = await fetch(await endpointUrl(endpoint), {
    method: "POST",
    headers,
    body: JSON.stringify(body),
    redirect: "manual",
    signal: AbortSignal.timeout(BYOK_TIMEOUT_MS),
  }).catch(() => {
    throw new ModelCallError("The AI endpoint could not be reached");
  });
  if (!response.ok) throw new ModelCallError(`The AI endpoint returned ${response.status}`);
  return readReply((await response.json()) as ChatReply);
}

function byokModel(endpoint: AiEndpoint): string {
  return endpoint.kind === "openai" ? endpoint.model : `${endpoint.provider}/${endpoint.model}`;
}

/**
 * One chat completion, tagged with the billed account so Gateway analytics line up with the usage ledger.
 * Without a bring-your-own endpoint the call runs on Workers AI through the instance's AI Gateway. A 'provider'
 * endpoint uses the owner's key through the instance Gateway's OpenAI-compatible endpoint; 'cf-gateway' and
 * 'openai' endpoints call the owner's own Gateway or OpenAI-compatible server directly.
 */
export async function chatCompletion(
  target: ChatTarget,
  messages: ChatMessage[],
): Promise<{ text: string; usage: TokenUsage }> {
  const gatewayId = env.MERGEFIX_GATEWAY_ID || "default";
  const maxTokens = target.maxTokens ?? MAX_OUTPUT_TOKENS;
  const byok = target.byok;
  if (byok && byok.kind !== "provider") {
    return customCompletion(byok, { model: byokModel(byok), messages, max_tokens: maxTokens }, target.metadata);
  }
  if (byok) {
    const response = await env.AI.gateway(gatewayId).run({
      provider: "compat",
      endpoint: "chat/completions",
      headers: {
        authorization: `Bearer ${byok.apiKey}`,
        "content-type": "application/json",
        "cf-aig-metadata": JSON.stringify(target.metadata),
      } as unknown as AIGatewayUniversalRequest["headers"],
      query: { model: byokModel(byok), messages, max_tokens: maxTokens },
    });
    if (!response.ok) throw new ModelCallError(`The model provider returned ${response.status}`);
    return readReply((await response.json()) as ChatReply);
  }
  const ai = env.AI as unknown as Unchecked;
  const reply = await ai.run(
    target.model,
    { messages, max_tokens: maxTokens },
    { gateway: { id: gatewayId, metadata: target.metadata } },
  );
  return readReply(reply as ChatReply);
}

/** The merge-fix provider: generation via chatCompletion, decisions via Clef, which always runs on Workers AI. */
export function gatewayProvider(input: {
  accountId: string;
  runId: string;
  byok: AiEndpoint | null;
}): MergeFixProvider {
  if (env.MERGEFIX_PROVIDER === "mock") return mockProvider();
  const gatewayId = env.MERGEFIX_GATEWAY_ID || "default";
  const metadata: GatewayMetadata = { account_id: input.accountId, run_id: input.runId, feature: "mergefix" };
  const ai = env.AI as unknown as Unchecked;
  return {
    generate({ messages }: { messages: ChatMessage[] }) {
      return chatCompletion({ model: env.MERGEFIX_MODEL, byok: input.byok, metadata }, messages);
    },
    async decide(request: ClefRequest): Promise<ClefResponse> {
      const model = request.model === "clef-flash" ? "@cf/cloudflare/clef-flash" : "@cf/cloudflare/clef";
      const reply = (await ai.run(model, request, { gateway: { id: gatewayId, metadata } })) as ClefResponse;
      if (!reply || typeof reply.answers !== "object") throw new ModelCallError("Clef returned no answers");
      return reply;
    },
  };
}

export function clefModel(): ClefRequest["model"] {
  return env.MERGEFIX_CLEF_MODEL === "clef-flash" ? "clef-flash" : "clef";
}