import sharp from "sharp";

import type {
  GenerationAdapter,
  GenerationAdapterRequest,
  GenerationAdapterResult,
  NormalizedGenerationOutput,
} from "./ai-adapter-contract.js";
import { gptImageRequestSizeForRatio, normalizeImageOutputToRatio } from "./image-output-normalizer.mjs";

const geminiProductModelId = "gemini-3.1-flash-image-preview";
const geminiProviderModelId = "gemini-3.1-flash-image";
const gptImageModelId = "gpt-image-2";
const geminiEndpoint = "https://oneapi.intelligrow.cn/v1/chat/completions";
const gptImageEndpoint = "https://oneapi.intelligrow.cn/v1/images/generations";
const gptImageReferenceEndpoint = "https://oneapi.intelligrow.cn/v1/images/edits";
const maximumResponseBytes = 32 * 1024 * 1024;
const requestTimeoutMilliseconds = 180_000;

type FetchLike = typeof fetch;

class OneApiRuntimeError extends Error {
  constructor(
    readonly category: "gateway_balance_insufficient" | "gateway_contract_invalid" | "model_disabled" | "reference_invalid" | "upstream_failed" | "upstream_timeout" | "unknown_non_retryable",
    readonly sourceCategory: string,
  ) {
    super(sourceCategory);
  }
}

function failure(error: unknown): GenerationAdapterResult {
  if (error instanceof OneApiRuntimeError) {
    return { category: error.category, sourceCategory: error.sourceCategory, status: "failed" };
  }
  if (error instanceof Error && error.name === "AbortError") {
    return { category: "upstream_timeout", sourceCategory: "upstream_timeout", status: "failed" };
  }
  return { category: "upstream_failed", sourceCategory: "upstream_failed", status: "failed" };
}

function mapHttpFailure(status: number) {
  if (status === 408 || status === 504) return new OneApiRuntimeError("upstream_timeout", `upstream_http_${status}`);
  if (status === 429) return new OneApiRuntimeError("gateway_balance_insufficient", "upstream_http_429");
  if (status >= 500) return new OneApiRuntimeError("upstream_failed", `upstream_http_${status}`);
  if (status === 400 || status === 404 || status === 422) return new OneApiRuntimeError("gateway_contract_invalid", `upstream_http_${status}`);
  return new OneApiRuntimeError("unknown_non_retryable", `upstream_http_${status}`);
}

async function readBoundedJson(response: Response) {
  const declaredLength = Number(response.headers.get("content-length") ?? 0);
  if (Number.isFinite(declaredLength) && declaredLength > maximumResponseBytes) {
    throw new OneApiRuntimeError("gateway_contract_invalid", "upstream_response_too_large");
  }
  if (!response.body) throw new OneApiRuntimeError("gateway_contract_invalid", "upstream_response_empty");
  const reader = response.body.getReader();
  const chunks: Buffer[] = [];
  let total = 0;
  try {
    while (true) {
      const next = await reader.read();
      if (next.done) break;
      const chunk = Buffer.from(next.value);
      total += chunk.length;
      if (total > maximumResponseBytes) {
        await reader.cancel();
        throw new OneApiRuntimeError("gateway_contract_invalid", "upstream_response_too_large");
      }
      chunks.push(chunk);
    }
    try {
      return JSON.parse(Buffer.concat(chunks).toString("utf8")) as unknown;
    } catch {
      throw new OneApiRuntimeError("gateway_contract_invalid", "upstream_response_invalid");
    }
  } finally {
    for (const chunk of chunks) chunk.fill(0);
  }
}

function extractGeminiImage(response: unknown) {
  if (!response || typeof response !== "object" || !("choices" in response) || !Array.isArray(response.choices)) {
    throw new OneApiRuntimeError("gateway_contract_invalid", "response_shape_invalid");
  }
  const choice = response.choices[0];
  const content = choice && typeof choice === "object" && "message" in choice && choice.message && typeof choice.message === "object"
    && "content" in choice.message && typeof choice.message.content === "string" ? choice.message.content : "";
  const matches = [...content.matchAll(/!\[[^\]]*\]\(\s*data:(image\/(?:jpeg|png|webp));base64,([A-Za-z0-9+/=\r\n]+)\s*\)/gi)];
  if (matches.length !== 1) throw new OneApiRuntimeError("gateway_contract_invalid", "response_single_image_required");
  return { bytes: Buffer.from(matches[0]![2]!, "base64"), declaredMimeType: matches[0]![1]!.toLowerCase() };
}

function extractGptImage(response: unknown) {
  if (!response || typeof response !== "object" || !("data" in response) || !Array.isArray(response.data)
    || response.data.length !== 1 || !response.data[0] || typeof response.data[0] !== "object"
    || !("b64_json" in response.data[0]) || typeof response.data[0].b64_json !== "string") {
    throw new OneApiRuntimeError("gateway_contract_invalid", "response_single_image_required");
  }
  return { bytes: Buffer.from(response.data[0].b64_json, "base64"), declaredMimeType: undefined };
}

