import { awaitEvent, emitEvent, listenTo } from "./events.js";
import {
  getPlannerCellDom,
  getPlannerCellsDom,
  getPlannerDom,
  getPlannerMultiblockSize,
  getPlannerMultiblockType
} from "./planner.js";
import { checkEquivalence, findFullEntry, findFullEntryFromAtlasMetadata } from "./ncpf-utils.js";
import { clearCellMetadata, deleteMetadata, getMetadata, setCellMetadata, setDomMetadata } from "./metadata.js";
import { registerLoadCallback } from "./main.js";
import { getAtlasMetadata } from "./atlas.js";
import { initialize3dArray } from "./js-utils.js";

const SFR_MODERATOR_MAX_LEN = 3;
const SFR_MODERATOR_REFLECTOR_MAX_LEN = 2;

registerLoadCallback(() => {
  listenTo("fileLoad:end", runSolver);
  listenTo("planner:cell:change", runSolver);
});

export function resetSolver(design = undefined, designConfiguration) {
  const planner = getPlannerDom();
  if (design && designConfiguration) {
    planner.dataset.type = design.type;
    if (design.type === "nuclearcraft:overhaul_sfr") {
      // design has coolant_recipe and block_recipes
      const designCoolantRecipe = designConfiguration.coolant_recipes[design.coolant_recipe];
      const fullCoolantRecipe = findFullEntry(
          "coolant_recipes",
          designCoolantRecipe,
          { configurationHint: design.type },
      );
      setDomMetadata(planner, { coolant_recipe: fullCoolantRecipe });

      const designBlockRecipes = design.block_recipes.flat(3);
      for (let x = 0; x < design.dimensions[0]; x++) {
        for (let y = 0; y < design.dimensions[1]; y++) {
          for (let z = 0; z < design.dimensions[2]; z++) {
            const block = designConfiguration.blocks[design.design[x][y][z]];
            const cellDom = getPlannerCellDom(x, y, z);
            deleteMetadata(cellDom);
            if (!block || !block.modules || !block.modules["ncpf:block_recipes"]) continue;
            const blockRecipeIndex = designBlockRecipes.shift();
            const blockRecipe = block.modules["ncpf:block_recipes"].recipes[blockRecipeIndex];
            setCellMetadata(cellDom, "block_recipe", blockRecipe);
          }
        }
      }
    }
  }

  emitEvent("solver:reset:done");
}

const sixAxes = {
  "-X": { x: -1, y: 0, z: 0 },
  "-Y": { x: 0, y: -1, z: 0 },
  "-Z": { x: 0, y: 0, z: -1 },
  "+X": { x: 1, y: 0, z: 0 },
  "+Y": { x: 0, y: 1, z: 0 },
  "+Z": { x: 0, y: 0, z: 1 },
};
function *travelUntil(grid, { x, y, z }, {x: dirX = 0, y: dirY = 0, z: dirZ = 0}, shouldStop = (pos) => false) {
  x = Number(x);
  y = Number(y);
  z = Number(z);
  while (!shouldStop({x, y, z})
        && x > 0 && y > 0 && z > 0
        && x < grid.length - 1 && y < grid[0].length - 1 && z < grid[0][0].length - 1) {
    x += dirX;
    y += dirY;
    z += dirZ;
    yield grid[x][y][z];
  }
}
function neighbor(grid, { x, y, z }, {x: dirX = 0, y: dirY = 0, z: dirZ = 0}) {
  x = Number(x);
  y = Number(y);
  z = Number(z);
  if (x > 0 && y > 0 && z > 0
        && x < grid.length - 1 && y < grid[0].length - 1 && z < grid[0][0].length - 1) {
    x += dirX;
    y += dirY;
    z += dirZ;
    return grid[x][y][z];
  }
}

async function debugBlockSolver(data) {
  const planner = getPlannerDom();

  if (planner.dataset.debugSolver) {
    emitEvent("solver:pause", data);
    await awaitEvent("solver:continue");
  }
}

