import * as tf from '@tensorflow/tfjs';
import { setWasmPaths } from '@tensorflow/tfjs-backend-wasm';
import modelJSON from './model.json';
import * as ccl from './ccl';

const charset = [' ', '0', '2', '4', '5', '8', 'A', 'D', 'G', 'H', 'J', 'K', 'M', 'N', 'P', 'R', 'S', 'T', 'V', 'W', 'X', 'Y'];
let weightsData: Uint8Array; // base64 encoded weights
let model: tf.LayersModel;

tf.enableProdMode();

const wasmToUrl = (wasm: any) => {
  const blb = new Blob([wasm], { type: 'application/wasm' });
  return URL.createObjectURL(blb);
};

const backendloaded = (async () => {
  try {
    // dead code elimination should occur here
    // eslint-disable-next-line camelcase
    if (execution_mode === 'userscript' || execution_mode === 'test') {
      weightsData = (await import('./group1-shard1of1.bin')).default;
      const tfwasmthreadedsimd = (await import('./tfjs-backend-wasm-threaded-simd.wasm')).default;
      const tfwasmsimd = (await import('./tfjs-backend-wasm-simd.wasm')).default;
      const tfwasm = (await import('./tfjs-backend-wasm.wasm')).default;
      setWasmPaths({
        'tfjs-backend-wasm.wasm': wasmToUrl(tfwasm),
        'tfjs-backend-wasm-simd.wasm': wasmToUrl(tfwasmsimd),
        'tfjs-backend-wasm-threaded-simd.wasm': wasmToUrl(tfwasmthreadedsimd)
      });
    } else {
      weightsData = new Uint8Array(await (await fetch(chrome.runtime.getURL('./group1-shard1of1.bin'))).arrayBuffer());
      const args = {
        'tfjs-backend-wasm.wasm': chrome.runtime.getURL('tfjs-backend-wasm.wasm'),
        'tfjs-backend-wasm-simd.wasm': chrome.runtime.getURL('tfjs-backend-wasm-simd.wasm'),
        'tfjs-backend-wasm-threaded-simd.wasm': chrome.runtime.getURL('tfjs-backend-wasm-threaded-simd.wasm')
      };
      setWasmPaths(args);
    }
    const l = await tf.setBackend('wasm');
    console.log('tf backend loaded', l);
  } catch (err) {
    console.log('tf err', err);
  }
})();

function toggle(obj: HTMLElement, v?: any) {
  if (v) obj.style.display = '';
  else obj.style.display = 'none';
}

const iohander: tf.io.IOHandler = {
  load: function () {
    return new Promise((resolve, reject) => {
      resolve({
        modelTopology: modelJSON.modelTopology,
        weightSpecs: modelJSON.weightsManifest[0].weights as tf.io.WeightsManifestEntry[],
        weightData: weightsData.buffer,
        format: modelJSON.format,
        generatedBy: modelJSON.generatedBy,
        convertedBy: modelJSON.convertedBy
      });
    });
  }
};

async function load() {
  model = await tf.loadLayersModel(iohander);
  return model;
}

function black(x: number) {
  return x < 64;
}

// Calculates "disorder" of the image. "Disorder" is the percentage of black pixels that have a
// non-black pixel below them. Minimizing this seems to be good enough metric for solving the slider.
function calculateDisorder(imgdata: ImageData) {
  const a = imgdata.data;
  const w = imgdata.width;
  const h = imgdata.height;
  const pic: number[] = [];
  const visited: number[] = [];

  for (let c = 0; c < w * h; c++) {
    if (visited[c]) continue;
    if (!black(a[c * 4])) continue;

    let blackCount = 0;
    const items: number[] = [];
    const toVisit: number[] = [c];
    while (toVisit.length > 0) {
      const cc = toVisit[toVisit.length - 1];
      toVisit.splice(toVisit.length - 1, 1);

      if (visited[cc]) continue;
      visited[cc] = 1;

      if (black(a[cc * 4])) {
        items.push(cc);

        blackCount++;
        toVisit.push(cc + 1);
        toVisit.push(cc - 1);
        toVisit.push(cc + w);
        toVisit.push(cc - w);
      }
    }

    if (blackCount >= 24) {
      items.forEach(function (x) {
        pic[x] = 1;
      });
    }
  }

  let res = 0;
  let total = 0;
  for (let c = 0; c < w * h - w; c++) {
    if (pic[c] !== pic[c + w]) res += 1;
    if (pic[c]) total += 1;
  }

  return res / (total === 0 ? 1 : total);
}

