import * as Plot from "@observablehq/plot";
import * as d3 from "d3";
import * as mobx from "mobx";
import * as React from "react";
import { $StateContext, AUTOGRID_MIN } from "./context";
import { events } from "./plot/events";
import { ellipse } from "./plot/ellipse";
import { img } from "./plot/img";
import { pixels } from "./plot/pixels";
import { binding, flatten, tw, useContainerWidth, joinClasses } from "./utils";

const Marks = { ...Plot, ellipse, events, img, pixels };
const { useEffect, useMemo } = React;
export const DEFAULT_PLOT_OPTIONS = { inset: 10 };

// Add per-mark defaults
const PER_MARK_DEFAULTS = {
  dot: { fill: "currentColor" },
  frame: { stroke: "#dddddd" },
};

function deepMergeLayers(target, source) {
  if (!source || typeof source !== "object") return target;

  return Object.keys(source).reduce(
    (result, key) => {
      if (key === "marks") {
        result[key] = [...result[key], ...source[key]];
      } else if (
        source[key] &&
        typeof source[key] === "object" &&
        key in result
      ) {
        result[key] = deepMergeLayers(result[key], source[key]);
      } else {
        result[key] = source[key];
      }
      return result;
    },
    { ...target },
  );
}

function mergePlotSpec(layers) {
  return layers.flat().reduce(
    (mergedSpec, layer) => {
      if (!layer) return mergedSpec;
      if (layer instanceof MarkSpec) {
        mergedSpec.marks.push(layer);
        return mergedSpec;
      }
      return deepMergeLayers(mergedSpec, layer.spec || layer);
    },
    { marks: [] },
  );
}

export class PlotSpec {
  constructor({ layers }) {
    this.spec = mergePlotSpec(layers);
  }

  render() {
    return <PlotWrapper spec={this.spec} />;
  }
}

export class MarkSpec {
  constructor(name, data = [], options) {
    if (!Marks[name]) {
      throw new Error(`Plot function "${name}" not found.`);
    }
    this.fn = Marks[name];

    options = { ...PER_MARK_DEFAULTS[name], ...options };

    // handle dimensional data passed in the 1st position
    if (data.dimensions) {
      options.dimensions = data.dimensions;
      data = data.value;
    }

    // flatten dimensional data
    if (options.dimensions) {
      data = flatten(data, options.dimensions);
    }

    this.data = data;
    this.options = options;
    this.computed = {};
  }

  render() {
    return <PlotWrapper spec={{ marks: [this] }} />;
  }

  compute(width) {
    if (width && this.computed.forWidth == width) {
      return this.computed;
    }
    this.extraMarks = [];
    let data = this.data;
    let options = { ...this.options };
    let computed = {
      forWidth: width,
      extraMarks: [],
      plotOptions: { ...DEFAULT_PLOT_OPTIONS },
      fn: this.fn,
    };

    // Below is where we add functionality to Observable.Plot by preprocessing
    // the data & options that are passed in.

    // handle columnar data in the 1st position
    function isArrayLike(value) {
      return Array.isArray(value) || ArrayBuffer.isView(value);
    }

    if (data.constructor === Object && !("length" in data)) {
      let length = null;
      data = { ...data };
      for (let [key, value] of Object.entries(data)) {
        if (isArrayLike(value)) {
          options[key] = value;
          delete data[key];
          length = value.length;
        }
        if (length !== null) {
          data.length = length;
        }
      }
    }
    // handle facetWrap option (grid) with minWidth consideration
    // see https://github.com/observablehq/plot/pull/892/files
    if (options.facetGrid) {
      // Check if data is an array of objects
      if (
        !Array.isArray(data) ||
        !data.every((item) => typeof item === "object" && item !== null)
      ) {
        throw new Error(
          "Invalid data format: facetGrid expects an array of objects",
        );
      }
      const facetGrid =
        typeof options.facetGrid === "string"
          ? [options.facetGrid, {}]
          : options.facetGrid;
      const [key, gridOpts] = facetGrid;
      const keys = Array.from(d3.union(data.map((d) => d[key])));
      const index = new Map(keys.map((key, i) => [key, i]));

      // Calculate columns based on minWidth and containerWidth
      const minWidth = gridOpts.minWidth || AUTOGRID_MIN;
      const columns = gridOpts.columns || Math.floor(width / minWidth);
      const fx = (key) => index.get(key) % columns;
      const fy = (key) => Math.floor(index.get(key) / columns);
      options.fx = (d) => fx(d[key]);
      options.fy = (d) => fy(d[key]);
      computed.plotOptions = {
        ...computed.plotOptions,
        fx: { axis: null },
        fy: { axis: null },
      };
      computed.extraMarks.push(
        Plot.text(keys, {
          fx,
          fy,
          frameAnchor: "top",
          dy: 4,
        }),
      );
    }
    computed.data = data;
    computed.options = options;
    this.computed = computed;
    return computed;
  }
}

function someObject(obj) {
  return Object.keys(obj).length === 0 ? null : obj;
}

export function readMark(mark, width) {
  if (!(mark instanceof MarkSpec)) {
    return mark;
  }
  let { fn, data, options, plotOptions, extraMarks } = mark.compute(width);
  mark = someObject(options) ? fn(data, options) : fn(data);
  mark.plotOptions = plotOptions;
  return [mark, ...extraMarks];
}

function prepareSpec(spec, containerWidth) {
  // Calculate width considering maxWidth constraint
  let width = spec.width ?? containerWidth;
  if (spec.maxWidth) {
    width = Math.min(width, spec.maxWidth);
  }

  const marks = spec.marks.flatMap((m) => readMark(m, width));
  spec = {
    ...spec,
    width: width,
    ...marks.reduce((acc, mark) => ({ ...acc, ...mark.plotOptions }), {}),
    marks: marks,
  };
  // if (!spec.height && containerHeight > 0 && !spec.aspectRatio) {
  //     spec.height = containerHeight
  // }

  if (spec.color_map) {
    const [domain, range] = [
      Object.keys(spec.color_map),
      Object.values(spec.color_map),
    ];
    spec.color = spec.color
      ? {
          ...spec.color,
          domain: [...(spec.color.domain || []), ...domain],
          range: [...(spec.color.range || []), ...range],
        }
      : { domain, range };
    delete spec.color_map;
  }
  return spec;
}

export function PlotWrapper({ spec }) {
  const $state = React.useContext($StateContext);
  return <PlotView spec={spec} $state={$state} />;
}
export function PlotView({ spec, $state }) {
  const [ref, containerWidth] = useContainerWidth();
  const done = useMemo(() => $state.beginUpdate("plot/PlotWrapper"), []);
  useEffect(() => {
    const parent = ref.current;
    if (parent && (containerWidth || spec.width)) {
      return mobx.autorun(() => {
        const preparedSpec = prepareSpec(spec, containerWidth);
        const startTime = performance.now();

        let plot = binding("$state", $state, () => Plot.plot(preparedSpec));
        const endTime = performance.now();
        // plot.setAttribute('data-render-time-ms', `${endTime - startTime}`);
        // console.log(`rendered plot in ${endTime - startTime}`)
        parent.innerHTML = "";
        if (spec.className) {
          plot.setAttribute("class", tw(spec.className));
        }
        if (spec.onPlotCreate) {
          spec.onPlotCreate(plot);
        }
        parent.appendChild(plot);
        done();
      });
    }
  }, [spec, containerWidth]);
  return (
    <div
      className={tw(`relative`)}
      style={spec.width && { width: spec.width }}
      ref={ref}
    ></div>
  );
}
