import * as RadixColors from "@radix-ui/colors";
import BezierEasing from "bezier-easing";
import Color from "colorjs.io";

import { ArrayOf12, grayScaleNames, scaleNames } from "../types";

const arrayOf12 = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11] as const;

export const lightColors = Object.fromEntries(
  scaleNames.map((scaleName) => [
    scaleName,
    Object.values(RadixColors[`${scaleName}P3`]).map((str) =>
      new Color(str).to("oklch")
    ),
  ])
) as Record<typeof scaleNames[number], ArrayOf12<Color>>;

const darkColors = Object.fromEntries(
  scaleNames.map((scaleName) => [
    scaleName,
    Object.values(RadixColors[`${scaleName}DarkP3`]).map((str) =>
      new Color(str).to("oklch")
    ),
  ])
) as Record<typeof scaleNames[number], ArrayOf12<Color>>;

const lightGrayColors = Object.fromEntries(
  grayScaleNames.map((scaleName) => [
    scaleName,
    Object.values(RadixColors[`${scaleName}P3`]).map((str) =>
      new Color(str).to("oklch")
    ),
  ])
) as Record<typeof grayScaleNames[number], ArrayOf12<Color>>;

const darkGrayColors = Object.fromEntries(
  grayScaleNames.map((scaleName) => [
    scaleName,
    Object.values(RadixColors[`${scaleName}DarkP3`]).map((str) =>
      new Color(str).to("oklch")
    ),
  ])
) as Record<typeof grayScaleNames[number], ArrayOf12<Color>>;

export const generateRadixColors = ({
  appearance,
  ...args
}: {
  appearance: "light" | "dark";
  accent: string;
  gray: string;
  background: string;
}) => {
  const allScales = appearance === "light" ? lightColors : darkColors;
  const grayScales = appearance === "light" ? lightGrayColors : darkGrayColors;
  const backgroundColor = new Color(args.background).to("oklch");

  const grayBaseColor = new Color(args.gray).to("oklch");
  const grayScaleColors = getScaleFromColor(
    grayBaseColor,
    grayScales,
    backgroundColor
  );

  const accentBaseColor = new Color(args.accent).to("oklch");

  let accentScaleColors = getScaleFromColor(
    accentBaseColor,
    allScales,
    backgroundColor
  );

  // Enforce srgb for the background color
  const backgroundHex = backgroundColor.to("srgb").toString({ format: "hex" });

  // Make sure we use the tint from the gray scale for when base is pure white or black
  const accentBaseHex = accentBaseColor.to("srgb").toString({ format: "hex" });
  if (accentBaseHex === "#000" || accentBaseHex === "#fff") {
    accentScaleColors = grayScaleColors.map((color) =>
      color.clone()
    ) as ArrayOf12<Color>;
  }

  const [accent9Color, accentContrastColor] = getStep9Colors(
    accentScaleColors,
    accentBaseColor
  );

  accentScaleColors[8] = accent9Color;
  accentScaleColors[9] = getButtonHoverColor(accent9Color, [accentScaleColors]);

  // Limit saturation of the text colors
  accentScaleColors[10].coords[1] = Math.min(
    Math.max(accentScaleColors[8].coords[1], accentScaleColors[7].coords[1]),
    accentScaleColors[10].coords[1]
  );
  accentScaleColors[11].coords[1] = Math.min(
    Math.max(accentScaleColors[8].coords[1], accentScaleColors[7].coords[1]),
    accentScaleColors[11].coords[1]
  );

  const accentScaleHex = accentScaleColors.map((color) =>
    color.to("srgb").toString({ format: "hex" })
  ) as ArrayOf12<string>;

  const accentScaleWideGamut = accentScaleColors.map(
    toOklchString
  ) as ArrayOf12<string>;

  const accentScaleAlphaHex = accentScaleHex.map((color) =>
    getAlphaColorSrgb(color, backgroundHex)
  ) as ArrayOf12<string>;

  const accentScaleAlphaWideGamutString = accentScaleHex.map((color) =>
    getAlphaColorP3(color, backgroundHex)
  ) as ArrayOf12<string>;

  const accentContrastColorHex = accentContrastColor
    .to("srgb")
    .toString({ format: "hex" });

  const grayScaleHex = grayScaleColors.map((color) =>
    color.to("srgb").toString({ format: "hex" })
  ) as ArrayOf12<string>;

  const grayScaleWideGamut = grayScaleColors.map(
    toOklchString
  ) as ArrayOf12<string>;

  const grayScaleAlphaHex = grayScaleHex.map((color) =>
    getAlphaColorSrgb(color, backgroundHex)
  ) as ArrayOf12<string>;

  const grayScaleAlphaWideGamutString = grayScaleHex.map((color) =>
    getAlphaColorP3(color, backgroundHex)
  ) as ArrayOf12<string>;

  const accentSurfaceHex =
    appearance === "light"
      ? getAlphaColorSrgb(accentScaleHex[1], backgroundHex, 0.8)
      : getAlphaColorSrgb(accentScaleHex[1], backgroundHex, 0.5);

  const accentSurfaceWideGamutString =
    appearance === "light"
      ? getAlphaColorP3(accentScaleWideGamut[1], backgroundHex, 0.8)
      : getAlphaColorP3(accentScaleWideGamut[1], backgroundHex, 0.5);

  return {
    accentScale: accentScaleHex,
    accentScaleAlpha: accentScaleAlphaHex,
    accentScaleWideGamut: accentScaleWideGamut,
    accentScaleAlphaWideGamut: accentScaleAlphaWideGamutString,
    accentContrast: accentContrastColorHex,

    grayScale: grayScaleHex,
    grayScaleAlpha: grayScaleAlphaHex,
    grayScaleWideGamut: grayScaleWideGamut,
    grayScaleAlphaWideGamut: grayScaleAlphaWideGamutString,

    graySurface: appearance === "light" ? "#ffffffcc" : "rgba(0, 0, 0, 0.05)",
    graySurfaceWideGamut:
      appearance === "light"
        ? "color(display-p3 1 1 1 / 80%)"
        : "color(display-p3 0 0 0 / 5%)",

    accentSurface: accentSurfaceHex,
    accentSurfaceWideGamut: accentSurfaceWideGamutString,

    background: backgroundHex,
  };
};

