import "server-only";

import { createHash } from "node:crypto";
import sharp from "sharp";
import { uploadFile } from "@/lib/storage/service";

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

interface FalImageParams {
  aspectRatio?: string;
  imageSize?: string;
  imageCount?: number;
  outputFormat?: "png" | "jpg" | "jpeg" | "webp";
  safetyTolerance?: string | number;
  enableWebSearch?: boolean;
}

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

const DEFAULT_ASPECT_RATIO = "auto";
const DEFAULT_RESOLUTION = "1K";
const DEFAULT_OUTPUT_FORMAT = "png";
const DEFAULT_SAFETY_TOLERANCE = "4";
const DEFAULT_TIMEOUT_MS = 180000;

function normalizeFalCredentials(apiKey: string) {
  const normalized = apiKey
    .replace(/\u00a0/g, " ")
    .replace(/\u3000/g, " ")
    .replace(/：/g, ":")
    .trim();
  const stripped = normalized.replace(/^bearer\s+/i, "").replace(/^key\s+/i, "");
  const kvMatch = stripped.match(/^(?:fal[_-]?key|api[_-]?key|token)\s*[:：]\s*(.+)$/i);
  if (kvMatch?.[1]) {
    return kvMatch[1].trim();
  }
  return stripped;
}

function normalizeUrlForCompare(url: string) {
  try {
    const parsed = new URL(url);
    parsed.search = "";
    parsed.hash = "";
    return parsed.toString();
  } catch {
    return url.split("?")[0]?.split("#")[0] || url;
  }
}

function filterSourceImages(urls: string[], sourceUrls: string[]) {
  if (sourceUrls.length === 0) return urls;
  const sourceSet = new Set(sourceUrls.map((url) => normalizeUrlForCompare(url)));
  return urls.filter((url) => !sourceSet.has(normalizeUrlForCompare(url)));
}

async function getImageContentHash(url: string): Promise<string | null> {
  try {
    const response = await fetch(url, { signal: AbortSignal.timeout(15_000) });
    if (!response.ok) return null;
    const buffer = Buffer.from(await response.arrayBuffer());
    return createHash("sha1").update(buffer).digest("hex");
  } catch {
    return null;
  }
}

async function filterSourceImagesByContent(urls: string[], sourceUrls: string[]) {
  if (sourceUrls.length === 0 || urls.length === 0) return urls;

  const sourceHashes = new Set<string>();
  for (const sourceUrl of sourceUrls) {
    const hash = await getImageContentHash(sourceUrl);
    if (hash) sourceHashes.add(hash);
  }
  if (sourceHashes.size === 0) return urls;

  const filtered: string[] = [];
  for (const url of urls) {
    const hash = await getImageContentHash(url);
    if (hash && sourceHashes.has(hash)) {
      continue;
    }
    filtered.push(url);
  }

  return filtered.length > 0 ? filtered : urls;
}

function normalizeAspectRatio(aspectRatio?: string) {
  if (!aspectRatio) return DEFAULT_ASPECT_RATIO;
  const ratio = aspectRatio.trim();
  const normalized = ratio.toLowerCase();

  if (
    [
      "auto",
      "21:9",
      "16:9",
      "3:2",
      "4:3",
      "5:4",
      "1:1",
      "4:5",
      "3:4",
      "2:3",
      "9:16",
    ].includes(normalized)
  ) {
    return normalized === "auto" ? "auto" : ratio;
  }

  return DEFAULT_ASPECT_RATIO;
}

function normalizeResolution(imageSize?: string) {
  if (!imageSize) return DEFAULT_RESOLUTION;
  const normalized = imageSize.trim().toUpperCase();
  if (["1K", "2K", "4K"].includes(normalized)) return normalized;
  return DEFAULT_RESOLUTION;
}

function normalizeOutputFormat(outputFormat?: string) {
  if (!outputFormat) return DEFAULT_OUTPUT_FORMAT;
  const normalized = outputFormat.trim().toLowerCase();
  if (normalized === "jpg" || normalized === "jpeg") return "jpeg";
  if (normalized === "webp") return "webp";
  return "png";
}

function normalizeSafetyTolerance(value?: string | number) {
  if (value === undefined || value === null) return DEFAULT_SAFETY_TOLERANCE;
  const normalized = String(value).trim();
  if (["1", "2", "3", "4", "5", "6"].includes(normalized)) return normalized;
  return DEFAULT_SAFETY_TOLERANCE;
}

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

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

  return "png";
}

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

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

async function downloadAndUploadImage(imageUrl: string): Promise<UploadedImage> {
  const response = await fetch(imageUrl);
  if (!response.ok) {
    throw new Error(`Failed to download image: ${response.status}`);
  }

  const arrayBuffer = await response.arrayBuffer();
  const buffer = Buffer.from(arrayBuffer);
  const thumbnailBuffer = await createThumbnail(buffer);

  const originalExt = detectImageExtension(imageUrl, response.headers.get("content-type"));
  const originalFilename = generateFilename("fal-image-orig", originalExt);
  const originalArrayBuffer = buffer.buffer.slice(
    buffer.byteOffset,
    buffer.byteOffset + buffer.byteLength
  ) as ArrayBuffer;
  const originalUploaded = await uploadFile(originalFilename, originalArrayBuffer);

  const thumbFilename = generateFilename("fal-image-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,
  };
}

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 = (url?: string) => {
    if (!url) return;
    const normalized = String(url).trim();

    if (normalized.startsWith("data:")) {
      throw new Error("FAL 不支持 Base64 图片传输，请先上传图片后再编辑");
    }

    if (!normalized.startsWith("http")) 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 };
}