export async function runSolver() {
  const startTime = performance.now()
  emitEvent("solver:begin");

  const multiblockType = getPlannerMultiblockType();

  const plannerBlocks = getPlannerCellsDom()
      .map(b => ({ dom: b, ...b.dataset }))
      .filter(b => b.tile !== "-1")
      .map(b => (b.metadata ? { ...b, metadata: getMetadata(b.metadata) } : b));
  const tileMap = Object.fromEntries([...new Set(plannerBlocks.map(b => b.tile))]
      .map(idx => [idx, getAtlasMetadata(idx)])
      .map(([idx, md]) => [idx, findFullEntry("blocks", md, { configurationHint: multiblockType })]));
  plannerBlocks.forEach(v => v.tile = tileMap[v.tile]);
  plannerBlocks.forEach(v => clearCellMetadata(v.dom, "active"));

  const plannerBlocksGrid = initialize3dArray(...getPlannerMultiblockSize());
  plannerBlocks.forEach(b => {
    plannerBlocksGrid[b.x][b.y][b.z] = b;
  });

  if (multiblockType === "nuclearcraft:overhaul_sfr") {
    const fuelCells = plannerBlocks.filter(b => b.tile.modules && b.tile.modules["nuclearcraft:overhaul_sfr:fuel_cell"]);
    // const neutronSources = plannerBlocks.filter(b => b.tile.modules && b.tile.modules["nuclearcraft:overhaul_sfr:neutron_source"]);
    // const reflectors = plannerBlocks.filter(b => b.tile.modules && b.tile.modules["nuclearcraft:overhaul_sfr:reflector"]);
    // const shields = plannerBlocks.filter(b => b.tile.modules && b.tile.modules["nuclearcraft:overhaul_sfr:neutron_shield"]);
    // const moderators = plannerBlocks.filter(b => b.tile.modules && b.tile.modules["nuclearcraft:overhaul_sfr:moderator"]);
    const heatSinks = plannerBlocks.filter(b => b.tile.modules && b.tile.modules["nuclearcraft:overhaul_sfr:heat_sink"]);
    // const irradiators = plannerBlocks.filter(b => b.tile.modules && b.tile.modules["nuclearcraft:overhaul_sfr:irradiator"]);
    // const conductors = plannerBlocks.filter(b => b.tile.modules && b.tile.modules["nuclearcraft:overhaul_sfr:conductor"]);

    fuelCells.forEach(cell => {
      const recipeList = cell.tile.modules["ncpf:block_recipes"].recipes;
      cell.recipe = recipeList.find(r => checkEquivalence(r, cell.metadata?.block_recipe));
      // TODO: handle blank recipe
      if (!cell.recipe) return; // FIXME
      cell.axes = Object.values(sixAxes).map(dir => [...travelUntil(
          plannerBlocksGrid,
          cell,
          dir,
          ({x, y, z}) => {
            const block = plannerBlocksGrid[x][y][z];
            if (block === cell) {
              return false;
            }
            if (!block) {
              return false;
            }
            if (block.tile.modules) {
              if (block.tile.modules["nuclearcraft:overhaul_sfr:fuel_cell"]) return true;
              if (block.tile.modules["nuclearcraft:overhaul_sfr:reflector"]) return true;
              if (block.tile.modules["nuclearcraft:overhaul_sfr:irradiator"]) return true;
            }
            return false;
          }
      )]);
      if (cell.recipe.modules["nuclearcraft:overhaul_sfr:fuel_stats"].self_priming) {
        cell.primer = cell;
        cell.active = true;
      } else if (cell.axes.some(a => (a[a.length - 1]?.tile?.modules ?? {})["nuclearcraft:overhaul_sfr:neutron_source"])) {
        const primerAxis = cell.axes.find(a => (a[a.length - 1]?.tile?.modules ?? {})["nuclearcraft:overhaul_sfr:neutron_source"]);
        cell.primer = primerAxis[primerAxis.length - 1];
        cell.active = true;
      }
      setCellMetadata(cell.dom, "primer", cell.primer);
    });

    await debugBlockSolver("primed");

    const seenModerators = new Set();
    const activeCells = fuelCells.filter(c => c.active);
    let activeCellIndex = 0;
    while (activeCellIndex < activeCells.length) {
      const cell = activeCells[activeCellIndex];
      if (!cell.axes) return; // ignore bad cells
      Object.keys(sixAxes).map((directionName, i) => {
        let valid = false;
        let moderatorPath = [];
        let length = 0;
        let reflectMultiplier = 1;
        let reflect = false;
        let fluxDestination;
        for (const block of cell.axes[i]) {
          if (!block) break;

          if (block.tile.modules["nuclearcraft:overhaul_sfr:moderator"]) {
            moderatorPath.push(block);
            seenModerators.add(block);
            length += 1;
            continue;
          }
          if (block.tile.modules["nuclearcraft:overhaul_sfr:neutron_shield"]) {
            length += 1;
            continue;
          }
          if (block.tile.modules["nuclearcraft:overhaul_sfr:fuel_cell"] || block.tile.modules["nuclearcraft:overhaul_sfr:irradiator"]) {
            valid = length <= SFR_MODERATOR_MAX_LEN;
            fluxDestination = block;
          }
          if (block.tile.modules["nuclearcraft:overhaul_sfr:reflector"]) {
            valid = length <= SFR_MODERATOR_REFLECTOR_MAX_LEN;
            reflect = true;
            reflectMultiplier = block.tile.modules["nuclearcraft:overhaul_sfr:reflector"].reflectivity;
            console.log(reflectMultiplier)
            fluxDestination = cell;
          }
          break;
        }
        if (moderatorPath.length < 1) {
          valid = false;
        }
        if (!valid) return;

        moderatorPath[0].active = true;
        let totalFlux = 0;
        for (const mod of moderatorPath) {
          totalFlux += mod.tile.modules["nuclearcraft:overhaul_sfr:moderator"].flux;
          if (!mod.flux) {
            mod.flux = {};
          }
          mod.flux[directionName] = totalFlux;
        }
        if (reflect) {
          totalFlux *= reflectMultiplier;
          const backAxis = directionName.replace(/^[+-]/, s => s === "+" ? "-" : "+");
          for (const mod of moderatorPath.toReversed()) {
            totalFlux += mod.tile.modules["nuclearcraft:overhaul_sfr:moderator"].flux * reflectMultiplier;
            mod.flux[backAxis] = totalFlux;
          }
        }
        fluxDestination.flux = (fluxDestination.flux ?? 0) + totalFlux;

        if (fluxDestination.tile.modules["nuclearcraft:overhaul_sfr:fuel_cell"] && !fluxDestination.active) {
          if (!fluxDestination.recipe) return; // FIXME
          const fuelStats = fluxDestination.recipe.modules["nuclearcraft:overhaul_sfr:fuel_stats"];
          if (fuelStats.criticality <= fluxDestination.flux) {
            fluxDestination.active = true;
            activeCells.push(fluxDestination);
          }
        }
      });

      activeCellIndex++;
    }
    seenModerators.forEach(moderator => {
      setCellMetadata(moderator.dom, "flux", moderator.flux);
      if (moderator.active) {
        setCellMetadata(moderator.dom, "active", moderator.active);
      }
    });
    fuelCells.forEach(cell => {
      if (cell.flux) {
        setCellMetadata(cell.dom, "flux", cell.flux);
      }
      setCellMetadata(cell.dom, "active", cell.active ?? false);
    });

    await debugBlockSolver("moderators");

    heatSinks.forEach(hs => {
      hs.neighbors = {};
      Object.entries(sixAxes).forEach(([name, dir]) => {
        hs.neighbors[name] = neighbor(plannerBlocksGrid, hs, dir);
      });
      hs.rules = hs.tile.modules["nuclearcraft:overhaul_sfr:heat_sink"].rules;
    });
    let heatSinkProcessed = false;
    const unprocessedHeatSinks = new Set(heatSinks);
    do {
      heatSinkProcessed = false;
      unprocessedHeatSinks.forEach(hs => {
        if (hs.active) {
          console.log("active hs in queue", hs);
          return;
        }
        const activeNeighbors = Object.entries(hs.neighbors).filter(([_, n]) => {
          if (!n) return false;
          if (n.active) return true;
          if (n.tile.modules) {
            return n.tile.modules["nuclearcraft:overhaul_sfr:reflector"]
                || n.tile.modules["nuclearcraft:overhaul_sfr:irradiator"]
                || n.tile.modules["nuclearcraft:overhaul_sfr:casing"]
                || n.tile.modules["nuclearcraft:overhaul_sfr:neutron_shield"];
          }
          return false;
        });
        if (activeNeighbors.length === 0) return;

        let valid = false;
        for (const rule of hs.rules) {
          if (checkHeatSinkRule(hs, rule, activeNeighbors)) {
            valid = true;
            break;
          }
        }
        if (valid) {
          heatSinkProcessed = true;
          unprocessedHeatSinks.delete(hs);
          hs.active = true;
          setCellMetadata(hs.dom, "active", true);
        }
      });
    } while (heatSinkProcessed);
    unprocessedHeatSinks.forEach(hs => {
      setCellMetadata(hs.dom, "active", false);
    });
    console.log(unprocessedHeatSinks);

    await debugBlockSolver("heatsinks");
  }

  const endTime = performance.now()
  console.log(`Solver took ${endTime - startTime}ms`);
}