function getStep9Colors(
  scale: ArrayOf12<Color>,
  accentBaseColor: Color
): [Color, Color] {
  const referenceBackgroundColor = scale[0];
  const distance = accentBaseColor.deltaEOK(referenceBackgroundColor) * 100;

  // If the accent base color is close to the page background color, it's likely
  // white on white or black on black, so we want to return something that makes sense instead
  if (distance < 25) {
    return [scale[8], getTextColor(scale[8])];
  }

  return [accentBaseColor, getTextColor(accentBaseColor)];
}

function getButtonHoverColor(source: Color, scales: ArrayOf12<Color>[]) {
  const [L, C, H] = source.coords;
  const newL = L > 0.4 ? L - 0.03 / (L + 0.1) : L + 0.03 / (L + 0.1);
  const newC = L > 0.4 && !isNaN(H) ? C * 0.93 + 0 : C;
  const buttonHoverColor = new Color("oklch", [newL, newC, H]);

  // Find closest in-scale color to donate the chroma and hue.
  // Especially useful when the source color is pure white or black,
  // but the gray scale is tinted.
  let closestColor = buttonHoverColor;
  let minDistance = Infinity;

  scales.forEach((scale) => {
    for (const color of scale) {
      const distance = buttonHoverColor.deltaEOK(color);
      if (distance < minDistance) {
        minDistance = distance;
        closestColor = color;
      }
    }
  });

  buttonHoverColor.coords[1] = closestColor.coords[1];
  buttonHoverColor.coords[2] = closestColor.coords[2];
  return buttonHoverColor;
}

