import { calculateBeaconDesign, calculateIsolatedBeaconDesign } from "./beacons.ts";
import {
  MODULE_BY_KEY, MODULE_PRIORITY, QUALITIES, QUALITY_BY_ID, displayName, iconFor,
} from "./factorio-data.ts";
import type {
  BeaconLayoutResult, CalculationResult, CalculatorInput, EffectKind, EffectTotals,
  Machine, MaterialFlow, ModuleFamily, ModuleSelection, QualityAmount, QualityId,
  RateUnit, Recipe,
} from "./types.ts";

const BEACON_PROFILE = [
  1, .7071, .5773, .5, .4472, .4082, .3779, .3535, .3333, .3162,
  .3015, .2886, .2773, .2672, .2581, .25, .2425, .2357, .2294, .2236,
];
const BEACON_HEATING_KW = 400;
const EMPTY_EFFECTS: EffectTotals = { consumption: 0, speed: 0, productivity: 0, pollution: 0, quality: 0 };
const UNIT_SECONDS: Record<RateUnit, number> = { s: 1, min: 60, h: 3600 };

function copyEffects(): EffectTotals { return { ...EMPTY_EFFECTS }; }
function addEffects(target: EffectTotals, source: Partial<Record<EffectKind, number>>, factor = 1) {
  (Object.keys(EMPTY_EFFECTS) as EffectKind[]).forEach((effect) => {
    target[effect] += (source[effect] ?? 0) * factor;
  });
}

export function moduleIsAllowed(selection: ModuleSelection, machine: Machine, recipe: Recipe, beacon = false): boolean {
  if (beacon) return selection.family === "speed" || selection.family === "efficiency";
  if (selection.family === "productivity") return machine.allowedEffects.includes("productivity") && recipe.allowProductivity;
  if (selection.family === "quality") return machine.allowedEffects.includes("quality") && recipe.allowQuality;
  if (selection.family === "speed") return machine.allowedEffects.includes("speed");
  return machine.allowedEffects.includes("consumption");
}
export function resolveAutoModules(
  families: ModuleFamily[], tier: 1 | 2 | 3, quality: QualityId, slots: number,
  machine: Machine, recipe: Recipe, beacon = false,
): Array<ModuleSelection | null> {
  const family = MODULE_PRIORITY.find((candidate) =>
    families.includes(candidate) && moduleIsAllowed({ family: candidate, tier, quality }, machine, recipe, beacon));
  return Array.from({ length: slots }, () => family ? { family, tier, quality } : null);
}
export function effectsForModules(
  selections: Array<ModuleSelection | null>, machine: Machine, recipe: Recipe, beacon = false,
): EffectTotals {
  const totals = copyEffects();
  for (const selection of selections) {
    if (!selection || !moduleIsAllowed(selection, machine, recipe, beacon)) continue;
    const definition = MODULE_BY_KEY.get(`${selection.family}:${selection.tier}`);
    const quality = QUALITY_BY_ID.get(selection.quality);
    if (!definition || !quality) continue;
    const scaled: Partial<Record<EffectKind, number>> = {};
    for (const [effect, rawValue] of Object.entries(definition.effects) as Array<[EffectKind, number]>) {
      const primary =
        (effect === "consumption" && rawValue < 0) ||
        (effect === "speed" && rawValue > 0) ||
        (effect === "productivity" && rawValue > 0) ||
        (effect === "pollution" && rawValue < 0) ||
        (effect === "quality" && rawValue > 0);
      scaled[effect] = rawValue * (primary ? quality.multiplier : 1);
    }
    addEffects(totals, scaled);
  }
  return totals;
}
export function beaconEffects(
  count: number, modules: Array<ModuleSelection | null>, machine: Machine, recipe: Recipe, beaconQuality: QualityId,
): EffectTotals {
  const totals = copyEffects();
  if (count <= 0) return totals;
  const effects = effectsForModules(modules, machine, recipe, true);
  const quality = QUALITY_BY_ID.get(beaconQuality) ?? QUALITY_BY_ID.get("normal")!;
  const profile = BEACON_PROFILE[Math.min(count, BEACON_PROFILE.length) - 1] ?? 1 / Math.sqrt(count);
  addEffects(totals, effects, count * profile * (1.5 + .2 * quality.level));
  totals.productivity = 0;
  return totals;
}

