import "server-only";

import JSZip from "jszip";
import { createMediaAssetRelation, getMediaAssetsByUserAndUrls, getSystemSetting, upsertMediaAsset } from "@/lib/db/queries";
import { db } from "@/lib/db";
import { generatedImage } from "@/lib/db/schema";
import { uploadFile } from "@/lib/storage/service";

const LAYER_SPLIT_CONFIG_KEY = "layer_split_config";
const DEFAULT_REPLICATE_BASE_URL = "https://api.replicate.com/v1";
const DEFAULT_LAYER_SPLIT_MODEL = "qwen/qwen-image-layered";
const DEFAULT_LAYERS = 4;
const DEFAULT_RESOLUTION = 640;
const DEFAULT_TRUE_CFG_SCALE = 4;
const DEFAULT_INFERENCE_STEPS = 30;
const DEFAULT_TIMEOUT_MS = 5 * 60 * 1000;

type LayerSplitConfig = {
  enabled?: boolean;
  provider?: "replicate";
  replicateApiToken?: string;
  replicateModel?: string;
  layers?: number;
  resolution?: number;
  trueCfgScale?: number;
  numInferenceSteps?: number;
  cfgNormalize?: boolean;
  useEnPrompt?: boolean;
  negativePrompt?: string;
  timeoutMs?: number;
};

type LayerSplitParams = {
  userId: string;
  imageUrl: string;
  prompt?: string;
  chatId?: string | null;
  messageId?: string | null;
  sourceNodeId?: string | null;
  skipGallerySave?: boolean;
};

type ReplicatePredictionResponse = {
  id?: string;
  status?: string;
  output?: unknown;
  error?: string;
  detail?: string;
  message?: string;
  urls?: {
    get?: string;
  };
};

export type LayerSplitResultItem = {
  originalUrl: string;
  thumbnailUrl: string;
  generatedImageId?: string | null;
  mediaAssetId?: string | null;
  title: string;
  layerIndex: number;
  totalLayers: number;
};

export type ExecuteLayerSplitJobResult = {
  layers: LayerSplitResultItem[];
  provider: "replicate";
  model: string;
  processingTime: number;
};

export class LayerSplitServiceError extends Error {
  status: number;
  detail?: string;

  constructor(message: string, status = 500, detail?: string) {
    super(message);
    this.name = "LayerSplitServiceError";
    this.status = status;
    this.detail = detail;
  }
}

function trimString(value: unknown) {
  return String(value || "").trim();
}

function splitReplicateModelIdentifier(modelId: string) {
  const [owner, ...rest] = modelId.split("/").filter(Boolean);
  const name = rest.join("/");
  if (!owner || !name) {
    throw new LayerSplitServiceError(`无效的 Replicate 模型标识：${modelId}`, 400);
  }
  return { owner, name };
}

function normalizeLayerSplitConfig(value: unknown): Required<LayerSplitConfig> {
  const raw =
    value && typeof value === "object" && !Array.isArray(value)
      ? (value as Record<string, unknown>)
      : {};

  return {
    enabled: raw.enabled === true,
    provider: "replicate",
    replicateApiToken: trimString(raw.replicateApiToken),
    replicateModel: trimString(raw.replicateModel) || DEFAULT_LAYER_SPLIT_MODEL,
    layers: Math.max(2, Math.min(12, Math.round(Number(raw.layers) || DEFAULT_LAYERS))),
    resolution: Math.max(256, Math.min(1536, Math.round(Number(raw.resolution) || DEFAULT_RESOLUTION))),
    trueCfgScale: Math.max(1, Math.min(12, Number(raw.trueCfgScale) || DEFAULT_TRUE_CFG_SCALE)),
    numInferenceSteps: Math.max(
      1,
      Math.min(80, Math.round(Number(raw.numInferenceSteps) || DEFAULT_INFERENCE_STEPS))
    ),
    cfgNormalize: raw.cfgNormalize !== false,
    useEnPrompt: raw.useEnPrompt !== false,
    negativePrompt: trimString(raw.negativePrompt),
    timeoutMs: Math.max(30_000, Math.min(15 * 60_000, Math.round(Number(raw.timeoutMs) || DEFAULT_TIMEOUT_MS))),
  };
}

