import { defineNuxtModule, useLogger } from "@nuxt/kit";
import { createHash } from "crypto";
import { readdir, readFile, writeFile } from "fs/promises";
import { join, relative, parse } from "path";
import { existsSync, watch } from "fs";
import { execSync } from "child_process";
import sharp from "sharp";

type AtlasRect = {
  x: number;
  y: number;
  width: number;
  height: number;
};

type ImageTree = {
  [key: string]: AtlasRect | ImageTree;
};

type ModelTree = {
  [key: string]: string | ModelTree;
};

type ImageData = {
  path: string;
  buffer: Buffer;
  width: number;
  height: number;
};

const IMAGE_EXTENSIONS = [".png", ".webp"];
const MODEL_EXTENSIONS = [".gltf"];
const MAX_WIDTH = 2048;

const toCamelCase = (str: string) => {
  const camel = str
    .replace(/[-_](.)/g, (_, c) => c.toUpperCase())
    .replaceAll(".", "_");
  // Prefix with underscore if starts with a digit to avoid octal literals
  return /^\d/.test(camel) ? `_${camel}` : camel;
};

const scanByExt = async (
  dir: string,
  extensions: string[],
): Promise<string[]> => {
  if (!existsSync(dir)) return [];
  const entries = await readdir(dir, { withFileTypes: true });
  const results: string[] = [];

  for (const entry of entries) {
    const fullPath = join(dir, entry.name);
    if (entry.isDirectory()) {
      results.push(...(await scanByExt(fullPath, extensions)));
    } else if (extensions.includes(parse(fullPath).ext.toLowerCase())) {
      results.push(fullPath);
    }
  }

  return results;
};

