import "server-only";

import sharp from "sharp";
import { uploadFile } from "@/lib/storage/service";

interface MidjourneyConfig {
  apiKey: string;
  baseURL: string;
}

export interface MidjourneyImageParams {
  aspectRatio?: string;
  imageSize?: string;
  imageCount?: number;
}

interface UploadedImage {
  thumbnail: string;
  original: string;
}

interface MidjourneyTaskButton {
  customId?: string;
  label?: string;
  emoji?: string;
  style?: number;
  type?: number;
}

interface MidjourneyTaskResult {
  id?: string;
  imageUrl?: string;
  status?: string;
  progress?: string;
  failReason?: string;
  description?: string;
  prompt?: string;
  buttons?: MidjourneyTaskButton[];
}

const DEFAULT_TIMEOUT_MS = 300000;
const POLL_INTERVAL_MS = 3500;
const DEFAULT_IMAGE_SIZE = "2K";

function normalizeBaseURL(baseURL: string) {
  return String(baseURL || "").trim().replace(/\/+$/, "");
}

function resolveMode(baseURL: string, modelId: string) {
  const normalizedBaseURL = normalizeBaseURL(baseURL).toLowerCase();
  const normalizedModelId = String(modelId || "").trim().toLowerCase();

  if (
    normalizedBaseURL.endsWith("/relax") ||
    normalizedBaseURL.endsWith("/mj-relax") ||
    normalizedModelId.includes("relax")
  ) {
    return "relax" as const;
  }

  return "fast" as const;
}

function buildEndpointCandidates(baseURL: string, mode: "fast" | "relax", suffix: string) {
  const normalizedBaseURL = normalizeBaseURL(baseURL);
  const normalizedSuffix = suffix.startsWith("/") ? suffix : `/${suffix}`;
  const directPattern = /\/(fast|relax|mj-fast|mj-relax)$/i;

  if (directPattern.test(normalizedBaseURL)) {
    return [`${normalizedBaseURL}${normalizedSuffix}`];
  }

  const modePrefix = mode === "relax" ? "relax" : "fast";
  const altPrefix = mode === "relax" ? "mj-relax" : "mj-fast";

  return [
    `${normalizedBaseURL}/${modePrefix}${normalizedSuffix}`,
    `${normalizedBaseURL}/${altPrefix}${normalizedSuffix}`,
  ];
}

function normalizeImageSizeLabel(imageSize?: string): "1K" | "2K" | "4K" {
  const normalized = String(imageSize || "").trim().toUpperCase();
  if (normalized === "1K" || normalized === "2K" || normalized === "4K") return normalized;
  return DEFAULT_IMAGE_SIZE;
}

function normalizeAspectRatio(value?: string) {
  const normalized = String(value || "").trim();
  if (!normalized || normalized.toLowerCase() === "auto") return undefined;
  if (/^\d+:\d+$/.test(normalized)) return normalized;
  return undefined;
}

function appendAspectRatio(prompt: string, aspectRatio?: string) {
  const normalizedPrompt = String(prompt || "").trim();
  const normalizedAspectRatio = normalizeAspectRatio(aspectRatio);
  if (!normalizedAspectRatio) return normalizedPrompt;
  if (/--ar\s+\d+:\d+/i.test(normalizedPrompt)) return normalizedPrompt;
  return `${normalizedPrompt} --ar ${normalizedAspectRatio}`.trim();
}

function detectImageExtension(imageUrl: string, contentType?: string | null) {
  const type = String(contentType || "").toLowerCase();
  if (type.includes("jpeg") || type.includes("jpg")) return "jpg";
  if (type.includes("webp")) return "webp";
  if (type.includes("png")) return "png";
  if (type.includes("gif")) return "gif";

  const cleanUrl = imageUrl.split("?")[0]?.toLowerCase() || "";
  if (cleanUrl.endsWith(".jpg") || cleanUrl.endsWith(".jpeg")) return "jpg";
  if (cleanUrl.endsWith(".webp")) return "webp";
  if (cleanUrl.endsWith(".gif")) return "gif";
  if (cleanUrl.endsWith(".png")) return "png";
  return "png";
}

async function createThumbnail(buffer: Buffer) {
  try {
    return await sharp(buffer)
      .resize(400, 400, { fit: "inside", withoutEnlargement: true })
      .jpeg({ quality: 80 })
      .toBuffer();
  } catch (error) {
    console.error("[Midjourney] Failed to create thumbnail:", error);
    return buffer;
  }
}

function generateFilename(prefix: string, extension: string) {
  const timestamp = Date.now();
  const random = Math.random().toString(36).slice(2, 8);
  return `${prefix}-${timestamp}-${random}.${extension}`;
}