async function getValidatedLayerSplitConfig() {
  const setting = await getSystemSetting({ key: LAYER_SPLIT_CONFIG_KEY });
  const config = normalizeLayerSplitConfig(setting?.value);

  if (!config.enabled) {
    throw new LayerSplitServiceError("图层拆分功能未启用，请先在本地AI设置中开启", 400);
  }

  if (!config.replicateApiToken) {
    throw new LayerSplitServiceError("未配置 Replicate API Token，请先在本地AI设置中保存", 400);
  }

  return config;
}

function toReplicateErrorMessage(payload: any, fallback: string) {
  const message = trimString(payload?.detail || payload?.error || payload?.message);
  return message || fallback;
}

function shouldRetryWithAlternativeInput(errorMessage: string) {
  const normalized = errorMessage.toLowerCase();
  return (
    normalized.includes("validation") ||
    normalized.includes("schema") ||
    normalized.includes("image") ||
    normalized.includes("input") ||
    normalized.includes("field")
  );
}

async function createPrediction(params: {
  apiKey: string;
  modelId: string;
  input: Record<string, unknown>;
}) {
  const { owner, name } = splitReplicateModelIdentifier(params.modelId);
  const endpoint = `${DEFAULT_REPLICATE_BASE_URL}/models/${encodeURIComponent(owner)}/${encodeURIComponent(name)}/predictions`;
  const response = await fetch(endpoint, {
    method: "POST",
    headers: {
      Authorization: `Bearer ${params.apiKey}`,
      "Content-Type": "application/json",
      Prefer: "wait=5",
    },
    body: JSON.stringify({ input: params.input }),
    cache: "no-store",
  });

  const payload = (await response.json().catch(() => null)) as ReplicatePredictionResponse | null;
  return { response, payload };
}

async function pollPrediction(params: {
  apiKey: string;
  taskId?: string;
  pollingUrl?: string;
  timeoutMs: number;
}): Promise<ReplicatePredictionResponse> {
  const deadline = Date.now() + params.timeoutMs;
  let lastMessage = "";

  while (Date.now() < deadline) {
    const taskId = trimString(params.taskId);
    const url =
      trimString(params.pollingUrl) ||
      `${DEFAULT_REPLICATE_BASE_URL}/predictions/${encodeURIComponent(taskId)}`;

    const response = await fetch(url, {
      method: "GET",
      headers: {
        Authorization: `Bearer ${params.apiKey}`,
        Accept: "application/json",
      },
      cache: "no-store",
    });
    const payload = (await response.json().catch(() => null)) as ReplicatePredictionResponse | null;
    const status = trimString(payload?.status).toLowerCase();

    if (status) {
      lastMessage = status;
    }

    if (response.ok && status === "succeeded") {
      return payload || {};
    }

    if (status === "failed" || status === "canceled") {
      throw new LayerSplitServiceError(
        "图层拆分任务失败",
        502,
        toReplicateErrorMessage(payload, "Replicate 任务执行失败")
      );
    }

    if (!response.ok) {
      throw new LayerSplitServiceError(
        `图层拆分轮询失败（${response.status}）`,
        502,
        toReplicateErrorMessage(payload, `Replicate 轮询失败（${response.status}）`)
      );
    }

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

  throw new LayerSplitServiceError(
    "图层拆分任务超时",
    504,
    `Replicate 轮询超时${lastMessage ? `（${lastMessage}）` : ""}`
  );
}

function collectArtifactStrings(value: unknown, bucket: string[] = []): string[] {
  if (typeof value === "string") {
    const next = trimString(value);
    if (next) bucket.push(next);
    return bucket;
  }

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

  if (value && typeof value === "object") {
    for (const [key, item] of Object.entries(value as Record<string, unknown>)) {
      if (
        key === "url" ||
        key === "href" ||
        key === "download" ||
        key === "file" ||
        key === "image" ||
        key === "get" ||
        key === "output" ||
        key === "outputs" ||
        key === "layers" ||
        key === "files" ||
        key === "result"
      ) {
        collectArtifactStrings(item, bucket);
      }
    }
  }

  return bucket;
}

function dedupeStrings(values: string[]) {
  return Array.from(new Set(values.map((value) => trimString(value)).filter(Boolean)));
}

function extFromFilename(filename: string): string {
  const clean = trimString(filename).split("?")[0];
  const ext = clean.split(".").pop()?.toLowerCase() || "";
  return ext;
}

function deriveFilenameFromUrl(url: string, fallback: string) {
  try {
    const parsed = new URL(url);
    const candidate = parsed.pathname.split("/").pop() || "";
    return candidate || fallback;
  } catch {
    return fallback;
  }
}

function detectMimeType(buffer: Buffer, fallback?: string): string {
  if (buffer.length >= 8 && buffer.subarray(0, 8).equals(Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]))) {
    return "image/png";
  }
  if (buffer.length >= 3 && buffer[0] === 0xff && buffer[1] === 0xd8 && buffer[2] === 0xff) {
    return "image/jpeg";
  }
  if (buffer.length >= 12 && buffer.toString("ascii", 0, 4) === "RIFF" && buffer.toString("ascii", 8, 12) === "WEBP") {
    return "image/webp";
  }
  if (buffer.length >= 4 && buffer.toString("ascii", 0, 4) === "PK\u0003\u0004") {
    return "application/zip";
  }
  if (buffer.length >= 4 && buffer.toString("ascii", 0, 4) === "8BPS") {
    return "image/vnd.adobe.photoshop";
  }

  const normalizedFallback = trimString(fallback).toLowerCase();
  if (normalizedFallback.startsWith("image/")) return normalizedFallback;
  if (normalizedFallback === "application/zip") return normalizedFallback;
  return "application/octet-stream";
}

