package puzzle;

import lombok.Data;
import puzzle.SwedishGenerator.Rng;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.*;
import java.time.OffsetDateTime;
import java.time.ZoneOffset;
import java.time.format.DateTimeFormatter;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Collectors;

import static puzzle.ExportFormat.*;
import static puzzle.SwedishGenerator.*;
import static puzzle.SwedishGenerator.loadWords;

public class Main {
   
   // ---------------- Top-level generatePuzzle ----------------
   public record PuzzleResult(SwedishGenerator swe, Dict dict, Gridded mask, FillResult filled) { }
   
   final static String         OUT_DIR     = envOrDefault("OUT_DIR", "/data/puzzle");
   final static Path           PUZZLE_DIR  = Paths.get(OUT_DIR, "puzzles");
   static final Path           INDEX_FILE  = PUZZLE_DIR.resolve("index.json");
   static final OffsetDateTime now         = OffsetDateTime.now(ZoneOffset.UTC);
   static final String         CREATED_AT  = now.format(DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss'Z'"));
   static final String         FILE_ID     = CREATED_AT.replace(":", "-") + "_" + System.currentTimeMillis() / 1000;
   static final String         FILE_NAME   = FILE_ID + ".json";
   static final Path           OUTPUT_PATH = PUZZLE_DIR.resolve(FILE_NAME);
   static final String         DATE_STRING = now.toLocalDate().toString();
   
   static final AtomicLong TOTAL_NODES      = new AtomicLong(0);
   static final AtomicLong TOTAL_BACKTRACKS = new AtomicLong(0);
   static final AtomicLong TOTAL_ATTEMPTS   = new AtomicLong(0);
   static final AtomicLong TOTAL_SUCCESS    = new AtomicLong(0);
   static final AtomicLong TOTAL_SIMPLICITY = new AtomicLong(0); // Scaled by 100 for precision
   
   @Data
   public static class Opts {
      
      public int     seed          = (int) (System.nanoTime() ^ System.currentTimeMillis());
      public int     pop           = 18;
      public int     gens          = 2000;
      public String  wordsPath     = "nl_score_hints.csv";
      public double  minSimplicity = 0; // 0 means no limit
      public int     threads       = Math.max(1, Runtime.getRuntime().availableProcessors());
      public int     tries         = threads;
      public boolean reindex       = false;
      public int     fillTimeout   = 20_000;
      public boolean verbose       = false;
   }
   
   public void main(String[] args) {
      var opts = parseArgs(args);
      
      if (opts.reindex) {
         
         section("Reindex");
         info("OutputDir  : " + OUT_DIR);
         rebuildIndex();
         return;
      }
      
      section("Puzzle Generator");
      info("OutputDir  : " + OUT_DIR);
      info("WordsFile  : " + opts.wordsPath);
      
      section("Settings");
      printSettings(opts);
      
      var res = generatePuzzle(opts);
      if (res == null) {
         err("Search status : UNSOLVED");
         err("Reason        : No solution found within tries.");
         System.exit(1);
         return;
      }
      
      section("Result");
      info(String.format(Locale.ROOT, "simplicity  : %.2f", res.filled().simplicity()));
      
      section("Mask");
      System.out.print(indentLines(res.mask().gridToString(), "  "));
      
      section("Grid (raw)");
      System.out.print(indentLines(res.filled().grid().gridToString(), "  "));
      
      section("Grid (human)");
      System.out.print(indentLines(res.filled().grid().renderHuman(), "  "));
      
      var exported = exportFormatFromFilled(res, 1, new Rewards(50, 2, 1));
      
      section("Clues");
      info("status       : generating...");
      info("generatedFor : " + exported.words().length);
      info("status       : done");
      
      section("Words");
      printWordsTable(exported.words());
      
      section("Gridv2");
      for (var row : exported.gridv2()) System.out.println("  " + row);
      
      var theme = "algemeen";
      section("Export");
      info("file         : " + OUTPUT_PATH);
      
      try {
         Files.createDirectories(PUZZLE_DIR);
         var json = toJson(exported, DATE_STRING, theme);
         Files.writeString(OUTPUT_PATH, json, StandardCharsets.UTF_8);
         
         // Update index.json
         var pathInIndex = "/puzzles/" + FILE_NAME;
         var indexRecord = toIndexRecordJson(FILE_ID, pathInIndex, DATE_STRING, theme, exported.difficulty(), CREATED_AT);
         if (1 != 1) updateIndex(PUZZLE_DIR.toString(), indexRecord);
         else rebuildIndex();
         info("indexUpdated : " + INDEX_FILE);
      } catch (IOException e) {
         err("Failed to write: " + FILE_NAME);
         err("Reason       : " + e.getMessage());
         System.exit(2);
      }
   }
   
   // ---------------- Output helpers ----------------
   
   private static void info(String msg) { System.out.println("[INFO ] " + msg); }
   private static void warn(String msg) { System.out.println("[WARN ] " + msg); }
   private static void err(String msg)  { System.err.println("[ERROR] " + msg); }
   
   private static void section(String title) {
      System.out.println();
      System.out.println(title);
   }
   
   private static String envOrDefault(String key, String def) {
      var v = System.getenv(key);
      return v == null || v.isBlank() ? def : v;
   }
   
   private static void printSettings(Opts o) {
      System.out.printf(Locale.ROOT, "  %-14s: %d%n", "seed", o.seed);
      System.out.printf(Locale.ROOT, "  %-14s: %d%n", "population", o.pop);
      System.out.printf(Locale.ROOT, "  %-14s: %d%n", "generations", o.gens);
      System.out.printf(Locale.ROOT, "  %-14s: %s%n", "wordsPath", o.wordsPath);
      System.out.printf(Locale.ROOT, "  %-14s: %.2f%n", "minSimplicity", o.minSimplicity);
      System.out.printf(Locale.ROOT, "  %-14s: %d%n", "threads", o.threads);
   }
   
   private static String fmtPoint(int r, int c) { return String.format(Locale.ROOT, "(%d,%d)", r, c); }
   
   private static void printWordsTable(WordOut[] words) {
      System.out.println("  #  WORD         CX DIR  START     ARROW     CLUE");
      for (int j = 0; j < words.length; j++) {
         var w = words[j];
         System.out.printf(
                 Locale.ROOT,
                 "  %-2d %-12s %-4s %-3s %-9s %-9s %s%n",
                 j,
                 safe(w.word(), 12),
                 safe("" + w.complex(), 4),
                 safe(w.direction(), 3),
                 fmtPoint(w.startRow(), w.startCol()),
                 fmtPoint(w.arrowRow(), w.arrowCol()),
                 Arrays.toString(w.lemma().clue().toArray(String[]::new)));
      }
   }
   
   private static String safe(String s, int max) {
      if (s == null) return "";
      if (s.length() <= max) return s;
      return s.substring(0, Math.max(0, max - 1)) + "…";
   }
   
   static String indentLines(String s, String indent) {
      if (s == null || s.isEmpty()) return "";
      var lines = s.split("\\R", -1);
      var sb    = new StringBuilder();
      for (var line : lines) sb.append(indent).append(line).append('\n');
      return sb.toString();
   }
   
   static void usage() {
      System.out.println("""
                                 Usage:
                                   java puzzle.Main [--seed N] [--pop N] [--gens N] [--tries N] [--words FILE] [--min-simplicity N.N] [--threads N] [--reindex]
                                 
                                 Defaults:
                                   --pop 18
                                   --gens 500
                                   --words nl_score_hints.csv
                                   --min-simplicity 0 (no limit)
                                   --threads %d
                                 """.formatted(Math.max(1, Runtime.getRuntime().availableProcessors())));
   }
   
   static Opts parseArgs(String[] argv) {
      var out = new Opts();
      for (var i = 0; i < argv.length; i++) {
         var a = argv[i];
         var 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 if (a.equals("--reindex")) {
            out.reindex = true;
         } else if (a.equals("--verbose")) {
            out.verbose = true;
         } else {
            throw new IllegalArgumentException("Unknown arg: " + a);
         }
      }
      return out;
   }
   
   // ---------------- Generation ----------------
   
   // Package-private method for testing
   PuzzleResult generatePuzzle(Opts opts) {
      
      var tLoad0 = System.nanoTime();
      var dict   = loadWords(opts.wordsPath);
      var tLoad1 = System.nanoTime();
      
      section("Load");
      info(String.format(Locale.ROOT, "words        : %,d", dict.wordz().length));
      info(String.format(Locale.ROOT, "loadTime     : %.3f s", (tLoad1 - tLoad0) / 1e9));
      
      section("Search");
      var tSearch0 = System.currentTimeMillis();
      var deadline = System.currentTimeMillis() + 40_000;
      
      PuzzleResult resFinal = null;
      if (opts.threads > 1) {
         info("mode         : multi-threaded (" + opts.threads + ")");
         var executor          = Executors.newFixedThreadPool(opts.threads);
         var completionService = new ExecutorCompletionService<PuzzleResult>(executor);
         int submitted         = 0;
         
         try {
            // Keep at least some tasks in flight
            for (int i = 0; i < opts.threads; i++) {
               final int attempt = ++submitted;
               completionService.submit(() -> attempt(new Rng(opts.seed + attempt), dict, opts));
            }
            
            while (System.currentTimeMillis() < deadline) {
               var future = completionService.poll(deadline - System.currentTimeMillis(), TimeUnit.MILLISECONDS);
               if (future == null) break;
               
               var result = future.get();
               if (result != null) {
                  info("status       : SOLVED");
                  resFinal = result;
                  break;
               }
               
               // Submit another task if we still have time
               if (System.currentTimeMillis() < deadline) {
                  final int attempt = ++submitted;
                  completionService.submit(() -> attempt(new Rng(opts.seed + attempt), dict, opts));
               }
            }
            if (resFinal == null) warn("status       : UNSOLVED (timeout)");
         } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            warn("status       : INTERRUPTED");
         } catch (ExecutionException e) {
            warn("status       : ERROR (" + e.getMessage() + ")");
         } finally {
            executor.shutdownNow();
         }
         
      } else {
         info("mode         : single-threaded");
         var rng     = new Rng(opts.seed);
         int attempt = 0;
         
         while (System.currentTimeMillis() < deadline) {
            attempt++;
            info("try          : " + attempt + " (remaining: " + (deadline - System.currentTimeMillis()) / 1000 + "s)");
            
            var result = attempt(rng, dict, opts);
            if (result != null) {
               info("status       : SOLVED");
               info("foundAtTry   : " + attempt);
               resFinal = result;
               break;
            }
         }
         
         if (resFinal == null) info("status       : UNSOLVED (timeout)");
      }
      
      var tSearch1   = System.currentTimeMillis();
      var searchTime = (tSearch1 - tSearch0) / 1000.0;
      
      section("Performance");
      info(String.format(Locale.ROOT, "totalNodes   : %,d", TOTAL_NODES.get()));
      info(String.format(Locale.ROOT, "totalBacks   : %,d", TOTAL_BACKTRACKS.get()));
      info(String.format(Locale.ROOT, "searchTime   : %.2f s", searchTime));
      info(String.format(Locale.ROOT, "nodes/sec    : %,d", (int) (TOTAL_NODES.get() / Math.max(0.001, searchTime))));
      
      section("Material");
      info(String.format(Locale.ROOT, "attempts     : %,d", TOTAL_ATTEMPTS.get()));
      info(String.format(Locale.ROOT, "successRate  : %.1f%%", TOTAL_ATTEMPTS.get() == 0 ? 0 : TOTAL_SUCCESS.get() * 100.0 / TOTAL_ATTEMPTS.get()));
      if (TOTAL_SUCCESS.get() > 0) {
         info(String.format(Locale.ROOT, "avgSimplic   : %.2f", TOTAL_SIMPLICITY.get() / 100.0 / TOTAL_SUCCESS.get()));
      }
      info(String.format(Locale.ROOT, "dictWords    : %,d", dict.wordz().length));
      
      return resFinal;
   }
   
