import type { Plugin, UserConfig } from "vite";
import type { OutputAsset, OutputChunk } from "rolldown";
import path from "node:path";
import fs from "node:fs";

type MinimalOutputAsset = Pick<OutputAsset, "type" | "fileName" | "source">;
type MinimalOutputChunk = Pick<OutputChunk, "type" | "fileName" | "code">;
type MinimalOutputBundle = Record<string, MinimalOutputAsset | MinimalOutputChunk>;

interface PublicDirFile {
  fullPath: string;
  relativePath: string;
}

const textEncoder = new TextEncoder();

function getPublicDirFiles(publicDir?: string): PublicDirFile[] {
  if (!publicDir) return [];

  const result: PublicDirFile[] = [];
  const files = fs.readdirSync(publicDir, { recursive: true, withFileTypes: true });
  for (const file of files) {
    if (!file.isFile()) continue;
    const fullPath = path.join(file.parentPath, file.name);
    let relativePath = path.relative(publicDir, fullPath).replaceAll(/\\/g, "/");
    if (relativePath.startsWith("./")) {
      relativePath = relativePath.substring(2);
    }

    result.push({
      fullPath,
      relativePath,
    });
  }

  return result;
}

function createVarName(fileName: string) {
  return fileName.replaceAll(/[\\/]/g, "__").replaceAll(/[\.-]/g, "_").toUpperCase();
}

function transformAsset(asset: MinimalOutputAsset) {
  const varName = createVarName(asset.fileName);

  let buffer: Uint8Array;
  if (typeof asset.source === "string") {
    buffer = textEncoder.encode(asset.source);
  } else {
    buffer = asset.source;
  }

  const bytes = [...buffer].map((v) => String(v)).join(",");

  return `constexpr const unsigned char ${varName}[] {${bytes}};
`;
}

function transformChunk(chunk: MinimalOutputChunk) {
  const varName = createVarName(chunk.fileName);
  const buffer = textEncoder.encode(chunk.code);
  const bytes = [...buffer].map((v) => String(v)).join(",");

  return `constexpr const unsigned char ${varName}[] {${bytes}};
`;
}

function transformPublicFile(publicFile: PublicDirFile) {
  const varName = createVarName(publicFile.relativePath);
  const bytes = [...fs.readFileSync(publicFile.fullPath)].map((v) => String(v)).join(",");

  return `constexpr const unsigned char ${varName}[] {${bytes}};
`;
}

function writeHeader(
  bundle: MinimalOutputBundle,
  outputDir?: string,
  options?: HeaderTransformationPluginOptions,
  publicDir?: string,
  devServerPort?: number,
) {
  const outputPath = options?.outputPath ?? path.join(outputDir ?? "dist", "ViteAssets.h");
  const outputPathParentDir = path.dirname(outputPath);

  fs.mkdirSync(outputPathParentDir, { recursive: true });

  const fd = fs.openSync(outputPath, "w");
  const includeFileEnumeration = options?.includeFileEnumeration ?? true;

  fs.writeSync(
    fd,
    `#pragma once

`,
  );

  if (includeFileEnumeration) {
    fs.writeSync(
      fd,
      `#include <cstdlib>
#include <type_traits>

`,
    );
  }

  fs.writeSync(
    fd,
    `constexpr auto VITE_DEV_SERVER = ${devServerPort ? "true" : "false"};
constexpr auto VITE_DEV_SERVER_PORT = ${devServerPort ? String(devServerPort) : "-1"};
`,
  );

  const fileNames: string[] = [];
  for (const curBundle of Object.values(bundle)) {
    if (curBundle.type === "asset") {
      fs.writeSync(fd, transformAsset(curBundle));
    } else {
      fs.writeSync(fd, transformChunk(curBundle));
    }
    fileNames.push(curBundle.fileName);
  }
  for (const publicDirFile of getPublicDirFiles(publicDir)) {
    fs.writeSync(fd, transformPublicFile(publicDirFile));
    fileNames.push(publicDirFile.relativePath);
  }

  if (includeFileEnumeration) {
    fs.writeSync(
      fd,
      `
struct UiFile
{
    const char* filename;
    const void* data;
    const size_t dataSize;
};

static inline const UiFile MOD_MAN_UI_FILES[] {
`,
    );

    let index = 0;
    for (const fileName of fileNames) {
      const varName = createVarName(fileName);

      let prefix = "  ";
      if (index > 0) {
        prefix = `,
  `;
      }

      fs.writeSync(
        fd,
        `${prefix}{ "${fileName}", ${varName}, std::extent_v<decltype(${varName})> }`,
      );
      index++;
    }

    fs.writeSync(
      fd,
      `
};
`,
    );

    fs.closeSync(fd);
  }
}

export interface HeaderTransformationPluginOptions {
  outputPath?: string;
  includeFileEnumeration?: boolean;
}

export default function headerTransformationPlugin(
  options?: HeaderTransformationPluginOptions,
): Plugin {
  let writeServerActive = false;
  let writeBundleActive = false;
  let publicDir: string | undefined = undefined;

  return {
    name: "vite-plugin-header-transformation",
    enforce: "post",
    config(userOptions: UserConfig, env) {
      if (env.command === "serve") {
        writeServerActive = true;
      } else {
        writeBundleActive = true;
      }

      if (typeof userOptions.publicDir === "string") {
        publicDir = userOptions.publicDir;
      }
    },
    configureServer(server) {
      if (!writeServerActive) {
        return;
      }
      server.httpServer?.once("listening", () => {
        writeHeader(
          {
            // We need at least one array entry for MSVC
            dummyfile: { type: "chunk", fileName: "dummyfile", code: "dummy" },
          },
          server.config.build.outDir,
          options,
          publicDir,
          server.config.server.port,
        );
      });
    },
    writeBundle(outputOptions, bundle) {
      if (!writeBundleActive) {
        return;
      }

      writeHeader(bundle, outputOptions.dir, options, publicDir);
    },
  };
}