async function uploadGeneratedBuffer(
  buffer: Buffer,
  extension: string,
  mimeType: string
): Promise<UploadedImage> {
  const thumbnailBuffer = await createThumbnail(buffer);

  const originalFilename = generateFilename("midjourney-orig", extension);
  const originalArrayBuffer = buffer.buffer.slice(
    buffer.byteOffset,
    buffer.byteOffset + buffer.byteLength
  ) as ArrayBuffer;
  const originalUploaded = await uploadFile(originalFilename, originalArrayBuffer);

  const thumbFilename = generateFilename("midjourney-thumb", "jpg");
  const thumbArrayBuffer = thumbnailBuffer.buffer.slice(
    thumbnailBuffer.byteOffset,
    thumbnailBuffer.byteOffset + thumbnailBuffer.byteLength
  ) as ArrayBuffer;
  const thumbUploaded = await uploadFile(thumbFilename, thumbArrayBuffer);

  return {
    thumbnail: thumbUploaded.url,
    original: originalUploaded.url,
  };
}

async function downloadAndUploadImage(imageUrl: string): Promise<UploadedImage> {
  const response = await fetch(imageUrl, { cache: "no-store" });
  if (!response.ok) {
    throw new Error(`Failed to download Midjourney image: ${response.status}`);
  }
  const buffer = Buffer.from(await response.arrayBuffer());
  const extension = detectImageExtension(imageUrl, response.headers.get("content-type"));
  const mimeType =
    response.headers.get("content-type") ||
    `image/${extension === "jpg" ? "jpeg" : extension}`;
  return uploadGeneratedBuffer(buffer, extension, mimeType);
}

function extractPromptAndImages(messages: Array<{ role: string; content: any }>) {
  const lastUserMessage = messages.filter((message) => message.role === "user").pop();
  if (!lastUserMessage) {
    throw new Error("No user message found");
  }

  let prompt = "";
  const imageUrls: string[] = [];

  const pushImageUrl = (value?: string) => {
    if (!value) return;
    const normalized = String(value).trim();
    if (!normalized) return;
    if (!imageUrls.includes(normalized)) {
      imageUrls.push(normalized);
    }
  };

  if (typeof lastUserMessage.content === "string") {
    prompt = lastUserMessage.content.trim();
  } else if (Array.isArray(lastUserMessage.content)) {
    for (const part of lastUserMessage.content) {
      if (part?.type === "text" && typeof part.text === "string") {
        prompt = `${prompt} ${part.text}`.trim();
      } else if (part?.type === "image_url" && part.image_url?.url) {
        pushImageUrl(part.image_url.url);
      } else if (part?.type === "file") {
        pushImageUrl(part.url || part.data);
      } else if (part?.type === "image") {
        pushImageUrl(part.image || part.url);
      }
    }
  }

  if (!prompt) {
    throw new Error("No prompt found in user message");
  }

  return { prompt, imageUrls };
}

async function loadInputImageAsDataUrl(source: string) {
  const normalized = String(source || "").trim();
  if (!normalized) {
    throw new Error("Midjourney 输入图片地址为空");
  }

  if (normalized.startsWith("data:")) {
    return normalized;
  }

  const response = await fetch(normalized, { cache: "no-store" });
  if (!response.ok) {
    throw new Error(`Failed to fetch Midjourney input image: ${response.status}`);
  }

  const mimeType = response.headers.get("content-type") || "image/png";
  const buffer = Buffer.from(await response.arrayBuffer());
  return `data:${mimeType};base64,${buffer.toString("base64")}`;
}

async function parseJsonResponse(response: Response) {
  const rawText = await response.text();
  try {
    return JSON.parse(rawText);
  } catch {
    return rawText;
  }
}

function extractErrorMessage(payload: any) {
  if (!payload) return "Midjourney 请求失败";
  if (typeof payload === "string") return payload;
  if (payload.failReason) return String(payload.failReason);
  if (payload.description) return String(payload.description);
  if (payload.message) return String(payload.message);
  if (payload.error?.message) return String(payload.error.message);
  try {
    return JSON.stringify(payload);
  } catch {
    return "Midjourney 请求失败";
  }
}

async function requestWithEndpointFallback(params: {
  baseURL: string;
  mode: "fast" | "relax";
  suffix: string;
  method: "GET" | "POST";
  apiKey: string;
  body?: string;
  contentType?: string;
}) {
  const { baseURL, mode, suffix, method, apiKey, body, contentType } = params;
  const endpoints = buildEndpointCandidates(baseURL, mode, suffix);
  let lastPayload: any = null;
  let lastStatus: number | null = null;
  let lastEndpoint = endpoints[0];

  for (const endpoint of endpoints) {
    lastEndpoint = endpoint;
    const response = await fetch(endpoint, {
      method,
      headers: {
        Authorization: `Bearer ${apiKey}`,
        ...(contentType ? { "Content-Type": contentType } : {}),
      },
      ...(body ? { body } : {}),
      signal: AbortSignal.timeout(45000),
    });
    const payload = await parseJsonResponse(response);
    lastPayload = payload;
    lastStatus = response.status;

    if (response.ok || response.status !== 404) {
      return { endpoint, response, payload };
    }
  }

  throw new Error(
    `Midjourney request failed: ${extractErrorMessage(lastPayload)} (${lastStatus || "no_status"}) @ ${lastEndpoint}`
  );
}