   static PuzzleResult attempt(Rng rng, Dict dict, Opts opts) {
      TOTAL_ATTEMPTS.incrementAndGet();
      var swe    = new SwedishGenerator();
      var mask   = swe.generateMask(rng, dict.lenCounts(), opts.pop, opts.gens, opts.verbose);
      var filled = swe.fillMask(rng, mask, dict.index(), 200, opts.fillTimeout, opts.verbose);
      
      TOTAL_NODES.addAndGet(filled.stats().nodes);
      TOTAL_BACKTRACKS.addAndGet(filled.stats().backtracks);
      if (filled.ok()) {
         TOTAL_SUCCESS.incrementAndGet();
         TOTAL_SIMPLICITY.addAndGet((long) (filled.simplicity() * 100));
      }
      
      var name       = Thread.currentThread().getName();
      var status     = filled.ok() ? "SUCCESS" : "FAILED";
      var simplicity = String.format(Locale.ROOT, "%.2f", filled.simplicity());
      var nps        = (int) (filled.stats().nodes / Math.max(0.001, filled.stats().seconds));
      
      System.out.printf(Locale.ROOT,
                        "[ATTEMPT] thread=%s | status=%s | nodes=%d | backtracks=%d | nps=%d | simplicity=%s | time=%.1fs%n",
                        name, status, filled.stats().nodes, filled.stats().backtracks, nps, simplicity, filled.stats().seconds
                       );
      
      if (filled.ok() && (opts.minSimplicity <= 0 || filled.simplicity() >= opts.minSimplicity)) {
         return new PuzzleResult(swe, dict, new Gridded(mask), filled);
      }
      
      if (opts.verbose && filled.ok()) {
         System.err.printf(Locale.ROOT,
                           "simplicity   : %.2f (below min %.2f)%n",
                           filled.simplicity(), opts.minSimplicity
                          );
      }
      return null;
   }
   
