import fs from 'node:fs';
import { readFile } from 'node:fs/promises';
import path from 'node:path';

import { z } from 'zod';

import type { AppContext } from '../context.js';
import { badRequest } from '../errors.js';
import { extractApiErrorMessage } from '../http.js';

export const sceneImageSchema = z.object({
  prompt: z.string().trim().min(1),
  negativePrompt: z.string().trim().optional().default(''),
  size: z.string().trim().optional().default('1280*720'),
  model: z.string().trim().optional().default(''),
  worldName: z.string().trim().optional().default(''),
  profileId: z.string().trim().optional().default(''),
  landmarkName: z.string().trim().optional().default(''),
  landmarkId: z.string().trim().optional().default(''),
  referenceImageSrc: z.string().trim().optional().default(''),
});
const TEXT_TO_IMAGE_SCENE_MODEL = 'wan2.2-t2i-flash';
const REFERENCE_IMAGE_SCENE_MODEL = 'qwen-image-2.0';

function parseImageDataUrl(source: string) {
  const matched = /^data:(image\/[^;]+);base64,(.+)$/u.exec(source);
  if (!matched) {
    return null;
  }

  return {
    buffer: Buffer.from(matched[2], 'base64'),
    mimeType: matched[1],
  };
}

async function resolveReferenceImageAsDataUrl(rootDir: string, source: string) {
  const trimmedSource = source.trim();
  if (!trimmedSource) {
    return '';
  }

  const parsedDataUrl = parseImageDataUrl(trimmedSource);
  if (parsedDataUrl) {
    return trimmedSource;
  }

  if (!trimmedSource.startsWith('/')) {
    throw badRequest('参考图必须是 Data URL 或 public 目录下的 URL。');
  }

  const normalizedSource = path.posix
    .normalize(trimmedSource)
    .replace(/^\/+/u, '');
  const absolutePath = path.resolve(
    rootDir,
    'public',
    ...normalizedSource.split('/'),
  );
  const publicRoot = path.resolve(rootDir, 'public');
  if (!absolutePath.startsWith(publicRoot)) {
    throw badRequest('参考图路径越界。');
  }

  const buffer = await readFile(absolutePath);
  const extension = path.extname(absolutePath).replace(/^\./u, '').toLowerCase();
  const mimeType = (() => {
    switch (extension) {
      case 'jpg':
      case 'jpeg':
        return 'image/jpeg';
      case 'webp':
        return 'image/webp';
      default:
        return 'image/png';
    }
  })();

  return `data:${mimeType};base64,${buffer.toString('base64')}`;
}

function collectStringsByKey(
  value: unknown,
  targetKey: string,
  results: string[],
) {
  if (typeof value === 'string') {
    return;
  }

  if (Array.isArray(value)) {
    value.forEach((entry) => collectStringsByKey(entry, targetKey, results));
    return;
  }

  if (!value || typeof value !== 'object') {
    return;
  }

  Object.entries(value).forEach(([key, nestedValue]) => {
    if (key === targetKey && typeof nestedValue === 'string' && nestedValue.trim()) {
      results.push(nestedValue.trim());
      return;
    }

    collectStringsByKey(nestedValue, targetKey, results);
  });
}

function findFirstStringByKey(value: unknown, targetKey: string) {
  const results: string[] = [];
  collectStringsByKey(value, targetKey, results);
  return results[0] ?? '';
}

function extractTaskId(payload: Record<string, unknown>) {
  return findFirstStringByKey(payload, 'task_id');
}

function extractImageUrls(payload: Record<string, unknown>) {
  const urls: string[] = [];
  collectStringsByKey(payload, 'image', urls);
  collectStringsByKey(payload, 'url', urls);
  return [...new Set(urls)];
}