export default defineNuxtModule({
  meta: {
    name: "asset-generator",
    configKey: "assetGenerator",
  },
  defaults: {},
  async setup(_, nuxt) {
    const logger = useLogger("asset-generator");
    const rootDir = nuxt.options.rootDir;
    const imagesDir = join(rootDir, "public/nds/images");
    const modelsDir = join(rootDir, "public/nds/models");
    const templateFile = join(rootDir, "app/composables/useAssets.ts.in");
    const outputFile = join(rootDir, "app/composables/useAssets.ts");
    const atlasOutputPath = join(rootDir, "public/nds/atlas.webp");

    const loadImages = async (
      paths: string[],
    ): Promise<Map<string, ImageData>> => {
      const images = new Map<string, ImageData>();

      await Promise.all(
        paths.map(async (path) => {
          try {
            const [metadata, buffer] = await Promise.all([
              sharp(path).metadata(),
              sharp(path).png().toBuffer(),
            ]);

            if (metadata.width && metadata.height) {
              images.set(path, {
                path,
                buffer,
                width: metadata.width,
                height: metadata.height,
              });
            }
          } catch {
            logger.warn(`Failed to load ${path}, skipping...`);
          }
        }),
      );

      return images;
    };

    const packImages = (images: ImageData[]) => {
      const sorted = images.toSorted((a, b) => b.height - a.height);
      const rects = new Map<string, AtlasRect>();
      let x = 0;
      let y = 0;
      let rowHeight = 0;
      let maxWidth = 0;

      for (const img of sorted) {
        if (x + img.width > MAX_WIDTH) {
          x = 0;
          y += rowHeight;
          rowHeight = 0;
        }

        rects.set(img.path, {
          x,
          y,
          width: img.width,
          height: img.height,
        });

        x += img.width;
        rowHeight = Math.max(rowHeight, img.height);
        maxWidth = Math.max(maxWidth, x);
      }

      return { rects, width: maxWidth, height: y + rowHeight };
    };

    const buildImageTree = (
      imagePaths: string[],
      rects: Map<string, AtlasRect>,
    ): ImageTree => {
      const tree: ImageTree = {};

      for (const path of imagePaths) {
        const parts = relative(imagesDir, path).split("/");
        const fileName = parse(parts.at(-1)!).name;

        let node = tree;
        for (const part of parts.slice(0, -1)) {
          const key = toCamelCase(part);
          node[key] ??= {};
          node = node[key] as ImageTree;
        }

        const rect = rects.get(path);
        if (rect) node[toCamelCase(fileName)] = rect;
      }

      return tree;
    };

    const buildModelTree = (modelPaths: string[]): ModelTree => {
      const tree: ModelTree = {};

      for (const path of modelPaths) {
        const parts = relative(modelsDir, path).split("/");
        const fileName = parse(parts.at(-1)!).name;

        let node = tree;
        for (const part of parts.slice(0, -1)) {
          const key = toCamelCase(part);
          node[key] ??= {};
          node = node[key] as ModelTree;
        }

        node[toCamelCase(fileName)] =
          `/nds/models/${relative(modelsDir, path)}`;
      }

      return tree;
    };

    const generateImageCode = (tree: ImageTree, indent = 0): string => {
      const entries = Object.entries(tree);
      if (!entries.length) return "{}";

      const sp = "  ".repeat(indent);
      const lines = entries.map(([key, value]) => {
        if (value && typeof value === "object" && "x" in value) {
          const { x, y, width, height } = value;
          return `${sp}  ${key}: {
${sp}    draw: (ctx, x, y, opts?) => drawAtlasImage(ctx, [${x}, ${y}, ${width}, ${height}], [x, y, ${width}, ${height}], opts),
${sp}    rect: { x: ${x}, y: ${y}, width: ${width}, height: ${height} },
${sp}  }`;
        }
        return `${sp}  ${key}: ${generateImageCode(value, indent + 1)}`;
      });

      return `{\n${lines.join(",\n")}\n${sp}}`;
    };

    const generateModelCode = (tree: ModelTree, indent = 0): string => {
      const entries = Object.entries(tree);
      if (!entries.length) return "{}";

      const sp = "  ".repeat(indent);
      const lines = entries.map(([key, value]) => {
        if (typeof value === "string") {
          return `${sp}  ${key}: createModel("${value}")`;
        }
        return `${sp}  ${key}: ${generateModelCode(value, indent + 1)}`;
      });

      return `{\n${lines.join(",\n")},\n${sp}}`;
    };

    const generateAssets = async () => {
      try {
        const imagePaths = await scanByExt(imagesDir, IMAGE_EXTENSIONS);
        const modelPaths = await scanByExt(modelsDir, MODEL_EXTENSIONS);

        const totalAssets = (imagePaths.length > 0 ? 1 : 0) + modelPaths.length;

        let imageCode = "{}";
        let atlasHash = "";
        let atlasWidth = 0;
        let atlasHeight = 0;

        if (imagePaths.length > 0) {
          const imageMap = await loadImages(imagePaths);
          const images = Array.from(imageMap.values());

          const { rects, width, height } = packImages(images);
          atlasWidth = width;
          atlasHeight = height;

          const atlasBuffer = await sharp({
            create: {
              width,
              height,
              channels: 4,
              background: { r: 0, g: 0, b: 0, alpha: 0 },
            },
          })
            .composite(
              images.map((img) => {
                const rect = rects.get(img.path)!;
                return { input: img.buffer, left: rect.x, top: rect.y };
              }),
            )
            .webp({ lossless: true })
            .toBuffer();

          await writeFile(atlasOutputPath, atlasBuffer);

          execSync(
            `magick "${atlasOutputPath}" -define webp:lossless=true -define webp:method=6 "${atlasOutputPath}"`,
            { stdio: "pipe" },
          );

          const finalAtlas = await readFile(atlasOutputPath);
          atlasHash = createHash("md5")
            .update(finalAtlas)
            .digest("hex")
            .slice(0, 8);

          const imageTree = buildImageTree(imagePaths, rects);
          imageCode = generateImageCode(imageTree);
        }

        const modelTree = buildModelTree(modelPaths);
        const modelCode = generateModelCode(modelTree);

        const template = await readFile(templateFile, "utf-8");
        const code = template
          .replace("{{TOTAL}}", totalAssets.toString())
          .replace("{{ATLAS_HASH}}", atlasHash)
          .replace("{{IMAGES}}", imageCode)
          .replace("{{MODELS}}", modelCode);

        await writeFile(outputFile, code, "utf-8");
        logger.success(
          `Generated assets with ${imagePaths.length} images (${atlasWidth}x${atlasHeight}) and ${modelPaths.length} models`,
        );
      } catch (error) {
        logger.error("Error generating assets:", error);
      }
    };

    await generateAssets();

    if (nuxt.options.dev) {
      const watchAndRegenerate = (
        dir: string,
        extensions: string[],
        type: string,
      ) => {
        watch(dir, { recursive: true }, (_, filePath) => {
          if (
            filePath &&
            extensions.includes(parse(filePath).ext.toLowerCase())
          ) {
            logger.info(`Detected ${type} change: ${filePath}`);
            generateAssets();
          }
        });
      };

      nuxt.hook("ready", () => {
        watchAndRegenerate(imagesDir, IMAGE_EXTENSIONS, "image");
        watchAndRegenerate(modelsDir, MODEL_EXTENSIONS, "model");
        watch(templateFile, () => {
          logger.info("Template file changed");
          generateAssets();
        });
      });
    }
  },
});