// returns ImageData from captcha's background image, foreground image, and offset (ranging from 0 to -50)
function imageFromCanvas(img: HTMLImageElement, bg: HTMLImageElement, off: number) {
  const h = img.height;
  const w = img.width;
  const th = 80;
  const ph = 0;
  const pw = 16;
  const scale = th / h;

  const canvas = document.createElement('canvas');
  canvas.width = w * scale + pw * 2;
  canvas.height = th;

  const ctx = canvas.getContext('2d', { willReadFrequently: true })!;

  ctx.fillStyle = 'rgb(238,238,238)';
  ctx.fillRect(0, 0, canvas.width, canvas.height);

  ctx.translate(canvas.width / 2, canvas.height / 2);

  const draw = function (off: number) {
    if (bg) {
      const border = 4;
      ctx.drawImage(
        bg,
        -off + border,
        0,
        w - border * 2,
        h,
        -w / 2 + border,
        -h / 2,
        w - border * 2,
        h
      );
    }
    ctx.drawImage(img, -w / 2, -h / 2, w, h);
  };

  // if off is not specified and background image is present, try to figure out
  // the best offset automatically; select the offset that has smallest value of
  // calculateDisorder for the resulting image
  if (bg && off == null) {
    let bestDisorder = 999;
    let bestImagedata: ImageData | null = null;
    let bestOff = -1;

    for (let off = 0; off >= -50; off--) {
      draw(off);

      let imgdata = ctx.getImageData(0, 0, canvas.width, canvas.height);
      const disorder = calculateDisorder(imgdata);

      if (disorder < bestDisorder) {
        bestDisorder = disorder;
        draw(off);
        imgdata = ctx.getImageData(0, 0, canvas.width, canvas.height);
        bestImagedata = imgdata;
        bestOff = off;
      }
    }

    // not the best idea to do this here
    setTimeout(function () {
      const bg = document.getElementById('t-bg');
      const slider = document.getElementById('t-slider') as HTMLInputElement;
      if (!bg || !slider) return;

      slider.value = '' + (-bestOff * 2);
      bg.style.backgroundPositionX = bestOff + 'px';
    }, 1);
    draw(bestOff);
    return bestImagedata;
  } else {
    draw(off);
    return ctx.getImageData(0, 0, canvas.width, canvas.height);
  }
}

// for debugging purposes
function imagedataToImage(imagedata: ImageData) {
  const canvas = document.createElement('canvas');
  const ctx = canvas.getContext('2d')!;
  canvas.width = imagedata.width;
  canvas.height = imagedata.height;
  ctx.putImageData(imagedata, 0, 0);

  const image = new Image();
  image.src = canvas.toDataURL();
  return image;
}

function toMonochrome(px: Uint8ClampedArray) {
  const ret = Array<number>(px.length >> 2);
  for (let i = 0; i < px.length; i += 4) {
    ret[i >> 2] = +(px[i] < 128);
  }
  return ret;
}

const greedyCTCDecode = (yPred: tf.Tensor<tf.Rank>) => tf.tidy(() => yPred.argMax(-1).arraySync());

function imgDisp(pixConv: (f: ArrayLike<number>, w: number, h: number, s: Uint8ClampedArray) => ArrayLike<number>, img: ArrayLike<number>, w: number, h: number, t?: boolean) {
  const dt = new ImageData(w, h);
  const rgba = pixConv(img, w, h, dt.data);
  const imgres = imagedataToImage(dt);
  document.body.appendChild(imgres);
  if (t) {
    imgres.style.transform = 'rotate(90deg) scaleY(-1)';
  }
}

let colors = [
  [255, 0, 0],   // Red
  [0, 255, 0],   // Green
  [0, 0, 255],   // Blue
  [255, 255, 0], // Yellow
  [255, 0, 255], // Magenta
  [0, 255, 255], // Cyan
  [128, 0, 0],   // Dark Red
  [0, 128, 0],   // Dark Green
  [0, 0, 128],   // Dark Blue
  [128, 128, 0], // Olive
  [128, 0, 128], // Purple
  [0, 128, 128], // Teal
  [192, 192, 192], // Silver
  [128, 128, 128], // Gray
  [255, 165, 0], // Orange
  [0, 128, 64]   // Medium Sea Green
];