export function qualityDistribution(baseQuality: QualityId, chance: number, maxQuality: QualityId = "legendary"): QualityAmount[] {
  const baseIndex = QUALITIES.findIndex((quality) => quality.id === baseQuality);
  const capIndex = Math.max(baseIndex, QUALITIES.findIndex((quality) => quality.id === maxQuality));
  if (baseIndex < 0) return [{ quality: "normal", amount: 1 }];
  if (baseIndex >= capIndex) return [{ quality: QUALITIES[baseIndex].id, amount: 1 }];
  const probability = Math.max(0, Math.min(1, chance));
  const result: QualityAmount[] = [{ quality: QUALITIES[baseIndex].id, amount: 1 - probability }];
  let remaining = probability;
  for (let index = baseIndex + 1; index <= capIndex; index += 1) {
    const atCap = index === capIndex;
    const amount = atCap ? remaining : remaining * .9;
    result.push({ quality: QUALITIES[index].id, amount });
    remaining *= .1;
  }
  return result.filter((entry) => entry.amount > 1e-12);
}

type CycleFlow = { inputs: MaterialFlow[]; outputs: MaterialFlow[] };
function buildCycleFlows(
  recipe: Recipe, targetId: string, targetType: "item" | "fluid", productivityBonus: number,
  qualityChance: number, inputQuality: QualityId, maxQuality: QualityId, operatingFuel?: { id: string; amount: number },
): CycleFlow {
  const inputAmounts = new Map<string, number>();
  const outputAmounts = new Map<string, Map<QualityId | null, number>>();
  const grossOutputs = new Map<string, number>();
  for (const ingredient of recipe.ingredients) {
    const key = `${ingredient.type}:${ingredient.name}`;
    inputAmounts.set(key, (inputAmounts.get(key) ?? 0) + ingredient.amount);
  }
  if (operatingFuel && operatingFuel.amount > 0) {
    const key = `item:${operatingFuel.id}`;
    inputAmounts.set(key, (inputAmounts.get(key) ?? 0) + operatingFuel.amount);
  }
  const distribution = qualityDistribution(inputQuality, recipe.allowQuality ? qualityChance : 0, maxQuality);
  for (const product of recipe.products) {
    const key = `${product.type}:${product.name}`;
    const eligible = Math.max(0, product.amount - product.ignoredByProductivity);
    const produced = (product.amount + eligible * productivityBonus) * product.probability;
    grossOutputs.set(key, (grossOutputs.get(key) ?? 0) + produced);
    const qualities = outputAmounts.get(key) ?? new Map<QualityId | null, number>();
    if (product.type === "fluid") qualities.set(null, (qualities.get(null) ?? 0) + produced);
    else distribution.forEach((entry) => qualities.set(entry.quality, (qualities.get(entry.quality) ?? 0) + produced * entry.amount));
    outputAmounts.set(key, qualities);
  }
  const inputs: MaterialFlow[] = [];
  for (const [key, grossInput] of inputAmounts) {
    const separator = key.indexOf(":");
    const type = key.slice(0, separator) as "item" | "fluid";
    const name = key.slice(separator + 1);
    const baseQuality = type === "fluid" ? null : inputQuality;
    const returned = Math.min(grossInput, outputAmounts.get(key)?.get(baseQuality) ?? 0);
    const net = Math.max(0, grossInput - returned);
    if (net > 1e-12) inputs.push({
      id: name, name: displayName(name), type, icon: iconFor(name, type),
      amounts: [{ quality: baseQuality, amount: net }], total: net, gross: grossInput, returned,
    });
  }
  const outputs: MaterialFlow[] = [];
  for (const [key, qualities] of outputAmounts) {
    const separator = key.indexOf(":");
    const type = key.slice(0, separator) as "item" | "fluid";
    const name = key.slice(separator + 1);
    const baseQuality = type === "fluid" ? null : inputQuality;
    const input = inputAmounts.get(key) ?? 0;
    if (input > 0) qualities.set(baseQuality, (qualities.get(baseQuality) ?? 0) - input);
    const amounts = [...qualities.entries()]
      .map(([quality, amount]) => ({ quality, amount: Math.max(0, amount) }))
      .filter((entry) => entry.amount > 1e-12);
    const total = amounts.reduce((sum, entry) => sum + entry.amount, 0);
    if (total > 1e-12) outputs.push({
      id: name, name: displayName(name), type, icon: iconFor(name, type), amounts, total,
      gross: grossOutputs.get(key) ?? total, returned: Math.min(input, grossOutputs.get(key) ?? 0),
    });
  }
  void targetId; void targetType;
  return { inputs, outputs };
}
function mergeCycleFlows(weighted: Array<{ flow: CycleFlow; weight: number }>): CycleFlow {
  const merge = (side: "inputs" | "outputs") => {
    const flows = new Map<string, MaterialFlow>();
    for (const { flow, weight } of weighted) {
      for (const entry of flow[side]) {
        const key = `${entry.type}:${entry.id}`;
        const current = flows.get(key) ?? {
          ...entry,
          amounts: [],
          total: 0,
          gross: entry.gross === undefined ? undefined : 0,
          returned: entry.returned === undefined ? undefined : 0,
        };
        current.total += entry.total * weight;
        if (entry.gross !== undefined)
          current.gross = (current.gross ?? 0) + entry.gross * weight;
        if (entry.returned !== undefined)
          current.returned = (current.returned ?? 0) + entry.returned * weight;
        for (const amount of entry.amounts) {
          const existing = current.amounts.find(
            (candidate) => candidate.quality === amount.quality,
          );
          if (existing) existing.amount += amount.amount * weight;
          else current.amounts.push({ ...amount, amount: amount.amount * weight });
        }
        flows.set(key, current);
      }
    }
    return [...flows.values()].map((flow) => ({
      ...flow,
      amounts: flow.amounts.filter((entry) => entry.amount > 1e-12),
    }));
  };
  return { inputs: merge("inputs"), outputs: merge("outputs") };
}

