#include <dirent.h>
#include <errno.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

/**
 * Kettle - Build system for HLSL and GLSL shaders
 * Copyright (c) 2025 Hunter Stasonis
 * Licensed under the MIT license
 */

/// Type of file
typedef enum {
  e_kettle_file_hlsl,
  e_kettle_file_glsl,
  e_kettle_file_unknown
} kettle_file;

typedef enum { e_kettle_compiler_dxc, e_kettle_compiler_glslc } kettle_compiler;

/**
 * File information
 */
typedef struct {
  const char *pFileName;
  kettle_file eFileType;
} kettle_file_info_t;

typedef struct {
  kettle_compiler eCompilerType;

  const char *pCompilerPath;
} kettle_compiler_t;

typedef struct {
  int numFiles;
  kettle_file_info_t *pFiles;

  int numIncludeDirs;

  /// List of directories to -I in this batch
  char **pIncludeDirs;

  /// Directory that all files are located in
  const char *pWorkingDir;
} kettle_compiler_batch_t;

typedef struct {
  int numCompilerBatches;

  /// Compiler batch list
  kettle_compiler_batch_t *pCompilerBatches;

  /// Current compiler batch
  kettle_compiler_batch_t *pCurrentCompilerBatch;

  /// List of compilers
  int numCompilers;
  kettle_compiler_t *pCompilers;
} kettle_input_t;

typedef struct {
} kettle_reg_t;

typedef struct {
  int numArgs;
  char **pArgs;
} kettle_arg_list_t;

kettle_file kettle_get_file_type(const char *pFileName) {
  if (strstr(pFileName, ".glsl") != NULL) {
    return e_kettle_file_glsl;
  } else if (strstr(pFileName, ".hlsl") != NULL) {
    return e_kettle_file_hlsl;
  } else {
    return e_kettle_file_unknown;
  }
}

void kettle_help() {
  printf("Kettle Build System\n");
  printf("Description: Compiles kettle information files and prints the result "
         "to stderr\n");
  printf("Author: Hunter Stasonis <webmaster @ interfiber.dev>\n");
  printf("Options: \n");
  printf("\t infile : Input kettle description file\n");
  printf("No arguments provided! Exiting\n");
}

/**
 * Remove leading spaces from a string
 * @param pInput Pointer to a string to trim
 * @return Newly allocated string with spaces removed from the start
 */
char *kettle_trim(const char *pInput) {
  int startIndex = 0;
  int size = strlen(pInput) + 1;

  // Find starting index: Loop over characters until we get a non-space
  for (int i = 0; i < size; i++) {
    if (pInput[i] != ' ') {
      startIndex = i;
      break;
    }
  }

  // Copy the rest of the string into the new allocated buffer

  int newSize = size - startIndex;

  char *newStr = (char *)malloc(newSize * sizeof(char));
  strncpy(newStr, pInput + startIndex, newSize);

  return newStr;
}

char *kettle_pop_index(const char *pInput, int index) {
  int size = strlen(pInput);
  int newSize = size - 1;

  char *newStr = (char *)malloc(newSize * sizeof(char));
  memmove(&newStr[index], &pInput[index + 1], strlen(pInput) - index);

  return newStr;
}

/**
 * Concat the contents of pInput with pInput2, automatically reallocating the
 * memory of pInput
 * @param pInput Input string buffer
 * @param pInput2 Input to append to pInput
 */
char *kettle_str_concat(char *pInput, const char *pInput2) {
  int newSize = strlen(pInput) + strlen(pInput2) + 1; // +1 for null byte
  newSize *= sizeof(char);

  pInput = realloc(pInput, newSize);

  strcat(pInput, pInput2);

  return pInput;
}

kettle_arg_list_t kettle_make_arg_list(const char *pInput) {
  int size = strlen(pInput);

  char *pCopyInput = (char *)malloc(size * sizeof(char));
  memcpy(pCopyInput, pInput, size * sizeof(char));

  char **args = NULL;
  int cArgSize = 0;
  int cArgIndex = 0;

  char *tok = strtok(pCopyInput, " ");
  while (tok != NULL) {
    cArgSize++;

    args = (char **)realloc(args, (cArgSize + 1) * sizeof(char *));
    args[cArgIndex] = tok;

    cArgIndex++;

    tok = strtok(NULL, " ");
  }

  kettle_arg_list_t argList;
  argList.numArgs = cArgSize;
  argList.pArgs = args;

  return argList;
}