const monoToPalette = (p: number[][], max: number) =>
  function (arr: ArrayLike<number>, w: number, h: number, res: Uint8ClampedArray) {
    let choice = p.slice(0);
    const choices = new Map<number, number[]>();
    for (let i = 0; i < arr.length; ++i) {
      let col: number[];
      if (choices.has(arr[i])) {
        col = choices.get(arr[i])!;
      } else {
        col = choice.shift()!;
        choices.set(arr[i], col);
        if (choice.length == 0)
          choice = p.slice(0);
      }
      [res[i * 4], res[i * 4 + 1], res[i * 4 + 2]] = col;
      res[i * 4 + 3] = 255;
    }
    return res;
  };


function monoToRgba(arr: ArrayLike<number>, w: number, h: number, res: Uint8ClampedArray) {
  for (let i = 0; i < arr.length; ++i) {
    res[i * 4] = res[i * 4 + 1] = res[i * 4 + 2] = arr[i] * 255;
    res[i * 4 + 3] = 255;
  }
  return res;
}

function processCTCDecodedSequence(decodedSequence: number[], blankLabel = 0) {
  const result = [];
  let prevLabel = blankLabel;

  for (const label of decodedSequence) {
    if (label !== blankLabel && label !== prevLabel) {
      result.push(label);
    }
    prevLabel = label;
  }

  return result;
}

function indicesToSymbols(decodedIndices: number[]) {
  return decodedIndices.map(index => charset[index - 1] || '');
}

async function predict(img: HTMLImageElement, bg: HTMLImageElement, off: number) {
  if (!model) {
    model = await load();
  }
  const image = imageFromCanvas(img, bg, off);
  if (!image)
    throw new Error("Failed to gen image");
  const mono = toMonochrome(image.data);
  console.log(mono.reduce((a, b) => a + b), 0);
  const labels = ccl.connectedComponentLabeling(mono, image.width, image.height);
  const props = ccl.computeBounds(labels, image.width, image.height);

  const sortedByArea = Object.entries(props).sort((a, b) => a[1].area - b[1].area);
  const n = 8;
  let eightBiggest = sortedByArea.slice(0, -n);
  //const filtered = new Float32Array(80 * 300);

  // TODO: maybe centering?
  //imgDisp(monoToPalette(colors, Math.max(...new Set(labels))), labels, image.width, image.height);

  for (const [label, region] of eightBiggest) {
    //if ((region.maxRow - region.minRow) <= 20) {
    //  continue;
    //}

    for (let y = region.minRow; y <= region.maxRow; ++y) {
      for (let x = region.minCol; x <= region.maxCol; ++x) {
        if (labels[y * image!.width + x] === +label) {
          labels[y * image!.width + x] = 0;
        }
      }
    }
  }

  eightBiggest = sortedByArea.slice(-n);
  //imgDisp(monoToPalette(colors, Math.max(...new Set(labels))), labels, image.width, image.height);

  for (const [label, region] of eightBiggest) {
    if ((region.maxRow - region.minRow) > 20) {
      continue;
    }

    for (let y = region.minRow; y <= region.maxRow; ++y) {
      for (let x = region.minCol; x <= region.maxCol; ++x) {
        if (labels[y * image!.width + x] === +label) {
          labels[y * image!.width + x] = 0;
        }
      }
    }
  }
  //imgDisp(monoToPalette(colors, Math.max(...new Set(labels))), labels, image.width, image.height);

  for (const [label, region] of eightBiggest) {
    if ((region.maxRow - region.minRow) <= 20) {
      continue;
    }

    for (let y = region.minRow; y <= region.maxRow; ++y) {
      for (let x = region.minCol; x <= region.maxCol; ++x) {
        if (labels[y * image!.width + x] === +label) {
          labels[y * image!.width + x] = 1;
        }
      }
    }
  }

  const filtered2 = tf.tensor3d(labels, [image.height, image.width, 1]).concat(tf.zeros([80, 300 - image.width, 1]), 1);
  //imgDisp(monoToPalette(colors, Math.max(...new Set(labels))), labels, image.width, image.height);
  //const tensor = tf.tensor3d(filtered, [80, 300, 1], 'float32');
  //const tr = [1, 0, 2];
  //console.log(tensor.shape, tensor.transpose(tr).shape);
  const prediction = model.predict(filtered2.transpose([1, 0, 2]).expandDims(0));

  let d: tf.TypedArray;

  if (!Array.isArray(prediction)) {
    const v = greedyCTCDecode(prediction) as number[][];
    console.log(v);
    const s = processCTCDecodedSequence(v[0], charset.length + 1);
    return indicesToSymbols(s).join('').trim();
  } else
    throw new Error("unexpected inference");

  // createSequence(d);
  return '';
}