function extFromMimeType(mimeType: string) {
  if (mimeType === "image/png") return "png";
  if (mimeType === "image/jpeg") return "jpg";
  if (mimeType === "image/webp") return "webp";
  if (mimeType === "application/zip") return "zip";
  if (mimeType === "image/vnd.adobe.photoshop") return "psd";
  return "bin";
}

async function fetchArtifact(url: string): Promise<{ buffer: Buffer; mimeType: string; filename: string }> {
  const response = await fetch(url, {
    headers: {
      "User-Agent": "AI-Chatbot-LayerSplit/1.0",
    },
    cache: "no-store",
  });

  if (!response.ok) {
    throw new LayerSplitServiceError(
      `下载图层拆分结果失败（${response.status}）`,
      502,
      `${response.status} ${response.statusText}`
    );
  }

  const contentType = trimString(response.headers.get("content-type"));
  const buffer = Buffer.from(await response.arrayBuffer());
  const filename = deriveFilenameFromUrl(url, `artifact.${extFromMimeType(contentType || "application/octet-stream")}`);
  const mimeType = detectMimeType(buffer, contentType);
  return { buffer, mimeType, filename };
}

async function extractLayerBuffersFromZip(zipBuffer: Buffer): Promise<Array<{ buffer: Buffer; filename: string; mimeType: string }>> {
  const zip = await JSZip.loadAsync(zipBuffer);
  const entries: Array<{ buffer: Buffer; filename: string; mimeType: string }> = [];

  for (const [name, file] of Object.entries(zip.files)) {
    if (file.dir) continue;
    const ext = extFromFilename(name);
    if (!["png", "jpg", "jpeg", "webp"].includes(ext)) continue;
    const fileBuffer = await file.async("nodebuffer");
    const mimeType = detectMimeType(fileBuffer, ext === "png" ? "image/png" : ext === "webp" ? "image/webp" : "image/jpeg");
    if (!mimeType.startsWith("image/")) continue;
    entries.push({
      buffer: fileBuffer,
      filename: name.split("/").pop() || `layer-${entries.length + 1}.${ext || "png"}`,
      mimeType,
    });
  }

  return entries;
}