function getScaleFromColor(
  source: Color,
  scales: Record<string, ArrayOf12<Color>>,
  backgroundColor: Color
) {
  const allColors: { scale: string; color: Color; distance: number }[] = [];

  Object.entries(scales).forEach(([name, scale]) => {
    for (const color of scale) {
      const distance = source.deltaEOK(color);
      allColors.push({ scale: name, distance, color });
    }
  });

  allColors.sort((a, b) => a.distance - b.distance);

  // Remove non-unique scales
  const closestColors = allColors.filter(
    (color, i, arr) =>
      i === arr.findIndex((value) => value.scale === color.scale)
  );

  // If the next two closest colors are both grays, remove the second one until it’s not a gray anymore.
  // This is because up next we will be comparing how close the two closest colors are to the source color,
  // and since the grays are all extremely close to each other, we won’t get any useful data from the second
  // closest color if it’s also a gray.
  const grayScaleNamesStr = grayScaleNames as readonly string[];
  const allAreGrays = closestColors.every((color) =>
    grayScaleNamesStr.includes(color.scale)
  );
  if (!allAreGrays && grayScaleNamesStr.includes(closestColors[0].scale)) {
    while (grayScaleNamesStr.includes(closestColors[1].scale)) {
      closestColors.splice(1, 1);
    }
  }

  const colorA = closestColors[0];
  const colorB = closestColors[1];

  // Light trigonometry ahead.
  //
  // We want to determine the color that is the closest to the source color. Sometimes it makes sense
  // to proportionally mix the two closest colors together, but sometimes it is not useful at all.
  // Color coords are spatial in 3D, however we can treat the data we have as a 2D projection that is good enough.
  //
  // Case 1:
  // If the distances between the source color, the 1st closest color (A) and the 2nd closest color (B) form
  // a triangle where NEITHER angle A nor B are larger than 90 degrees, then we want to mix the 1st and the 2nd
  // closest colors in the same proportion as distances AD and BD are to each other. Mixing the two would result
  // in a color that would be closer to the source color than either of the two original closest colors.
  // Example: source color is a desaturated blue, which is between "indigo" and "slate" scales.
  //
  //        C ← Source color
  //       /|⟍
  //      / |  ⟍
  //   b /  |    ⟍  a
  //    /   |      ⟍
  //   /    |        ⟍
  //  A --- D -------- B
  //        ↑
  //        The color we want to use as the base, which is a mix of A and B.
  //
  // Case 2:
  // If the distances between the source color, the 1st closest color (A) and the 2nd closest color (B) form
  // a triangle where EITHER angle A or B are larger than 90 degrees, then we don’t care about point B because it’s
  // directionally the same as A, as mixing A and B can’t provide us with a color that is any closer to the source.
  // Example: source color is a saturated blue, with "blue" being the closest scale, and "indigo" just being further.
  //
  //      C ← Source color
  //       \⟍
  //        \  ⟍
  //         \    ⟍  a
  //        b \      ⟍
  //           \        ⟍
  //            A ------- B
  //            ↑
  //            The color we want to use as the base, which is not influenced by B.

  // We’ll need all the lengths of the triangle sides, named after the angles they look at:
  const a = colorB.distance;
  const b = colorA.distance;
  const c = colorA.color.deltaEOK(colorB.color);

  // We can get the ratios of AD to BD lengths with trigonometry using tangents,
  // as the ratio of the tangents of the opposite angles will match.
  const cosA = (b ** 2 + c ** 2 - a ** 2) / (2 * b * c);
  const radA = Math.acos(cosA);
  const sinA = Math.sin(radA);

  const cosB = (a ** 2 + c ** 2 - b ** 2) / (2 * a * c);
  const radB = Math.acos(cosB);
  const sinB = Math.sin(radB);

  // Tangent of angle C in the ACD triangle
  const tanC1 = cosA / sinA;

  // Tangent of angle C in the BCD triangle
  const tanC2 = cosB / sinB;

  // The ratio of the tangents corresponds to the ratio of the distances AD to BD
  // In the end, it means how much of scale B we want to mix into scale A.
  // If it’s "0" or less, this is an obtuse triangle from case 2, and we use just scale A.
  const ratio = Math.max(0, tanC1 / tanC2) * 0.5;

  // The base scale is going to be a mix of the two closest scales, with the mix ratio we determined before
  const scaleA = scales[colorA.scale];
  const scaleB = scales[colorB.scale];
  const scale = arrayOf12.map((i) =>
    new Color(Color.mix(scaleA[i], scaleB[i], ratio)).to("oklch")
  ) as ArrayOf12<Color>;

  // Get the closest color from the pre-mixed scale we created
  const baseColor = scale
    .slice()
    .sort((a, b) => source.deltaEOK(a) - source.deltaEOK(b))[0];

  // Note the chroma difference between the source color and the base color
  const ratioC = source.coords[1] / baseColor.coords[1];

  // Modify hue and chroma of the scale to match the source color
  scale.forEach((color) => {
    color.coords[1] = Math.min(
      source.coords[1] * 1.5,
      color.coords[1] * ratioC
    );
    color.coords[2] = source.coords[2];
  });

  // Light mode
  if (scale[0].coords[0] > 0.5) {
    const lightnessScale = scale.map(({ coords }) => coords[0]);
    const backgroundL = Math.max(0, Math.min(1, backgroundColor.coords[0]));
    const newLightnessScale = transposeProgressionStart(
      backgroundL,
      // Add white as the first "step" of the light scale
      [1, ...lightnessScale],
      lightModeEasing
    );

    // Remove the step we added
    newLightnessScale.shift();

    newLightnessScale.forEach((lightness, i) => {
      scale[i].coords[0] = lightness;
    });

    return scale;
  }

  // Dark mode
  const ease: typeof darkModeEasing = [...darkModeEasing];
  const referenceBackgroundColorL = scale[0].coords[0];
  const backgroundColorL = Math.max(0, Math.min(1, backgroundColor.coords[0]));

  // If background is lighter than step 0, we want to gradually change the easing to linear
  const ratioL = backgroundColorL / referenceBackgroundColorL;

  if (ratioL > 1) {
    const maxRatio = 1.5;

    for (let i = 0; i < ease.length; i++) {
      const metaRatio = (ratioL - 1) * (maxRatio / (maxRatio - 1));
      ease[i] = ratioL > maxRatio ? 0 : Math.max(0, ease[i] * (1 - metaRatio));
    }
  }

  const lightnessScale = scale.map(({ coords }) => coords[0]);
  const backgroundL = backgroundColor.coords[0];
  const newLightnessScale = transposeProgressionStart(
    backgroundL,
    lightnessScale,
    ease
  );

  newLightnessScale.forEach((lightness, i) => {
    scale[i].coords[0] = lightness;
  });

  return scale;
}

