import sys, os
import argparse
from datetime import datetime

# Check if the script is run with the correct arguments
parser = argparse.ArgumentParser(description="Generate chunk header files")
parser.add_argument('--output', required=True, help='Dir to write headers')
parser.add_argument('--input', required=True, help='Input directory containing language files')
args = parser.parse_args()

# Ensure outdir exists
outputFile = args.output
outputDir = args.output
os.makedirs(outputDir, exist_ok=True)

inputDir = args.input
# Scan for .json files in the input directory
if not os.path.exists(inputDir):
  print(f"Error: Input directory '{inputDir}' does not exist.")
  sys.exit(1)

jsonFiles = [f for f in os.listdir(inputDir) if f.endswith('.json')]
if not jsonFiles or len(jsonFiles) == 0:
  print(f"Error: No JSON files found in '{inputDir}'.")
  sys.exit(1)

# take JSON from form { "a": { "b": { "c": "d" } } } to "a.b.c": "d"
def flattenJson(y):
  keyValues = {}
  for key, value in y.items():
    if isinstance(value, dict):
      # If the value is a dictionary, recurse into it
      subKeyValues = flattenJson(value)
      for subKey, subValue in subKeyValues.items():
        keyValues[f"{key}.{subKey}"] = subValue
    else:
      # If the value is not a dictionary, add it to the keyValues
      keyValues[key] = value
  return keyValues

def escapeString(s):
  # Escape double quotes and backslashes in the string
  return s.replace('\\', '\\\\').replace('"', '\\"')

# For each language file...
now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
isFirstLanguage = True

# Because I code in english, I am going to reorder the langs so it is first.
jsonFiles.sort(key=lambda x: x.lower() if x.lower() == 'en.json' else "zz_" + x.lower())

keysExpected = []

languages = []
for jsonFile in jsonFiles:
  inputFile = os.path.join(inputDir, jsonFile)
  languageName = os.path.splitext(jsonFile)[0]
  langUpper = languageName.upper()
  outputFile = os.path.join(outputDir, f"{languageName}.h")

  # Read the JSON file
  with open(inputFile, 'r', encoding='utf-8') as f:
    content = f.read()

  # Write the header file
  with open(outputFile, 'w', encoding='utf-8') as f:
    f.write(f"// Generated from {jsonFile} on {now}\n")
    f.write("#pragma once\n")
    f.write("#include \"dusk.h\"\n\n")
    f.write(f"// Language: {languageName} from {jsonFile}\n")
    
    keyValues = flattenJson(eval(content))

    if 'meta.language.name' not in keyValues:
      print(f"Error: 'meta.language.name' not found in {jsonFile}.")
      sys.exit(1)

    f.write(f"#define LANGUAGE_{langUpper}_CODE \"{languageName}\"\n")
    f.write(f"#define LANGUAGE_{langUpper}_NAME \"{keyValues['meta.language.name']}\"\n")
    f.write(f"#define LANGUAGE_{langUpper}_COUNT {len(keyValues)}\n\n")

    # Write keys
    f.write(f"static const char_t *LANGUAGE_{langUpper}_KEYS[] = {{\n")
    for key in keyValues.keys():
      f.write(f'  "{escapeString(key)}",\n')
    f.write("};\n\n")

    # Write values
    f.write(f"static const char_t *LANGUAGE_{langUpper}_VALUES[] = {{\n")
    for value in keyValues.values():
      f.write(f'  "{escapeString(value)}",\n')
    f.write("};\n\n")

    languages.append(langUpper)

    if isFirstLanguage:
      # For the first language, we also write the keysExpected
      keysExpected = list(keyValues.keys())
    else:
      for key in keysExpected:
        if key in keyValues:
          continue
        print(f"Error, expected language translation key: '{key}' was not found in {jsonFile}.")
        sys.exit(1)

# Now write the main header file
mainOutputFile = os.path.join(outputDir, "languages.h")
with open(mainOutputFile, 'w', encoding='utf-8') as f:
  f.write("// Generated from languagecompile.py\n")
  f.write("#pragma once\n")
  f.write("#include \"dusk.h\"\n")
  for lang in languages:
    f.write(f'#include "locale/language/{lang.lower()}.h"\n')
  f.write("\n")

  f.write(f"#define LANGUAGE_COUNT {len(languages)}\n\n")

  index = 0
  for lang in languages:
    f.write(f"#define LANGUAGE_{lang} {index}\n")
    index += 1
  f.write("\n")

  f.write("static const char_t* LANGUAGE_CODES[] = {\n")
  for lang in languages:
    f.write(f'  LANGUAGE_{lang}_CODE,\n')
  f.write("};\n\n")

  f.write("static const char_t* LANGUAGE_NAMES[] = {\n")
  for lang in languages:
    f.write(f'  LANGUAGE_{lang}_NAME,\n')
  f.write("};\n\n")

  f.write("static const char_t** LANGUAGE_KEYS[] = {\n")
  for lang in languages:
    f.write(f'  LANGUAGE_{lang}_KEYS,\n')
  f.write("};\n\n")

  f.write("static const int LANGUAGE_COUNTS[] = {\n")
  for lang in languages:
    f.write(f'  LANGUAGE_{lang}_COUNT,\n')
  f.write("};\n\n")

  f.write("static const char_t** LANGUAGE_VALUES[] = {\n")
  for lang in languages:
    f.write(f'  LANGUAGE_{lang}_VALUES,\n')
  f.write("};\n\n")