async function resolveLayerArtifacts(payload: ReplicatePredictionResponse): Promise<Array<{ buffer: Buffer; filename: string; mimeType: string }>> {
  const artifacts = dedupeStrings(collectArtifactStrings(payload?.output ?? payload));

  if (artifacts.length === 0) {
    throw new LayerSplitServiceError("图层拆分未返回任何可用结果", 502);
  }

  const imageEntries: Array<{ buffer: Buffer; filename: string; mimeType: string }> = [];

  for (const artifact of artifacts) {
    if (artifact.startsWith("data:")) {
      const match = artifact.match(/^data:([^;,]+)?;base64,(.+)$/i);
      if (!match) continue;
      const mimeType = detectMimeType(Buffer.from(match[2], "base64"), match[1] || "image/png");
      const buffer = Buffer.from(match[2], "base64");
      if (mimeType.startsWith("image/")) {
        imageEntries.push({
          buffer,
          mimeType,
          filename: `layer-${imageEntries.length + 1}.${extFromMimeType(mimeType)}`,
        });
        continue;
      }
      if (mimeType === "application/zip") {
        const zipped = await extractLayerBuffersFromZip(buffer);
        imageEntries.push(...zipped);
      }
      continue;
    }

    if (!/^https?:\/\//i.test(artifact)) continue;

    const downloaded = await fetchArtifact(artifact);
    if (downloaded.mimeType.startsWith("image/")) {
      imageEntries.push(downloaded);
      continue;
    }

    if (downloaded.mimeType === "application/zip") {
      const zipped = await extractLayerBuffersFromZip(downloaded.buffer);
      imageEntries.push(...zipped);
      continue;
    }
  }

  if (imageEntries.length === 0) {
    throw new LayerSplitServiceError("图层拆分返回了结果文件，但未找到可回画布的图片图层", 502);
  }

  return imageEntries;
}

async function createThumbnail(buffer: Buffer, mimeType: string): Promise<Buffer> {
  try {
    const sharp = await import("sharp");
    const pipeline = sharp.default(buffer).resize(500, null, {
      fit: "inside",
      withoutEnlargement: true,
    });

    if (mimeType === "image/jpeg") {
      return await pipeline.jpeg({ quality: 82 }).toBuffer();
    }
    if (mimeType === "image/webp") {
      return await pipeline.webp({ quality: 82 }).toBuffer();
    }
    return await pipeline.png({ compressionLevel: 8 }).toBuffer();
  } catch (error) {
    console.error("[LayerSplit] Failed to create thumbnail:", error);
    return buffer;
  }
}

async function uploadBufferPair(buffer: Buffer, mimeType: string, prefix: string) {
  const ext = extFromMimeType(mimeType);
  const timestamp = Date.now();
  const random = Math.random().toString(36).slice(2, 8);
  const thumbFilename = `${prefix}-thumb-${timestamp}-${random}.${ext}`;
  const origFilename = `${prefix}-orig-${timestamp}-${random}.${ext}`;
  const thumbnailBuffer = await createThumbnail(buffer, mimeType);

  const thumbnailArrayBuffer = thumbnailBuffer.buffer.slice(
    thumbnailBuffer.byteOffset,
    thumbnailBuffer.byteOffset + thumbnailBuffer.byteLength
  ) as ArrayBuffer;
  const originalArrayBuffer = buffer.buffer.slice(
    buffer.byteOffset,
    buffer.byteOffset + buffer.byteLength
  ) as ArrayBuffer;

  const [thumbnailResult, originalResult] = await Promise.all([
    uploadFile(thumbFilename, thumbnailArrayBuffer),
    uploadFile(origFilename, originalArrayBuffer),
  ]);

  return {
    thumbnail: thumbnailResult.url,
    original: originalResult.url,
  };
}

async function runReplicateLayerSplit(params: {
  apiKey: string;
  modelId: string;
  imageUrl: string;
  config: Required<LayerSplitConfig>;
  prompt?: string;
}): Promise<ReplicatePredictionResponse> {
  const baseInput: Record<string, unknown> = {
    max_num_layers: params.config.layers,
    resolution: params.config.resolution,
    true_cfg_scale: params.config.trueCfgScale,
    num_inference_steps: params.config.numInferenceSteps,
    cfg_normalize: params.config.cfgNormalize,
    use_en_prompt: params.config.useEnPrompt,
  };

  const negativePrompt = trimString(params.config.negativePrompt);
  if (negativePrompt) {
    baseInput.negative_prompt = negativePrompt;
  }

  const prompt = trimString(params.prompt);
  if (prompt) {
    baseInput.prompt = prompt;
  }

  const candidateInputs = [
    { ...baseInput, image: params.imageUrl },
    { ...baseInput, input_image: params.imageUrl },
    { ...baseInput, images: [params.imageUrl] },
  ];

  let lastPayload: ReplicatePredictionResponse | null = null;
  let lastError = "";

  for (let index = 0; index < candidateInputs.length; index += 1) {
    const { response, payload } = await createPrediction({
      apiKey: params.apiKey,
      modelId: params.modelId,
      input: candidateInputs[index],
    });

    lastPayload = payload;

    if (
      response.ok &&
      (trimString(payload?.status) === "succeeded" || trimString(payload?.id) || trimString(payload?.urls?.get))
    ) {
      if (trimString(payload?.status).toLowerCase() === "succeeded") {
        return payload || {};
      }
      return await pollPrediction({
        apiKey: params.apiKey,
        taskId: trimString(payload?.id),
        pollingUrl: trimString(payload?.urls?.get),
        timeoutMs: params.config.timeoutMs,
      });
    }

    const errorMessage = toReplicateErrorMessage(payload, `Replicate 请求失败（${response.status}）`);
    lastError = errorMessage;

    if (response.status === 401 || response.status === 403) {
      throw new LayerSplitServiceError("Replicate API Token 无效或无权限", 401, errorMessage);
    }

    if (response.status === 404) {
      throw new LayerSplitServiceError(`Replicate 模型不存在：${params.modelId}`, 400, errorMessage);
    }

    if (index < candidateInputs.length - 1 && shouldRetryWithAlternativeInput(errorMessage)) {
      continue;
    }
  }

  throw new LayerSplitServiceError(
    "图层拆分请求失败",
    502,
    lastError || toReplicateErrorMessage(lastPayload || {}, "Replicate 请求失败")
  );
}

export async function executeLayerSplitJob(params: LayerSplitParams): Promise<ExecuteLayerSplitJobResult> {
  const startedAt = Date.now();
  const config = await getValidatedLayerSplitConfig();

  const payload = await runReplicateLayerSplit({
    apiKey: config.replicateApiToken,
    modelId: config.replicateModel,
    imageUrl: params.imageUrl,
    config,
    prompt: params.prompt,
  });

  const layerBuffers = await resolveLayerArtifacts(payload);
  const totalLayers = layerBuffers.length;
  const sourceAssets = await getMediaAssetsByUserAndUrls({
    userId: params.userId,
    urls: [params.imageUrl],
  }).catch(() => []);

  const results: LayerSplitResultItem[] = [];

  for (let index = 0; index < layerBuffers.length; index += 1) {
    const layer = layerBuffers[index];
    const uploadResult = await uploadBufferPair(layer.buffer, layer.mimeType, `layer-split-${index + 1}`);
    let generatedImageId: string | null = null;
    let mediaAssetId: string | null = null;

    if (!params.skipGallerySave) {
      const [savedImage] = await db
        .insert(generatedImage)
        .values({
          userId: params.userId,
          chatId: params.chatId || null,
          messageId: params.messageId || null,
          thumbnailUrl: uploadResult.thumbnail,
          originalUrl: uploadResult.original,
          prompt: params.prompt || `图层拆分 ${index + 1}/${totalLayers}`,
          modelId: config.replicateModel,
          createdAt: new Date(),
        })
        .returning({ id: generatedImage.id });
      generatedImageId = savedImage?.id || null;

      const media = await upsertMediaAsset({
        userId: params.userId,
        chatId: params.chatId || null,
        mediaType: "image",
        sourceType: "edited",
        status: "ready",
        originalUrl: uploadResult.original,
        previewUrl: uploadResult.original,
        thumbnailUrl: uploadResult.thumbnail,
        displayName: `图层拆分 ${index + 1}/${totalLayers}`.slice(0, 80),
        prompt: params.prompt || `图层拆分 ${index + 1}/${totalLayers}`,
        modelId: config.replicateModel,
        sourceMessageId: params.messageId || null,
        sourceTable: "GeneratedImage",
        sourceRowId: generatedImageId,
        sourceNodeId: params.sourceNodeId || null,
        primaryParentAssetId: sourceAssets[0]?.id ?? null,
        metadata: {
          layerIndex: index + 1,
          totalLayers,
          provider: "replicate",
          model: config.replicateModel,
          sourceImageUrl: params.imageUrl,
          sourceFilename: layer.filename,
        },
      });
      mediaAssetId = media?.id || null;

      if (media && sourceAssets.length > 0) {
        for (const sourceAsset of sourceAssets) {
          await createMediaAssetRelation({
            userId: params.userId,
            fromAssetId: sourceAsset.id,
            toAssetId: media.id,
            relationType: "edited_from",
            chatId: params.chatId || null,
            sourceMessageId: params.messageId || null,
            sourceNodeId: params.sourceNodeId || null,
            metadata: {
              tool: "layer_split",
              layerIndex: index + 1,
              totalLayers,
            },
          }).catch(() => null);
        }
      }
    }

    results.push({
      originalUrl: uploadResult.original,
      thumbnailUrl: uploadResult.thumbnail,
      generatedImageId,
      mediaAssetId,
      title: `图层${index + 1}`,
      layerIndex: index + 1,
      totalLayers,
    });
  }

  return {
    layers: results,
    provider: "replicate",
    model: config.replicateModel,
    processingTime: (Date.now() - startedAt) / 1000,
  };
}