package puzzle;

import puzzle.SwedishGenerator.PuzzleResult;
import puzzle.SwedishGenerator.Rng;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.time.OffsetDateTime;
import java.time.ZoneOffset;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.Locale;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executors;
import static puzzle.SwedishGenerator.fillMask;
import static puzzle.SwedishGenerator.generateMask;
import static puzzle.SwedishGenerator.loadScores;
import static puzzle.SwedishGenerator.loadWords;

public class Main {
   // ---------------- CLI ----------------

   public static class Opts {
      public int seed = 1234;
      public int pop = 18;
      public int gens = 300;
      public String wordsPath = "/data/puzzle/pool.txt";
      public double minSimplicity = 0; // 0 means no limit
      public int threads = Math.max(1, Runtime.getRuntime().availableProcessors());
      public int tries = threads;
   }
   
   static void usage() {
      System.out.println("""
      Usage:
        java SwedishGenerator [--seed N] [--pop N] [--gens N] [--tries N] [--words word-list.txt] [--min-simplicity N.N] [--threads N]

      Defaults:
        --seed 1234
        --pop 18
        --gens 1000
        --tries 5
        --words /data/puzzle/pool.txt
        --min-simplicity 0 (no limit)
        --threads %d
      """.formatted(Math.max(1, Runtime.getRuntime().availableProcessors())));
   }
   
   static Opts parseArgs(String[] argv) {
      var out = new Opts();
      for (int i = 0; i < argv.length; i++) {
         String a = argv[i];
         String v = (i + 1 < argv.length) ? argv[i + 1] : null;
         if (a.equals("--help") || a.equals("-h")) {
            usage();
            System.exit(0);
         }
         if (a.equals("--seed")) { out.seed = Integer.parseInt(v); i++; }
         else if (a.equals("--pop")) { out.pop = Integer.parseInt(v); i++; }
         else if (a.equals("--gens")) { out.gens = Integer.parseInt(v); i++; }
         else if (a.equals("--tries")) { out.tries = Integer.parseInt(v); i++; }
         else if (a.equals("--words")) { out.wordsPath = v; i++; }
         else if (a.equals("--min-simplicity")) { out.minSimplicity = Double.parseDouble(v); i++; }
         else if (a.equals("--threads")) { out.threads = Integer.parseInt(v); i++; }
         else throw new IllegalArgumentException("Unknown arg: " + a);
      }
      return out;
   }
   