async function createSceneImageTask(params: {
  baseUrl: string;
  apiKey: string;
  payload: z.infer<typeof sceneImageSchema>;
}) {
  const { baseUrl, apiKey, payload } = params;
  const response = await fetch(`${baseUrl}/services/aigc/image-generation/generation`, {
    method: 'POST',
    headers: {
      Authorization: `Bearer ${apiKey}`,
      'Content-Type': 'application/json',
      'X-DashScope-Async': 'enable',
    },
    body: JSON.stringify({
      model: payload.model,
      input: {
        messages: [
          {
            role: 'user',
            content: [{ text: payload.prompt }],
          },
        ],
      },
      parameters: {
        n: 1,
        size: payload.size,
        prompt_extend: true,
        watermark: false,
        ...(payload.negativePrompt
          ? { negative_prompt: payload.negativePrompt }
          : {}),
      },
    }),
  });
  const responseText = await response.text();

  if (!response.ok) {
    return {
      ok: false as const,
      errorMessage: extractApiErrorMessage(
        responseText,
        '创建场景图片生成任务失败',
      ),
    };
  }

  return {
    ok: true as const,
    payload: JSON.parse(responseText) as Record<string, unknown>,
  };
}

async function createSceneImageFromReference(params: {
  baseUrl: string;
  apiKey: string;
  payload: z.infer<typeof sceneImageSchema>;
  referenceImage: string;
}) {
  const { baseUrl, apiKey, payload, referenceImage } = params;
  const response = await fetch(
    `${baseUrl}/services/aigc/multimodal-generation/generation`,
    {
      method: 'POST',
      headers: {
        Authorization: `Bearer ${apiKey}`,
        'Content-Type': 'application/json',
      },
      body: JSON.stringify({
        model: payload.model,
        input: {
          messages: [
            {
              role: 'user',
              content: [{ image: referenceImage }, { text: payload.prompt }],
            },
          ],
        },
        parameters: {
          n: 1,
          size: payload.size,
          prompt_extend: true,
          watermark: false,
          ...(payload.negativePrompt
            ? { negative_prompt: payload.negativePrompt }
            : {}),
        },
      }),
    },
  );
  const responseText = await response.text();

  if (!response.ok) {
    return {
      ok: false as const,
      errorMessage: extractApiErrorMessage(
        responseText,
        '创建参考图场景编辑任务失败',
      ),
    };
  }

  const responsePayload = JSON.parse(responseText) as Record<string, unknown>;
  const imageUrl = extractImageUrls(responsePayload)[0] ?? '';
  if (!imageUrl) {
    return {
      ok: false as const,
      errorMessage: '参考图场景编辑未返回图片地址',
    };
  }

  return {
    ok: true as const,
    imageUrl,
    actualPrompt: findFirstStringByKey(responsePayload, 'actual_prompt').trim(),
    taskId: `scene-edit-${Date.now()}`,
  };
}

function ensurePayload(
  payload: z.infer<typeof sceneImageSchema>,
  _defaultModel: string,
) {
  if (!payload.landmarkName && !payload.landmarkId) {
    throw badRequest('landmarkName 或 landmarkId 至少要提供一个');
  }

  const referenceImageSrc =
    typeof payload.referenceImageSrc === 'string'
      ? payload.referenceImageSrc.trim()
      : '';

  return {
    ...payload,
    referenceImageSrc,
    model: referenceImageSrc
      ? REFERENCE_IMAGE_SCENE_MODEL
      : TEXT_TO_IMAGE_SCENE_MODEL,
  };
}