function getTextColor(background: Color) {
  const white = new Color("oklch", [1, 0, 0]);

  if (Math.abs(white.contrastAPCA(background)) < 40) {
    const [L, C, H] = background.coords;
    return new Color("oklch", [0.25, Math.max(0.08 * C, 0.04), H]);
  }

  return white;
}

// target = background * (1 - alpha) + foreground * alpha
// alpha = (target - background) / (foreground - background)
// Expects 0-1 numbers for the RGB channels
function getAlphaColor(
  targetRgb: number[],
  backgroundRgb: number[],
  rgbPrecision: number,
  alphaPrecision: number,
  targetAlpha?: number
) {
  const [tr, tg, tb] = targetRgb.map((c) => Math.round(c * rgbPrecision));
  const [br, bg, bb] = backgroundRgb.map((c) => Math.round(c * rgbPrecision));

  if (
    tr === undefined ||
    tg === undefined ||
    tb === undefined ||
    br === undefined ||
    bg === undefined ||
    bb === undefined
  ) {
    throw Error("Color is undefined");
  }

  // Is the background color lighter, RGB-wise, than target color?
  // Decide whether we want to add as little color or as much color as possible,
  // darkening or lightening the background respectively.
  // If at least one of the bits of the target RGB value
  // is lighter than the background, we want to lighten it.
  let desiredRgb = 0;
  if (tr > br) {
    desiredRgb = rgbPrecision;
  } else if (tg > bg) {
    desiredRgb = rgbPrecision;
  } else if (tb > bb) {
    desiredRgb = rgbPrecision;
  }

  const alphaR = (tr - br) / (desiredRgb - br);
  const alphaG = (tg - bg) / (desiredRgb - bg);
  const alphaB = (tb - bb) / (desiredRgb - bb);

  const isPureGray = [alphaR, alphaG, alphaB].every(
    (alpha) => alpha === alphaR
  );

  // No need for precision gymnastics with pure grays, and we can get cleaner output
  if (!targetAlpha && isPureGray) {
    // Convert back to 0-1 values
    const V = desiredRgb / rgbPrecision;
    return [V, V, V, alphaR] as const;
  }

  const clampRgb = (n: number) =>
    isNaN(n) ? 0 : Math.min(rgbPrecision, Math.max(0, n));
  const clampA = (n: number) =>
    isNaN(n) ? 0 : Math.min(alphaPrecision, Math.max(0, n));
  const maxAlpha = targetAlpha ?? Math.max(alphaR, alphaG, alphaB);

  const A = clampA(Math.ceil(maxAlpha * alphaPrecision)) / alphaPrecision;
  let R = clampRgb(((br * (1 - A) - tr) / A) * -1);
  let G = clampRgb(((bg * (1 - A) - tg) / A) * -1);
  let B = clampRgb(((bb * (1 - A) - tb) / A) * -1);

  R = Math.ceil(R);
  G = Math.ceil(G);
  B = Math.ceil(B);

  const blendedR = blendAlpha(R, A, br);
  const blendedG = blendAlpha(G, A, bg);
  const blendedB = blendAlpha(B, A, bb);

  // Correct for rounding errors in light mode
  if (desiredRgb === 0) {
    if (tr <= br && tr !== blendedR) {
      R = tr > blendedR ? R + 1 : R - 1;
    }

    if (tg <= bg && tg !== blendedG) {
      G = tg > blendedG ? G + 1 : G - 1;
    }

    if (tb <= bb && tb !== blendedB) {
      B = tb > blendedB ? B + 1 : B - 1;
    }
  }

  // Correct for rounding errors in dark mode
  if (desiredRgb === rgbPrecision) {
    if (tr >= br && tr !== blendedR) {
      R = tr > blendedR ? R + 1 : R - 1;
    }

    if (tg >= bg && tg !== blendedG) {
      G = tg > blendedG ? G + 1 : G - 1;
    }

    if (tb >= bb && tb !== blendedB) {
      B = tb > blendedB ? B + 1 : B - 1;
    }
  }

  // Convert back to 0-1 values
  R = R / rgbPrecision;
  G = G / rgbPrecision;
  B = B / rgbPrecision;

  return [R, G, B, A] as const;
}