   // ---------------- Export (unchanged logic) ----------------
   
   private static String toJson(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().length; i++) {
         var w     = puzzle.words()[i];
         var clues = w.clue().toArray(String[]::new);
         Arrays.sort(clues, Comparator.comparingInt(String::length));
         sb.append("    {\n");
         sb.append("      \"word\": \"").append(escapeJson(w.word())).append("\",\n");
         sb.append("      \"clue\": [").append(Arrays.stream(clues).map(ss -> "\"" + escapeJson(ss) + "\"").collect(Collectors.joining(","))).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("      \"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().length - 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 toIndexRecordJson(String id, String path, String date, String theme, int difficulty, String createdAt) {
      return String.format(
              Locale.ROOT,
              "{\"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) {
      var indexPath = Paths.get(outDir, "index.json");
      try {
         var content = Files.exists(indexPath) ? Files.readString(indexPath, StandardCharsets.UTF_8).trim() : "";
         
         if (content.isEmpty() || content.equals("[]")) {
            content = "[\n  " + newRecordJson + "\n]";
         } else {
            var 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);
         info("indexUpdated : " + indexPath);
      } catch (IOException e) {
         err("Failed to update index.json: " + e.getMessage());
      }
   }
   
   private void rebuildIndex() {
      
      if (!Files.exists(PUZZLE_DIR)) {
         err("Puzzles directory does not exist: " + PUZZLE_DIR);
         return;
      }
      
      info("Rebuilding index from: " + PUZZLE_DIR);
      
      var records = new ArrayList<String>();
      try (var stream = Files.list(PUZZLE_DIR)) {
         stream.filter(p -> p.toString().endsWith(".json") && !p.getFileName().toString().equals("index.json"))
               .sorted(Comparator.comparing(Path::getFileName).reversed())
               .forEach(path -> {
                  try {
                     var filename = path.getFileName().toString();
                     var id       = filename.substring(0, filename.length() - 5);
                     var content  = Files.readString(path, StandardCharsets.UTF_8);
                     
                     var date       = extractValue(content, "date");
                     var theme      = extractValue(content, "theme");
                     var difficulty = 1;
                     try { difficulty = Integer.parseInt(extractValue(content, "difficulty")); } catch (Exception ignored) { }
                     
                     var createdAt = id;
                     if (id.length() >= 20 && id.charAt(10) == 'T') {
                        var parts  = id.split("_");
                        var dtPart = parts[0]; // 2025-12-24T04-25-06Z
                        if (dtPart.length() >= 19) {
                           createdAt = dtPart.substring(0, 13) + ":" + dtPart.substring(14, 16) + ":" + dtPart.substring(17);
                        }
                     }
                     
                     var pathInIndex = "/puzzles/" + filename;
                     records.add(toIndexRecordJson(id, pathInIndex, date, theme, difficulty, createdAt));
                  } catch (IOException e) {
                     err("Failed to read " + path + ": " + e.getMessage());
                  }
               });
      } catch (IOException e) {
         err("Failed to list puzzles: " + e.getMessage());
         return;
      }
      
      var indexPath = PUZZLE_DIR.resolve("index.json");
      var content   = "[\n  " + String.join(",\n  ", records) + "\n]";
      try {
         Files.writeString(indexPath, content, StandardCharsets.UTF_8);
         info("Successfully rebuilt index.json with " + records.size() + " records.");
      } catch (IOException e) {
         err("Failed to write index.json: " + e.getMessage());
      }
   }
   
   private static String extractValue(String json, String key) {
      var pattern = java.util.regex.Pattern.compile("\"" + key + "\":\\s*\"?([^\",\\n\\r}]*)\"?");
      var matcher = pattern.matcher(json);
      if (matcher.find()) return matcher.group(1).trim();
      return "";
   }
}