   public static void main(String[] args) {
      var opts = parseArgs(args);
      var res = generatePuzzle(opts);
      if (res == null) {
         System.out.println("No solution found within tries.");
         System.exit(1);
      }

      System.out.println("\n=== GENERATED MASK ===");
      System.out.println(SwedishGenerator.gridToString(res.mask()));

      System.out.println("\n=== FILLED PUZZLE (RAW) ===");
      System.out.println(SwedishGenerator.gridToString(res.filled().grid));

      System.out.println("\n=== FILLED PUZZLE (HUMAN) ===");
      System.out.println(SwedishGenerator.renderHuman(res.filled().grid));
      System.out.printf(Locale.ROOT, "Puzzle Simplicity: %.2f%n", res.filled().simplicity);
      var out = ExportFormat.exportFormatFromFilled(res, 1, new ExportFormat.Rewards(50, 2, 1));
      
      // Generate clues via LLM
      System.out.println("Generating clues for " + out.words().size() + " words...");
      out = ClueGenerator.applyClues(out);

      System.out.println("gridv2:");
      for (String row : out.gridv2()) System.out.println(row);
      System.out.println("words: " + out.words().size());
      for (var w : out.words()) {
         var simplicityOfWord =
         System.out.printf("%s %s start=(%d,%d) arrow=(%d,%d)%n",
                           w.word(), w.direction(), w.startRow(), w.startCol(), w.arrowRow(), w.arrowCol());
      }

      // Export to JSON file
      var now = OffsetDateTime.now(ZoneOffset.UTC);
      var createdAt = now.format(DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss'Z'"));
      var dateStr = now.toLocalDate().toString();
      var id = createdAt.replace(":", "-") + "_" + (System.currentTimeMillis() / 1000);
      var theme = "algemeen";
      var filename = id + ".json";
      var outDir = "/data/puzzle/puzzles";
      var outputPath = Paths.get(outDir, filename);

      try {
         Files.createDirectories(Paths.get(outDir));
         var json = toJson(out, dateStr, theme);
         Files.writeString(outputPath, json, StandardCharsets.UTF_8);
         System.out.println("\nSaved to: " + outputPath);

         // Update index.json
         var pathInIndex = "/puzzles/" + filename;
         var indexRecord = toIndexRecordJson(id, pathInIndex, dateStr, theme, out.difficulty(), createdAt);
         updateIndex(outDir, indexRecord);
      } catch (IOException e) {
         System.err.println("Failed to write " + filename + ": " + e.getMessage());
      }
   }

   // Package-private method for testing
   static PuzzleResult generatePuzzle(Opts opts) {
      var llmScores = loadScores();
      var tLoad0    = System.nanoTime();
      var dict      = loadWords(opts.wordsPath, llmScores);
      var tLoad1    = System.nanoTime();
      System.out.printf(Locale.ROOT, "LOAD_WORDS: %.3fs%n %s words%n", (tLoad1 - tLoad0) / 1e9, dict.words.size());
      
      if (opts.threads > 1) {
         System.out.println("Running in multi-threaded mode with " + opts.threads + " threads...");
         var executor = Executors.newFixedThreadPool(opts.threads);
         try {
            var tasks = new ArrayList<Callable<PuzzleResult>>();
            for (int i = 1; i <= opts.tries; i++) {
               final int attempt = i;
               tasks.add(() -> {
                  var threadRng = new Rng(opts.seed + attempt);
                  var mask      = generateMask(threadRng, dict.lenCounts, opts.pop, opts.gens, false);
                  var filled    = fillMask(threadRng, mask, dict.index, llmScores, 200, 60000, false);
                  
                  if (filled.ok && (opts.minSimplicity <= 0 || filled.simplicity >= opts.minSimplicity)) {
                     System.out.println("\nSolution found on attempt " + attempt);
                     return new PuzzleResult(mask, filled);
                  }
                  throw new RuntimeException("No solution found in attempt " + attempt);
               });
            }
            return executor.invokeAny(tasks);
         } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
         } catch (ExecutionException e) {
            // all failed
         } finally {
            executor.shutdownNow();
         }
         return null;
      } else {
         var rng = new Rng(opts.seed);
         for (var attempt = 1; attempt <= opts.tries; attempt++) {
            System.out.println("\nAttempt " + attempt + "/" + opts.tries);
            
            var tMask0 = System.nanoTime();
            var mask   = generateMask(rng, dict.lenCounts, opts.pop, opts.gens, true);
            var tMask1 = System.nanoTime();
            System.out.printf(Locale.ROOT, "MASK: %.3fs%n", (tMask1 - tMask0) / 1e9);
            
            var tFill0 = System.nanoTime();
            var filled = fillMask(rng, mask, dict.index, llmScores, 200, 60000, true);
            var tFill1 = System.nanoTime();
            System.out.printf(Locale.ROOT, "FILL: %.3fms | Simplicity: %.2f%n", (tFill1 - tFill0) / 1e6, filled.simplicity);
            
            if (filled.ok && (opts.minSimplicity <= 0 || filled.simplicity >= opts.minSimplicity)) {
               return new PuzzleResult(mask, filled);
            }
            if (filled.ok) {
               System.out.printf(Locale.ROOT, "Puzzle simplicity %.2f is below min %.2f, retrying...%n",
                                 filled.simplicity, opts.minSimplicity);
            }
         }
      }
      
      return null;
   }