async function normalizeOutput(bytes: Buffer, declaredMimeType: string | undefined, ratio: GenerationAdapterRequest["ratio"]): Promise<NormalizedGenerationOutput> {
  try {
    const metadata = await sharp(bytes, { failOn: "error", limitInputPixels: 40_000_000 }).metadata();
    const mimeType = metadata.format === "png" ? "image/png" : metadata.format === "jpeg" ? "image/jpeg" : metadata.format === "webp" ? "image/webp" : undefined;
    if (!mimeType || !metadata.width || !metadata.height || (declaredMimeType && declaredMimeType !== mimeType)) {
      throw new OneApiRuntimeError("gateway_contract_invalid", "response_media_invalid");
    }
    const normalized = await normalizeImageOutputToRatio({ bytes, mimeType, pixelHeight: metadata.height, pixelWidth: metadata.width, ratio });
    return { bytes: normalized.bytes, mimeType: normalized.mimeType, pixelHeight: normalized.pixelHeight, pixelWidth: normalized.pixelWidth };
  } catch (error) {
    if (error instanceof OneApiRuntimeError) throw error;
    throw new OneApiRuntimeError("gateway_contract_invalid", "response_media_invalid");
  }
}

function validateRequest(request: GenerationAdapterRequest) {
  if (!request.prompt.trim() || request.prompt.length > 1_000) throw new OneApiRuntimeError("gateway_contract_invalid", "prompt_invalid");
  const references = request.referenceImages ?? [];
  if (references.length !== request.referenceAssetIds.length || references.length > 2) {
    throw new OneApiRuntimeError("reference_invalid", "reference_count_invalid");
  }
  const totalBytes = references.reduce((total, reference) => total + reference.bytes.length, 0);
  if (totalBytes > 20 * 1024 * 1024 || references.some((reference) => reference.bytes.length === 0 || reference.bytes.length > 10 * 1024 * 1024)) {
    throw new OneApiRuntimeError("reference_invalid", "reference_size_invalid");
  }
  return references;
}

function buildRequest(request: GenerationAdapterRequest) {
  const references = validateRequest(request);
  if (request.modelId === geminiProductModelId) {
    const content = references.length === 0
      ? request.prompt
      : [
          { text: request.prompt, type: "text" },
          ...references.map((reference) => ({
            image_url: { url: `data:${reference.mimeType};base64,${reference.bytes.toString("base64")}` },
            type: "image_url",
          })),
        ];
    return {
      body: JSON.stringify({
        extra_body: { google: { image_config: { aspect_ratio: request.ratio, image_size: "1K" } } },
        messages: [{ content, role: "user" }],
        model: geminiProviderModelId,
        stream: false,
      }),
      contentType: "application/json",
      endpoint: geminiEndpoint,
      parser: extractGeminiImage,
    };
  }
  if (request.modelId !== gptImageModelId) throw new OneApiRuntimeError("model_disabled", "model_not_supported");
  if (references.length > 0) {
    const form = new FormData();
    form.append("model", gptImageModelId);
    form.append("prompt", request.prompt);
    form.append("response_format", "b64_json");
    form.append("size", gptImageRequestSizeForRatio(request.ratio));
    references.forEach((reference, index) => form.append("image[]", new Blob([reference.bytes], { type: reference.mimeType }), `reference-${index + 1}.png`));
    return { body: form, contentType: undefined, endpoint: gptImageReferenceEndpoint, parser: extractGptImage };
  }
  return {
    body: JSON.stringify({ model: gptImageModelId, prompt: request.prompt, response_format: "b64_json", size: gptImageRequestSizeForRatio(request.ratio) }),
    contentType: "application/json",
    endpoint: gptImageEndpoint,
    parser: extractGptImage,
  };
}

export class OneApiGenerationAdapter implements GenerationAdapter {
  private readonly credential: Buffer;
  private readonly fetchImpl: FetchLike;
  private disposed = false;

  constructor(input: { credential: Buffer; fetch?: FetchLike }) {
    if (input.credential.length < 8) throw new Error("ai_gateway_credential_invalid");
    this.credential = Buffer.from(input.credential);
    this.fetchImpl = input.fetch ?? fetch;
  }

  async start(request: GenerationAdapterRequest): Promise<GenerationAdapterResult> {
    if (this.disposed) return { category: "upstream_failed", sourceCategory: "adapter_disposed", status: "failed" };
    const controller = new AbortController();
    const timeout = setTimeout(() => controller.abort(), requestTimeoutMilliseconds);
    let sourceBytes: Buffer | undefined;
    try {
      const providerRequest = buildRequest(request);
      const headers = new Headers({ authorization: `Bearer ${this.credential.toString("utf8")}` });
      if (providerRequest.contentType) headers.set("content-type", providerRequest.contentType);
      const response = await this.fetchImpl(providerRequest.endpoint, {
        body: providerRequest.body,
        headers,
        method: "POST",
        redirect: "error",
        signal: controller.signal,
      });
      if (!response.ok) throw mapHttpFailure(response.status);
      const parsed = await readBoundedJson(response);
      const extracted = providerRequest.parser(parsed);
      sourceBytes = extracted.bytes;
      const output = await normalizeOutput(sourceBytes, extracted.declaredMimeType, request.ratio);
      return { outputs: [output], status: "completed" };
    } catch (error) {
      return failure(error);
    } finally {
      clearTimeout(timeout);
      sourceBytes?.fill(0);
    }
  }

  dispose() {
    if (this.disposed) return;
    this.disposed = true;
    this.credential.fill(0);
  }
}