async function saveSceneImageAsset(params: {
  context: AppContext;
  payload: z.infer<typeof sceneImageSchema>;
  imageUrl: string;
  taskId: string;
  actualPrompt: string;
}) {
  const { context, payload, imageUrl, taskId, actualPrompt } = params;
  const imageResponse = await fetch(imageUrl);
  if (!imageResponse.ok) {
    throw badRequest('下载生成图片失败');
  }

  const imageBuffer = Buffer.from(await imageResponse.arrayBuffer());
  const contentType = imageResponse.headers.get('content-type') || '';
  const extension = contentType.includes('png')
    ? 'png'
    : contentType.includes('webp')
      ? 'webp'
      : 'jpg';
  const assetId = `custom-scene-${Date.now()}`;
  const worldSegment = (payload.profileId || payload.worldName || 'world')
    .replace(/[^\w\u4e00-\u9fa5-]+/gu, '-')
    .slice(0, 48);
  const landmarkSegment = (payload.landmarkId || payload.landmarkName || 'landmark')
    .replace(/[^\w\u4e00-\u9fa5-]+/gu, '-')
    .slice(0, 48);
  const relativeDir = path.join(
    'generated-custom-world-scenes',
    worldSegment || 'world',
    landmarkSegment || 'landmark',
    assetId,
  );
  const outputDir = path.join(context.config.publicDir, relativeDir);
  fs.mkdirSync(outputDir, { recursive: true });
  const fileName = `scene.${extension}`;
  fs.writeFileSync(path.join(outputDir, fileName), imageBuffer);

  const imageSrc = `/${relativeDir.replace(/\\/gu, '/')}/${fileName}`;
  fs.writeFileSync(
    path.join(outputDir, 'manifest.json'),
    `${JSON.stringify(
      {
        assetId,
        taskId,
        model: payload.model,
        size: payload.size,
        prompt: payload.prompt,
        negativePrompt: payload.negativePrompt,
        actualPrompt,
        imageSrc,
        worldName: payload.worldName,
        landmarkName: payload.landmarkName,
        createdAt: new Date().toISOString(),
      },
      null,
      2,
    )}\n`,
  );

  return {
    ok: true,
    imageSrc,
    assetId,
    taskId,
    model: payload.model,
    size: payload.size,
    prompt: payload.prompt,
    actualPrompt,
  };
}

export async function generateSceneImage(
  context: AppContext,
  input: z.infer<typeof sceneImageSchema>,
) {
  const payload = ensurePayload(input, context.config.dashScope.imageModel);
  const baseUrl = context.config.dashScope.baseUrl.replace(/\/+$/u, '');
  const referenceImage = payload.referenceImageSrc.trim()
    ? await resolveReferenceImageAsDataUrl(
        context.config.projectRoot,
        payload.referenceImageSrc,
      )
    : '';

  if (referenceImage) {
    const referenceResult = await createSceneImageFromReference({
      baseUrl,
      apiKey: context.config.dashScope.apiKey,
      payload,
      referenceImage,
    });

    if (!referenceResult.ok) {
      throw badRequest(referenceResult.errorMessage);
    }

    return saveSceneImageAsset({
      context,
      payload,
      imageUrl: referenceResult.imageUrl,
      taskId: referenceResult.taskId,
      actualPrompt: referenceResult.actualPrompt,
    });
  }

  const createTaskResult = await createSceneImageTask({
    baseUrl,
    apiKey: context.config.dashScope.apiKey,
    payload,
  });

  if (!createTaskResult.ok) {
    throw badRequest(createTaskResult.errorMessage);
  }

  const createPayload = createTaskResult.payload;
  const taskId = extractTaskId(createPayload);
  if (!taskId) {
    throw badRequest('场景图片生成任务未返回 task_id');
  }

  const deadline = Date.now() + context.config.dashScope.requestTimeoutMs;
  let imageUrl = '';
  let actualPrompt = '';

  while (Date.now() < deadline) {
    const pollResponse = await fetch(`${baseUrl}/tasks/${taskId}`, {
      headers: {
        Authorization: `Bearer ${context.config.dashScope.apiKey}`,
      },
    });
    const pollText = await pollResponse.text();
    if (!pollResponse.ok) {
      throw badRequest(
        extractApiErrorMessage(pollText, '查询场景图片任务失败'),
      );
    }

    const pollPayload = JSON.parse(pollText) as Record<string, unknown>;
    const status = findFirstStringByKey(pollPayload, 'task_status').trim();
    if (status === 'SUCCEEDED') {
      imageUrl = extractImageUrls(pollPayload)[0] ?? '';
      actualPrompt = findFirstStringByKey(pollPayload, 'actual_prompt').trim();
      break;
    }
    if (status === 'FAILED' || status === 'UNKNOWN') {
      throw badRequest(
        extractApiErrorMessage(pollText, '场景图片生成任务失败'),
      );
    }

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

  if (!imageUrl) {
    throw badRequest('场景图片生成超时或未返回图片地址');
  }

  return saveSceneImageAsset({
    context,
    payload,
    imageUrl,
    taskId,
    actualPrompt,
  });
}