   private static String toJson(ExportFormat.ExportedPuzzle puzzle, String date, String theme) {
      var sb = new StringBuilder();
      sb.append("{\n");
      sb.append("  \"date\": \"").append(escapeJson(date)).append("\",\n");
      sb.append("  \"theme\": \"").append(escapeJson(theme)).append("\",\n");
      sb.append("  \"difficulty\": ").append(puzzle.difficulty()).append(",\n");
      sb.append("  \"rewards\": {\n");
      sb.append("    \"coins\": ").append(puzzle.rewards().coins()).append(",\n");
      sb.append("    \"stars\": ").append(puzzle.rewards().stars()).append(",\n");
      sb.append("    \"hints\": ").append(puzzle.rewards().hints()).append("\n");
      sb.append("  },\n");
      sb.append("  \"gridv2\": [\n");
      for (var i = 0; i < puzzle.gridv2().size(); i++) {
         sb.append("    \"").append(escapeJson(puzzle.gridv2().get(i))).append("\"");
         if (i < puzzle.gridv2().size() - 1) sb.append(",");
         sb.append("\n");
      }
      sb.append("  ],\n");
      sb.append("  \"words\": [\n");
      for (var i = 0; i < puzzle.words().size(); i++) {
         var w = puzzle.words().get(i);
         sb.append("    {\n");
         sb.append("      \"word\": \"").append(escapeJson(w.word())).append("\",\n");
         sb.append("      \"clue\": \"").append(escapeJson(w.clue())).append("\",\n");
         sb.append("      \"startRow\": ").append(w.startRow()).append(",\n");
         sb.append("      \"startCol\": ").append(w.startCol()).append(",\n");
         sb.append("      \"direction\": \"").append(escapeJson(w.direction())).append("\",\n");
         sb.append("      \"answer\": \"").append(escapeJson(w.answer())).append("\",\n");
         sb.append("      \"arrowRow\": ").append(w.arrowRow()).append(",\n");
         sb.append("      \"arrowCol\": ").append(w.arrowCol()).append(",\n");
         sb.append("      \"isReversed\": ").append(w.isReversed()).append("\n");
         sb.append("    }");
         if (i < puzzle.words().size() - 1) sb.append(",");
         sb.append("\n");
      }
      sb.append("  ]\n");
      sb.append("}\n");
      return sb.toString();
   }

   private static String escapeJson(String s) {
      return s.replace("\\", "\\\\")
              .replace("\"", "\\\"")
              .replace("\n", "\\n")
              .replace("\r", "\\r")
              .replace("\t", "\\t");
   }

   private static String safeSlug(String s) {
      return s.toLowerCase().replaceAll("[^a-z0-9]+", "-").replaceAll("^-|-$", "");
   }

   private static String toIndexRecordJson(String id, String path, String date, String theme, int difficulty, String createdAt) {
      return String.format(
          "{\"id\":\"%s\",\"path\":\"%s\",\"date\":\"%s\",\"theme\":\"%s\",\"difficulty\":%d,\"createdAt\":\"%s\"}",
          escapeJson(id), escapeJson(path), escapeJson(date), escapeJson(theme), difficulty, escapeJson(createdAt)
      );
   }

   private static void updateIndex(String outDir, String newRecordJson) {
      Path indexPath = Paths.get(outDir, "index.json");
      try {
         String content;
         if (Files.exists(indexPath)) {
            content = Files.readString(indexPath, StandardCharsets.UTF_8).trim();
         } else {
            content = "";
         }

         if (content.isEmpty() || content.equals("[]")) {
            content = "[\n  " + newRecordJson + "\n]";
         } else {
             int firstBracket = content.indexOf('[');
             if (firstBracket != -1) {
                content = content.substring(0, firstBracket + 1) + "\n  " + newRecordJson + "," + content.substring(firstBracket + 1);
             } else {
                content = "[\n  " + newRecordJson + "\n]";
             }
         }
         Files.writeString(indexPath, content, StandardCharsets.UTF_8);
         System.out.println("Updated index.json at: " + indexPath);
      } catch (IOException e) {
         System.err.println("Failed to update index.json: " + e.getMessage());
      }
   }
}