kettle_compiler_t *kettle_input_get_compiler(const kettle_input_t *pInput,
                                             kettle_compiler eCompilerType) {
  for (int i = 0; i < pInput->numCompilers; i++) {
    if (pInput->pCompilers[i].eCompilerType == eCompilerType)
      return &pInput->pCompilers[i];
  }

  return NULL;
}

void kettle_handle_cmd(kettle_arg_list_t argList, kettle_input_t *pInput) {
  const char *pCmd = argList.pArgs[0];

  if (strcmp(pCmd, "AddVariant") == 0 && argList.numArgs == 2) {
    printf("> Adding variant with define: %s\n", argList.pArgs[1]);
  } else if (strcmp(pCmd, "AddCompiler") == 0 && argList.numArgs == 3) {
    printf("> Adding compiler with type '%s' with path '%s'\n",
           argList.pArgs[1], argList.pArgs[2]);

    kettle_compiler_t compiler;
    compiler.pCompilerPath = argList.pArgs[2];

    if (strcmp(argList.pArgs[1], "dxc") == 0) {
      compiler.eCompilerType = e_kettle_compiler_dxc;
    } else if (strcmp(argList.pArgs[1], "glslc") == 0) {
      compiler.eCompilerType = e_kettle_compiler_glslc;
    } else {
      printf("> Invalid compiler identifier for AddCompiler call. Expected "
             "'glslc', or 'dxc'\n");
      exit(EXIT_FAILURE);
    }

    // Grow compiler array

    pInput->numCompilers++;
    pInput->pCompilers = (kettle_compiler_t *)realloc(
        pInput->pCompilers, sizeof(kettle_compiler_t) * (pInput->numCompilers));

    pInput->pCompilers[pInput->numCompilers - 1] = compiler;

  } else if (strcmp(pCmd, "CompilerBatch") == 0 && argList.numArgs == 3) {
    printf("> Starting compiler batch with type '%s' with workingDir of '%s'\n",
           argList.pArgs[1], argList.pArgs[2]);

    kettle_compiler_batch_t batch;
    memset(&batch, 0, sizeof(batch));

    batch.pWorkingDir = argList.pArgs[2];

    // Grow compiler batch array
    pInput->numCompilerBatches++;

    pInput->pCompilerBatches = (kettle_compiler_batch_t *)realloc(
        pInput->pCompilerBatches,
        sizeof(kettle_compiler_batch_t) * pInput->numCompilerBatches);

    pInput->pCompilerBatches[pInput->numCompilerBatches - 1] = batch;

    // Activate this compiler batch
    pInput->pCurrentCompilerBatch =
        &pInput->pCompilerBatches[pInput->numCompilerBatches - 1];

  } else if (strcmp(pCmd, "EndCompilerBatch") == 0) {
    printf("> Ending active compiler batch\n");

    if (pInput->pCurrentCompilerBatch == NULL) {
      printf("> Error: Cannot end compiler batch when no batch is active!\n");
      exit(EXIT_FAILURE);
    }

    // Deactivate current compiler batch
    pInput->pCurrentCompilerBatch = NULL;
  } else if (strcmp(pCmd, "Include") == 0 && argList.numArgs == 2) {
    if (pInput->pCurrentCompilerBatch == NULL) {
      printf("> Error: No active compiler batch for 'Include' command!\n");
      exit(EXIT_FAILURE);
    }

    // The new include dir may not exist, this will be checked during actual
    // shader build time.

    // Insert new include dir into the list

    pInput->pCurrentCompilerBatch->numIncludeDirs++;

    pInput->pCurrentCompilerBatch->pIncludeDirs = (char **)realloc(
        pInput->pCurrentCompilerBatch->pIncludeDirs,
        sizeof(char *) * pInput->pCurrentCompilerBatch->numIncludeDirs);

    pInput->pCurrentCompilerBatch
        ->pIncludeDirs[pInput->pCurrentCompilerBatch->numIncludeDirs - 1] =
        argList.pArgs[1];
  } else if (strcmp(pCmd, "AddShader") == 0 && argList.numArgs == 2) {
    if (pInput->pCurrentCompilerBatch == NULL) {
      printf("> Error: Cannot add shader when no compiler batch is active!\n");
      exit(EXIT_FAILURE);
    }

    const char *pFileName = argList.pArgs[1];

    kettle_file_info_t fileInf;
    memset(&fileInf, 0, sizeof(fileInf));
    fileInf.eFileType = kettle_get_file_type(pFileName);

    if (fileInf.eFileType == e_kettle_file_unknown) {
      printf("> Error: Cannot determine file type of path '%s'!\n", pFileName);
      exit(EXIT_FAILURE);
    }

    fileInf.pFileName = pFileName;

    // Allocate a new file

    pInput->pCurrentCompilerBatch->numFiles++;

    pInput->pCurrentCompilerBatch->pFiles = (kettle_file_info_t *)realloc(
        pInput->pCurrentCompilerBatch->pFiles,
        sizeof(kettle_file_info_t) * pInput->pCurrentCompilerBatch->numFiles);

    pInput->pCurrentCompilerBatch
        ->pFiles[pInput->pCurrentCompilerBatch->numFiles - 1] = fileInf;
  } else {
    printf("> Warning: Cannot process unknown compiler command '%s'\n", pCmd);
  }
}