export function normalizedInputQualityDistribution(
  input: Pick<CalculatorInput, "inputQuality" | "inputQualityDistribution" | "maxQuality">,
): QualityAmount[] {
  const baseIndex = Math.max(0, QUALITIES.findIndex((quality) => quality.id === input.inputQuality));
  const requestedCap = QUALITIES.findIndex(
    (quality) => quality.id === (input.maxQuality ?? "legendary"),
  );
  const capIndex = Math.max(baseIndex, requestedCap < 0 ? QUALITIES.length - 1 : requestedCap);
  const totals = new Map<QualityId, number>();
  for (const entry of input.inputQualityDistribution ?? []) {
    if (entry.quality === null || !Number.isFinite(entry.amount) || entry.amount <= 0)
      continue;
    const index = QUALITIES.findIndex((quality) => quality.id === entry.quality);
    if (index < baseIndex || index > capIndex) continue;
    totals.set(entry.quality, (totals.get(entry.quality) ?? 0) + entry.amount);
  }
  const total = [...totals.values()].reduce((sum, amount) => sum + amount, 0);
  if (total <= 1e-12)
    return [{ quality: input.inputQuality, amount: 1 }];
  return QUALITIES.slice(baseIndex, capIndex + 1)
    .map((quality) => ({
      quality: quality.id,
      amount: (totals.get(quality.id) ?? 0) / total,
    }))
    .filter((entry) => entry.amount > 1e-12);
}

