import type { DiagramIRV2 } from "../../model/diagram-ir.js";

export interface HeatmapRequest {
  metric: string;
  min: number;
  max: number;
  values: Readonly<Record<string, number>>;
  colors?: readonly string[];
}
export interface HeatmapLegendStop { value: number; color: string; fontColor: string }
export interface HeatmapLegend { metric: string; min: number; max: number; stops: HeatmapLegendStop[] }
export interface HeatmapResult { diagram: DiagramIRV2; legend: HeatmapLegend }

export const DEFAULT_HEATMAP_COLORS = Object.freeze(["#440154", "#3B528B", "#21918C", "#5EC962", "#FDE725"]);
const HEX_COLOR = /^#[0-9A-Fa-f]{6}$/;

function setHeatmapColors(style: string | undefined, color: string, fontColor: string): string {
  const entries = (style ?? "").split(";").filter(Boolean);
  const replacements = new Map([["fillColor", color], ["fontColor", fontColor]]);
  const targetKeys = new Set(replacements.keys());
  const output = entries.map((entry) => {
    const key = entry.split("=", 1)[0];
    if (!targetKeys.has(key)) return entry;
    const replacement = replacements.get(key);
    if (replacement === undefined) return undefined;
    replacements.delete(key);
    return `${key}=${replacement}`;
  }).filter((entry): entry is string => entry !== undefined);
  for (const [key, value] of replacements) output.push(`${key}=${value}`);
  return `${output.join(";")};`;
}

function contrastRatio(left: string, right: string): number {
  const luminance = (color: string): number => [1, 3, 5]
    .map((offset) => Number.parseInt(color.slice(offset, offset + 2), 16) / 255)
    .map((value) => value <= 0.04045 ? value / 12.92 : ((value + 0.055) / 1.055) ** 2.4)
    .reduce((total, value, index) => total + value * [0.2126, 0.7152, 0.0722][index], 0);
  const [high, low] = [luminance(left), luminance(right)].sort((a, b) => b - a);
  return (high + 0.05) / (low + 0.05);
}

const READABLE_FONT_COLORS = ["#FFFFFF", "#1A1A1A"] as const;

function readableFontColor(background: string): string {
  const color = READABLE_FONT_COLORS.find((candidate) => contrastRatio(candidate, background) >= 4.5);
  if (!color) throw new Error(`Heatmap palette color ${background} has no readable font color`);
  return color;
}

function exactNumber(value: number): number {
  return Number(value.toFixed(12));
}

export function applyMetricsHeatmap(ir: DiagramIRV2, request: HeatmapRequest): HeatmapResult {
  if (typeof request.metric !== "string" || request.metric.trim() === "") throw new Error("Heatmap metric must be a non-empty string");
  if (!Number.isFinite(request.min) || !Number.isFinite(request.max) || request.min >= request.max) throw new Error("Heatmap min must be less than max");
  if (!request.values || typeof request.values !== "object" || Array.isArray(request.values)) throw new Error("Heatmap values must be an object");
  const colors = [...(request.colors ?? DEFAULT_HEATMAP_COLORS)];
  if (colors.length < 2 || colors.some((color) => typeof color !== "string" || !HEX_COLOR.test(color))) throw new Error("Heatmap colors must contain at least two #RRGGBB colors");
  const nodeKeys = new Set(ir.pages.flatMap((page) => page.nodes.map((node) => `${page.id}/${node.id}`)));
  for (const [key, value] of Object.entries(request.values)) {
    if (!nodeKeys.has(key)) throw new Error(`Heatmap value has unknown node: ${key}`);
    if (typeof value !== "number" || !Number.isFinite(value) || value < request.min || value > request.max) {
      throw new Error(`Heatmap value for ${key} must be numeric and bounded by min and max`);
    }
  }
  const stops = colors.map((color, index) => ({
    value: exactNumber(request.min + (request.max - request.min) * index / (colors.length - 1)),
    color,
    fontColor: readableFontColor(color),
  }));
  const legend: HeatmapLegend = { metric: request.metric, min: request.min, max: request.max, stops };
  const diagram = structuredClone(ir);
  for (const page of diagram.pages) for (const node of page.nodes) {
    const value = request.values[`${page.id}/${node.id}`];
    if (value === undefined) continue;
    const nextStopIndex = stops.findIndex((stop) => value < stop.value);
    const colorIndex = nextStopIndex === -1 ? stops.length - 1 : Math.max(0, nextStopIndex - 1);
    node.style = setHeatmapColors(node.style, colors[colorIndex], stops[colorIndex].fontColor);
  }
  diagram.extensions = { ...(diagram.extensions ?? {}), heatmap: legend };
  return { diagram, legend };
}

export const transformMetricsHeatmap = applyMetricsHeatmap;