async function submitImagineTask(params: {
  baseURL: string;
  apiKey: string;
  mode: "fast" | "relax";
  prompt: string;
  imageUrls: string[];
}) {
  const { baseURL, apiKey, mode, prompt, imageUrls } = params;
  const base64Array = await Promise.all(imageUrls.map((imageUrl) => loadInputImageAsDataUrl(imageUrl)));

  const { payload } = await requestWithEndpointFallback({
    baseURL,
    mode,
    suffix: "/mj/submit/imagine",
    method: "POST",
    apiKey,
    body: JSON.stringify({
      base64Array,
      prompt,
    }),
    contentType: "application/json",
  });

  const code = Number((payload as any)?.code ?? 0);
  const taskId = String((payload as any)?.result || "").trim();
  if (![1, 22].includes(code) || !taskId) {
    throw new Error(extractErrorMessage(payload));
  }

  return taskId;
}

async function fetchMidjourneyTask(params: {
  baseURL: string;
  apiKey: string;
  mode: "fast" | "relax";
  taskId: string;
}) {
  const { payload } = await requestWithEndpointFallback({
    baseURL: params.baseURL,
    mode: params.mode,
    suffix: `/mj/task/${params.taskId}/fetch`,
    method: "GET",
    apiKey: params.apiKey,
  });

  return (payload || {}) as MidjourneyTaskResult;
}

async function waitForMidjourneyTask(params: {
  baseURL: string;
  apiKey: string;
  mode: "fast" | "relax";
  taskId: string;
}) {
  const deadlineAt = Date.now() + DEFAULT_TIMEOUT_MS;
  let lastTask: MidjourneyTaskResult | null = null;

  while (Date.now() < deadlineAt) {
    const task = await fetchMidjourneyTask(params);
    lastTask = task;

    const status = String(task.status || "").toUpperCase();
    if (status === "SUCCESS" && task.imageUrl) {
      return task;
    }
    if (status === "FAILURE") {
      throw new Error(task.failReason || task.description || "Midjourney 任务失败");
    }

    await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL_MS));
  }

  throw new Error(
    `Midjourney task timeout: ${params.taskId} (${lastTask?.status || "unknown"})`
  );
}

function formatMarkdown(image: UploadedImage, imageSize: string) {
  return `![thumb|||${image.original}|||${imageSize}](${image.thumbnail})`;
}

export function createMidjourney({ apiKey, baseURL }: MidjourneyConfig) {
  return {
    provider: "midjourney",
    chat: (modelId: string) => ({
      provider: "midjourney",
      doGenerate: async (options: {
        messages: Array<{ role: string; content: any }>;
        imageParams?: MidjourneyImageParams;
        overrideModelId?: string;
      }) => {
        const { prompt, imageUrls } = extractPromptAndImages(options.messages);
        const activeModelId = options.overrideModelId || modelId;
        const mode = resolveMode(baseURL, activeModelId);
        const imageSize = normalizeImageSizeLabel(options.imageParams?.imageSize);
        const finalPrompt = appendAspectRatio(prompt, options.imageParams?.aspectRatio);

        console.log("[Midjourney] Submitting request", {
          model: activeModelId,
          mode,
          imageCount: imageUrls.length,
          promptPreview: finalPrompt.slice(0, 120),
        });

        const taskId = await submitImagineTask({
          baseURL,
          apiKey,
          mode,
          prompt: finalPrompt,
          imageUrls,
        });

        console.log("[Midjourney] Task submitted", { taskId, mode });
        const taskResult = await waitForMidjourneyTask({
          baseURL,
          apiKey,
          mode,
          taskId,
        });

        if (!taskResult.imageUrl) {
          throw new Error(taskResult.failReason || taskResult.description || "Midjourney 未返回图片地址");
        }

        const uploadedImage = await downloadAndUploadImage(taskResult.imageUrl);
        const text = formatMarkdown(uploadedImage, imageSize);

        return {
          text,
          finishReason: "stop" as const,
          usage: {
            promptTokens: 0,
            completionTokens: text.length,
          },
          providerMetadata: {
            midjourney: {
              taskId,
              mode,
              progress: taskResult.progress || "",
              buttons: taskResult.buttons || [],
              imageUrl: taskResult.imageUrl,
            },
          },
        };
      },
      doStream: async function* (options: {
        messages: Array<{ role: string; content: any }>;
        imageParams?: MidjourneyImageParams;
        overrideModelId?: string;
      }) {
        const result = await this.doGenerate(options);
        yield {
          type: "text-delta",
          textDelta: result.text,
        };
      },
    }),
  };
}