function createSequence(prediction: any) {
  const csl = charset.length;
  const sequence: Record<string, number>[] = [];

  // for each prediction
  for (let pos = 0; pos < prediction.length; pos += csl) {
    // look at the probabilities for the 22 token characters
    const preds = prediction.slice(pos, pos + csl);
    const max = Math.max(...preds);

    const seqElem: Record<string, number> = {};

    for (let i = 0; i < csl; i++) {
      const p = preds[i] / max; // normalize probability
      const c = charset[i + 1];

      if (p >= 0.05) { // if it's probable enough
        seqElem[c || ''] = p; // save its probability, to give alternative solutions
      }
    }

    sequence.push(seqElem);
  }

  return sequence;
}


async function imageFromUri(uri: string) {
  if (uri.startsWith('url("')) {
    uri = uri.substr(5, uri.length - 7);
  }
  // eslint-disable-next-line camelcase
  if (execution_mode !== 'test' && !uri.startsWith('data:')) {
    return null;
  }

  const img = new Image();
  await new Promise((r) => {
    img.onload = r;
    img.src = uri;
  });
  return img;
}

async function predictUri(uri: string, uribg: string, bgoff: string | null) {
  const img = await imageFromUri(uri);
  const bg = uribg ? await imageFromUri(uribg) : null;
  const off = bgoff ? parseInt(bgoff) : null;

  return await predict(img!, bg!, off!);
}

const solveButton = document.createElement('input');
solveButton.id = 't-auto-solve';
solveButton.value = 'Solve';
solveButton.type = 'button';
solveButton.style.fontSize = '11px';
solveButton.style.padding = '0 2px';
solveButton.style.margin = '0px 0px 0px 6px';
solveButton.style.height = '18px';
solveButton.onclick = async function () {
  solve(true);
};

const altsDiv = document.createElement('div');
altsDiv.id = 't-auto-options';
altsDiv.style.margin = '0';
altsDiv.style.padding = '0';

let storedPalceholder: string;

let overrides = {};

function placeAfter(elem: HTMLElement, sibling: HTMLElement) {
  if (elem.parentElement !== sibling.parentElement) {
    setTimeout(function () {
      sibling.parentElement?.insertBefore(elem, sibling.nextElementSibling);
    }, 1);
  }
}

let previousText: string | null = null;
async function solve(force?: any) {
  const resp = document.getElementById('t-resp') as HTMLInputElement;
  if (!resp) return;

  const bg = document.getElementById('t-bg');
  if (!bg) return;

  const fg = document.getElementById('t-fg');
  if (!fg) return;

  const help = document.getElementById('t-help');
  if (!help) return;

  await backendloaded;

  placeAfter(solveButton, resp);
  placeAfter(altsDiv, help);

  // palememe
  setTimeout(function () {
    toggle(solveButton, bg.style.backgroundImage);
  }, 1);

  const text = fg.style.backgroundImage;
  if (!text) {
    altsDiv.innerHTML = '';
    return;
  }

  if (text === previousText && !force) return;
  previousText = text;

  altsDiv.innerHTML = '';
  if (!storedPalceholder)
    storedPalceholder = resp.placeholder;
  resp.placeholder = 'solving captcha...';

  overrides = {};

  const result = await predictUri(
    text,
    bg.style.backgroundImage,
    force ? bg.style.backgroundPositionX : null
  );
  resp.placeholder = storedPalceholder;

  resp.value = result;
}

const observer = new MutationObserver(async function (mutationsList, observer) {
  solve(false);
});

//window['solve'] = solve;

observer.observe(document.body, {
  attributes: true,
  childList: true,
  subtree: true
});