function checkHeatSinkRule(hs, rule, activeNeighbors) {
  switch (rule.type) {
    case "and":
      return rule.rules.every(r => checkHeatSinkRule(hs, r, activeNeighbors));

    case "or":
      return rule.rules.some(r => checkHeatSinkRule(hs, r, activeNeighbors));

    case "between": {
      const blocks = activeNeighbors
          .map(([,b]) => b.tile)
          .filter(b => {
            if (rule.block.type === "module") {
              return b.modules && rule.block.name in b.modules;
            }
            return checkEquivalence(rule.block, b);
          });
      return rule.min <= blocks.length && blocks.length <= rule.max;
    }
    case "axial": {
      const axialPairs = activeNeighbors
          .filter(([,b]) => {
            if (rule.block.type === "module") {
              return b.tile.modules && rule.block.name in b.tile.modules;
            }
            return checkEquivalence(rule.block, b.tile);
          })
          .reduce((acc, [dir, block]) => {
            acc[dir[1]] = acc[dir[1]] ?? [];
            acc[dir[1]].push(block);
            return acc;
          }, {});
      const count = Object.values(axialPairs).filter(p => p.length === 2).length;
      return rule.min <= count && count <= rule.max;
    }
    default:
      console.warn("unknown rule: %o", rule);
      return false;
  }
}