function buildCompositeCycleFlows(
  input: CalculatorInput,
  productivityBonus: number,
  qualityChance: number,
  operatingFuel?: { id: string; amount: number },
): CycleFlow {
  const distribution = normalizedInputQualityDistribution(input);
  return mergeCycleFlows(
    distribution.map((entry) => ({
      weight: entry.amount,
      flow: buildCycleFlows(
        input.recipe,
        input.product.id,
        input.product.type,
        productivityBonus,
        qualityChance,
        entry.quality as QualityId,
        input.maxQuality ?? "legendary",
        operatingFuel,
      ),
    })),
  );
}
function scaledFlows(flows: MaterialFlow[], multiplier: number): MaterialFlow[] {
  return flows.map((flow) => ({
    ...flow, total: flow.total * multiplier,
    gross: flow.gross === undefined ? undefined : flow.gross * multiplier,
    returned: flow.returned === undefined ? undefined : flow.returned * multiplier,
    amounts: flow.amounts.map((entry) => ({ ...entry, amount: entry.amount * multiplier })),
  }));
}
function makeLayout(input: CalculatorInput, count: number): BeaconLayoutResult {
  return input.isolatedBeacons
    ? calculateIsolatedBeaconDesign(input.machine, count, input.beaconCount)
    : calculateBeaconDesign(input.machine, count, input.beaconCount);
}
function layoutCapacity(layout: BeaconLayoutResult, input: CalculatorInput, machineEffects: EffectTotals): number {
  const machineQuality = QUALITY_BY_ID.get(input.machineQuality) ?? QUALITY_BY_ID.get("normal")!;
  const capacity = (count: number) => {
    const transmitted = beaconEffects(count, input.beaconModules, input.machine, input.recipe, input.beaconQuality);
    const speedFactor = Math.max(.2, 1 + machineEffects.speed + transmitted.speed);
    return input.machine.speed * machineQuality.multiplier * speedFactor / input.recipe.energy;
  };
  if (!layout.counts.length) return 0;
  const sample = layout.counts.reduce((sum, count) => sum + capacity(count), 0);
  return layout.extrapolated ? sample / layout.counts.length * layout.machineCount : sample;
}
function averageEnergyModifier(layout: BeaconLayoutResult, input: CalculatorInput, machineEffects: EffectTotals): number {
  if (!layout.counts.length) return Math.max(.2, 1 + machineEffects.consumption);
  return layout.counts.reduce((sum, count) => {
    const transmitted = beaconEffects(count, input.beaconModules, input.machine, input.recipe, input.beaconQuality);
    return sum + Math.max(.2, 1 + machineEffects.consumption + transmitted.consumption);
  }, 0) / layout.counts.length;
}
function findMachineCount(initial: number, cycles: number, input: CalculatorInput, machineEffects: EffectTotals) {
  let count = Math.max(1, initial);
  let layout = makeLayout(input, count);
  let capacity = layoutCapacity(layout, input, machineEffects);
  for (let i = 0; i < 20 && capacity + 1e-10 < cycles; i += 1) {
    count = Math.max(count + 1, capacity > 0 ? Math.ceil(count * cycles / capacity) : count + 1);
    layout = makeLayout(input, count);
    capacity = layoutCapacity(layout, input, machineEffects);
  }
  return { count, layout, capacity };
}
function invalidResult(message: string): CalculationResult {
  return {
    valid: false, message, cyclesPerSecond: 0, ratePerSecond: 0, inputs: [], outputs: [], targetPerCycle: 0,
    machineEffects: copyEffects(), isolatedBeaconEffects: copyEffects(), craftingSpeed: 0, productivityBonus: 0,
    qualityChance: 0, theoreticalMachines: 0, machinesToBuild: 0, utilization: 0, idealIsolatedMachines: 0,
    beaconLayout: calculateBeaconDesign({ maxBeacons: 0, width: 1, height: 1 } as Machine, 0, 0),
    electricTargetKw: 0, electricFullLoadKw: 0, machineElectricTargetKw: 0, beaconElectricKw: 0,
    fuelTargetKw: 0, nutrientItemsPerSecond: 0, idleDrainKw: 0, heatingKw: 0,
    qualityDistribution: [], warnings: [],
  };
}