kettle_input_t kettle_read_input(const char *pInput) {
  FILE *fInput = fopen(pInput, "r");

  kettle_input_t input;
  memset(&input, 0, sizeof(input));

  // Read file line, by line
  char lineTmp[256];

  while (fgets(lineTmp, sizeof(lineTmp), fInput)) {
    char *line = kettle_trim(lineTmp);
    int lineLength = strlen(line);
    line[strcspn(line, "\n")] = 0;

    lineLength = strlen(line);

    if (lineLength == 0)
      continue;

    if (lineLength >= 2) {
      if (line[0] == '/' && line[1] == '/') {
        continue;
      } else if (line[0] == '$') {
        const char *pCmd = kettle_pop_index(line, 0);
        const kettle_arg_list_t argList = kettle_make_arg_list(pCmd);

        kettle_handle_cmd(argList, &input);
      } else {
        printf("Error: Line '%s' is neither a comment, or a compiler "
               "instruction!\n",
               line);

        exit(EXIT_FAILURE);
      }
    }
  }

  fclose(fInput);

  return input;
}

const char *kettle_hlsl_compile(const kettle_compiler_t *pCompiler,
                                const kettle_compiler_batch_t *pBatch,
                                const kettle_file_info_t *pFileInfo) {
  char *pCommand =
      (char *)malloc((strlen(pCompiler->pCompilerPath) + 1) * sizeof(char));

  strcpy(pCommand, pCompiler->pCompilerPath);

  // This might be a little slow, but we need to create the command

  // We force the DXC compiler to output spirv, this could change later
  pCommand = kettle_str_concat(pCommand, " -spirv");

  return pCommand;
}

int main(int argc, char **argv) {
  if (argc <= 1) {
    kettle_help();
    exit(EXIT_FAILURE);
  }

  const char *inFile = argv[1];

  printf("> Will process input file: %s\n", inFile);

  kettle_input_t inputProcessed = kettle_read_input(inFile);

  printf("> Will process %i compiler batche(s)...\n",
         inputProcessed.numCompilerBatches);

  for (int i = 0; i < inputProcessed.numCompilerBatches; ++i) {
    const kettle_compiler_batch_t batch = inputProcessed.pCompilerBatches[i];

    printf("> Change directory to: %s\n", batch.pWorkingDir);

    if (chdir(batch.pWorkingDir) != 0) {
      printf("> Error: Failed to chdir into new directory: %s\n",
             strerror(errno));
      exit(EXIT_FAILURE);
    }

    for (int x = 0; x < batch.numFiles; ++x) {
      kettle_file_info_t file = batch.pFiles[x];

      kettle_compiler_t *pCompiler =
          kettle_input_get_compiler(&inputProcessed, e_kettle_compiler_dxc);

      if (pCompiler == NULL) {
        printf("> Error: Failed to find compatible DXC compiler!\n");
        exit(EXIT_FAILURE);
      }

      printf("> %s\n", kettle_hlsl_compile(pCompiler, &batch, &file));
    }
  }
}