// Important – I empirically discovered that this rounding is how the browser actually overlays
// transparent RGB bits over each other. It does NOT round the whole result altogether.
function blendAlpha(
  foreground: number,
  alpha: number,
  background: number,
  round = true
) {
  if (round) {
    return (
      Math.round(background * (1 - alpha)) + Math.round(foreground * alpha)
    );
  }

  return background * (1 - alpha) + foreground * alpha;
}

function getAlphaColorSrgb(
  targetColor: string,
  backgroundColor: string,
  targetAlpha?: number
) {
  const [r, g, b, a] = getAlphaColor(
    new Color(targetColor).to("srgb").coords,
    new Color(backgroundColor).to("srgb").coords,
    255,
    255,
    targetAlpha
  );

  return formatHex(new Color("srgb", [r, g, b], a).toString({ format: "hex" }));
}

function getAlphaColorP3(
  targetColor: string,
  backgroundColor: string,
  targetAlpha?: number
) {
  const [r, g, b, a] = getAlphaColor(
    new Color(targetColor).to("p3").coords,
    new Color(backgroundColor).to("p3").coords,
    // Not sure why, but the resulting P3 alpha colors are blended in the browser most precisely when
    // rounded to 255 integers too. Is the browser using 0-255 rather than 0-1 under the hood for P3 too?
    255,
    1000,
    targetAlpha
  );

  return (
    new Color("p3", [r, g, b], a)
      .toString({ precision: 4 })
      // Important: in non-browser environments colorjs.io outputs a different format for some reason
      .replace("color(p3 ", "color(display-p3 ")
  );
}

// Format shortform hex to longform
function formatHex(str: string) {
  if (!str.startsWith("#")) {
    return str;
  }

  if (str.length === 4) {
    const hash = str.charAt(0);
    const r = str.charAt(1);
    const g = str.charAt(2);
    const b = str.charAt(3);
    return hash + r + r + g + g + b + b;
  }

  if (str.length === 5) {
    const hash = str.charAt(0);
    const r = str.charAt(1);
    const g = str.charAt(2);
    const b = str.charAt(3);
    const a = str.charAt(4);
    return hash + r + r + g + g + b + b + a + a;
  }

  return str;
}

const darkModeEasing = [1, 0, 1, 0] as [number, number, number, number];
const lightModeEasing = [0, 2, 0, 2] as [number, number, number, number];

function transposeProgressionStart(
  to: number,
  arr: number[],
  curve: [number, number, number, number]
) {
  return arr.map((n, i, arr) => {
    const lastIndex = arr.length - 1;
    const diff = arr[0] - to;
    const fn = BezierEasing(...curve);
    return n - diff * fn(1 - i / lastIndex);
  });
}

function transposeProgressionEnd(
  to: number,
  arr: number[],
  curve: [number, number, number, number]
) {
  return arr.map((n, i, arr) => {
    const lastIndex = arr.length - 1;
    const diff = arr[lastIndex] - to;
    const fn = BezierEasing(...curve);
    return n - diff * fn(i / lastIndex);
  });
}

// Convert to OKLCH string with percentage for the lightness channel
// https://github.com/radix-ui/themes/issues/420
function toOklchString(color: Color) {
  const L = +(color.coords[0] * 100).toFixed(1);
  return color
    .to("oklch")
    .toString({ precision: 4 })
    .replace(/(\S+)(.+)/, `oklch(${L}%$2`);
}