export function calculate(input: CalculatorInput): CalculationResult {
  if (!Number.isFinite(input.amount) || input.amount < 0) return invalidResult("Enter a non-negative amount.");
  const machineEffects = effectsForModules(input.machineModules, input.machine, input.recipe);
  const isolatedBeaconEffects = beaconEffects(input.beaconCount, input.beaconModules, input.machine, input.recipe, input.beaconQuality);
  const productivityBonus = Math.min(3, input.machine.baseProductivity +
    (input.recipe.allowProductivity ? machineEffects.productivity : 0));
  const calculationMode = input.calculationMode ?? "normal";
  const qualityChance = calculationMode === "quality" && input.recipe.allowQuality
    ? Math.max(0, Math.min(1, machineEffects.quality + isolatedBeaconEffects.quality))
    : 0;
  const machineQuality = QUALITY_BY_ID.get(input.machineQuality) ?? QUALITY_BY_ID.get("normal")!;
  const isolatedSpeedFactor = Math.max(.2, 1 + machineEffects.speed + isolatedBeaconEffects.speed);
  const isolatedCraftingSpeed = input.machine.speed * machineQuality.multiplier * isolatedSpeedFactor;
  const isolatedEnergyModifier = Math.max(.2, 1 + machineEffects.consumption + isolatedBeaconEffects.consumption);
  const operatingFuel = input.machine.fuelItem && input.machine.fuelValueKj && isolatedCraftingSpeed > 0
    ? {
        id: input.machine.fuelItem,
        amount: input.machine.energyKw * isolatedEnergyModifier * input.recipe.energy /
          isolatedCraftingSpeed / input.machine.fuelValueKj,
      }
    : undefined;
  const flows = buildCompositeCycleFlows(
    input,
    productivityBonus,
    qualityChance,
    operatingFuel,
  );
  const targetFlow = flows.outputs.find((flow) => flow.id === input.product.id && flow.type === input.product.type);
  if (!targetFlow) return invalidResult("This route does not create a positive net amount of the selected product.");
  const targetPerCycle = input.product.type === "fluid" || input.outputQuality === "total"
    ? targetFlow.total
    : targetFlow.amounts.find((entry) => entry.quality === input.outputQuality)?.amount ?? 0;
  if (targetPerCycle <= 1e-12) return invalidResult(`This setup produces no ${input.outputQuality} ${input.product.name}.`);

  let cyclesPerSecond = 0;
  let limitingInput: string | undefined;
  if (input.mode === "output") cyclesPerSecond = input.amount / UNIT_SECONDS[input.unit] / targetPerCycle;
  else {
    const limits = input.inputConstraints.flatMap((constraint) => {
      const flow = flows.inputs.find((candidate) => `${candidate.type}:${candidate.id}` === constraint.ingredientKey);
      return flow && flow.total > 0 && constraint.amount >= 0
        ? [{ cycles: constraint.amount / UNIT_SECONDS[input.unit] / flow.total, name: flow.name }]
        : [];
    });
    if (!limits.length) return invalidResult("Add at least one net input constraint for this route.");
    const limit = limits.reduce((a, b) => b.cycles < a.cycles ? b : a);
    cyclesPerSecond = limit.cycles;
    limitingInput = limit.name;
  }
  if (!Number.isFinite(cyclesPerSecond) || cyclesPerSecond < 0) return invalidResult("The selected rate cannot be calculated.");
  const ratePerSecond = cyclesPerSecond * targetPerCycle;
  const perMachineCycles = isolatedCraftingSpeed / input.recipe.energy;
  const theoreticalMachines = perMachineCycles > 0 ? cyclesPerSecond / perMachineCycles : 0;
  const isolatedBuild = cyclesPerSecond > 0 ? Math.max(1, Math.ceil(theoreticalMachines - 1e-12)) : 0;
  let machinesToBuild = isolatedBuild;
  let beaconLayout = makeLayout(input, machinesToBuild);
  let capacity = layoutCapacity(beaconLayout, input, machineEffects);
  if (!input.isolatedBeacons && input.beaconCount > 0 && machinesToBuild > 1 && capacity + 1e-10 < cyclesPerSecond) {
    const found = findMachineCount(machinesToBuild, cyclesPerSecond, input, machineEffects);
    machinesToBuild = found.count; beaconLayout = found.layout; capacity = found.capacity;
  }
  const utilization = capacity > 0 ? Math.min(1, cyclesPerSecond / capacity) : 0;
  const energyModifier = averageEnergyModifier(beaconLayout, input, machineEffects);
  const activeFullKw = input.machine.energyKw * energyModifier * machinesToBuild;
  const idleDrainKw = (input.machine.drainKw ?? (input.machine.energyKind === "electric" ? input.machine.energyKw / 30 : 0)) * machinesToBuild;
  const machineElectricTargetKw = input.machine.energyKind === "electric" ? activeFullKw * utilization + idleDrainKw : 0;
  const machineElectricFullKw = input.machine.energyKind === "electric" ? activeFullKw + idleDrainKw : 0;
  const beaconQuality = QUALITY_BY_ID.get(input.beaconQuality) ?? QUALITY_BY_ID.get("normal")!;
  const beaconElectricKw = beaconLayout.uniqueBeacons * 480 * beaconQuality.beaconPowerMultiplier;
  const fuelTargetKw = input.machine.energyKind === "electric" || input.machine.energyKind === "passive" ? 0 : activeFullKw * utilization;
  const nutrientItemsPerSecond = input.machine.energyKind === "nutrients" ? fuelTargetKw / 2000 : 0;
  const heatingKw = (input.machine.heatingKw ?? 0) * machinesToBuild + beaconLayout.uniqueBeacons * BEACON_HEATING_KW;
  const warnings: string[] = [];
  const ignoredMachine = input.machineModules.filter((m) => m && !moduleIsAllowed(m, input.machine, input.recipe)).length;
  const ignoredBeacon = input.beaconModules.filter((m) => m && !moduleIsAllowed(m, input.machine, input.recipe, true)).length;
  if (ignoredMachine) warnings.push(`${ignoredMachine} incompatible machine module slot(s) were ignored.`);
  if (ignoredBeacon) warnings.push(`${ignoredBeacon} incompatible beacon module slot(s) were ignored.`);
  if (!input.isolatedBeacons && input.beaconCount > 0 && machinesToBuild > 1) {
    warnings.push(`Compact geometry provides ${beaconLayout.average.toFixed(2)} active beacons per machine on average (${beaconLayout.min}–${beaconLayout.max}).`);
  }
  if (machinesToBuild > isolatedBuild) warnings.push(`Shared-layout coverage increases the build count from ${isolatedBuild} to ${machinesToBuild}.`);
  if (beaconLayout.extrapolated) warnings.push("The full totals are exact or extrapolated; the drawing is capped for readability.");
  if (qualityChance > 0) warnings.push("Quality outputs are expected long-run averages.");
  if (input.recipe.products.some((p) => p.probability < 1)) warnings.push("Probabilistic outputs are shown as long-run expectations.");
  if (isolatedSpeedFactor <= .2000001) warnings.push("Crafting speed is clamped to Factorio's 20% minimum.");
  if (input.machine.auxiliaryActiveKw) warnings.push(`The rocket silo's separate ${input.machine.auxiliaryActiveKw / 1000} MW launch draw is not part of continuous crafting power.`);
  const outputDistribution = targetFlow.amounts.map((entry) => ({ ...entry, amount: entry.amount * cyclesPerSecond }));
  return {
    valid: true, cyclesPerSecond, ratePerSecond, inputs: scaledFlows(flows.inputs, cyclesPerSecond),
    outputs: scaledFlows(flows.outputs, cyclesPerSecond), targetPerCycle, limitingInput, machineEffects,
    isolatedBeaconEffects, craftingSpeed: isolatedCraftingSpeed, productivityBonus, qualityChance,
    theoreticalMachines, machinesToBuild, utilization, idealIsolatedMachines: theoreticalMachines,
    beaconLayout, electricTargetKw: machineElectricTargetKw + beaconElectricKw,
    electricFullLoadKw: machineElectricFullKw + beaconElectricKw, machineElectricTargetKw,
    beaconElectricKw, fuelTargetKw, nutrientItemsPerSecond, idleDrainKw, heatingKw,
    qualityDistribution: outputDistribution, warnings,
  };
}
export function toDisplayRate(perSecond: number, unit: RateUnit): number { return perSecond * UNIT_SECONDS[unit]; }
export function fromDisplayRate(value: number, unit: RateUnit): number { return value / UNIT_SECONDS[unit]; }
export function unitLabel(unit: RateUnit): string { return unit === "s" ? "/s" : unit === "min" ? "/min" : "/h"; }