function collectImageUrls(payload: any): string[] {
  const urls = new Set<string>();

  const isLikelyImageUrl = (value: string) =>
    /^https?:\/\//i.test(value) &&
    (/\.(png|jpg|jpeg|webp|gif)(\?|$)/i.test(value) ||
      value.includes("image") ||
      value.includes("fal.media"));

  const walk = (value: any, parentKey = "") => {
    if (typeof value === "string") {
      if (isLikelyImageUrl(value)) {
        urls.add(value);
      }
      return;
    }

    if (Array.isArray(value)) {
      for (const item of value) {
        walk(item, parentKey);
      }
      return;
    }

    if (value && typeof value === "object") {
      for (const [key, nestedValue] of Object.entries(value)) {
        if (key === "url" && typeof nestedValue === "string" && /^https?:\/\//i.test(nestedValue)) {
          if (/image|img|file|output|result|artifact/i.test(parentKey) || isLikelyImageUrl(nestedValue)) {
            urls.add(nestedValue);
          }
        }
        walk(nestedValue, key);
      }
    }
  };

  walk(payload);
  return Array.from(urls);
}

function formatMarkdown(images: UploadedImage[], resolution: string) {
  return images
    .map((image) => `![thumb|||${image.original}|||${resolution}](${image.thumbnail})`)
    .join("\n\n");
}

export function createFalImage({ apiKey, baseURL }: FalImageConfig) {
  return {
    provider: "fal-image",
    chat: (modelId: string) => ({
      provider: "fal-image",
      doGenerate: async (options: {
        messages: Array<{ role: string; content: any }>;
        imageParams?: FalImageParams;
        overrideModelId?: string;
      }) => {
        const { prompt, imageUrls } = extractPromptAndImages(options.messages);
        const aspectRatio = normalizeAspectRatio(options.imageParams?.aspectRatio);
        const resolution = normalizeResolution(options.imageParams?.imageSize);
        const outputFormat = normalizeOutputFormat(options.imageParams?.outputFormat);
        const safetyTolerance = normalizeSafetyTolerance(options.imageParams?.safetyTolerance);
        const imageCount = Math.max(1, Math.min(6, Number(options.imageParams?.imageCount || 1)));
        const activeModelId = options.overrideModelId || modelId;

        console.log("[Fal Image] Submitting request", {
          model: activeModelId,
          baseURL,
          imageCount,
          aspectRatio,
          resolution,
        });

        const fal = await import("@fal-ai/serverless-client");
        fal.config({
          credentials: normalizeFalCredentials(apiKey),
        });

        const input: Record<string, any> = {
          prompt,
          num_images: imageCount,
          aspect_ratio: aspectRatio,
          resolution,
          output_format: outputFormat,
          safety_tolerance: safetyTolerance,
        };

        if (imageUrls.length > 0) {
          input.image_urls = imageUrls;
        }

        if (typeof options.imageParams?.enableWebSearch === "boolean") {
          input.enable_web_search = options.imageParams.enableWebSearch;
        }

        const result: any = await fal.subscribe(activeModelId, {
          input,
          logs: true,
          timeout: DEFAULT_TIMEOUT_MS,
          onQueueUpdate: (update: any) => {
            if (update?.status === "IN_PROGRESS" && Array.isArray(update.logs)) {
              for (const log of update.logs) {
                if (log?.message) {
                  console.log(`[Fal Image][${activeModelId}]`, log.message);
                }
              }
            }
          },
        });

        let outputUrls = filterSourceImages(collectImageUrls(result), imageUrls);
        outputUrls = await filterSourceImagesByContent(outputUrls, imageUrls);

        if (outputUrls.length === 0 && result?.data) {
          outputUrls = filterSourceImages(collectImageUrls(result.data), imageUrls);
          outputUrls = await filterSourceImagesByContent(outputUrls, imageUrls);
        }

        if (outputUrls.length === 0) {
          throw new Error("FAL 未返回可用图片");
        }

        const uploadedImages: UploadedImage[] = [];
        for (const outputUrl of outputUrls.slice(0, imageCount)) {
          uploadedImages.push(await downloadAndUploadImage(outputUrl));
        }

        if (uploadedImages.length === 0) {
          throw new Error("FAL 未返回可用图片");
        }

        const text = formatMarkdown(uploadedImages, resolution);
        return {
          text,
          finishReason: "stop" as const,
          usage: {
            promptTokens: 0,
            completionTokens: text.length,
          },
        };
      },
      doStream: async function* (options: {
        messages: Array<{ role: string; content: any }>;
        imageParams?: FalImageParams;
        overrideModelId?: string;
      }) {
        const result = await this.doGenerate(options);
        yield {
          type: "text-delta",
          textDelta: result.text,
        };
      },
    }),
  };
}