import * as THREE from "three";
import h337 from "heatmap.js-fix";
import { deepAssign, deepEqual } from "@/utils";

type HeatmapRenderer = {
    canvas: HTMLCanvasElement;
};

type HeatmapInstance = {
    setData: (data: { max: number; min?: number; data: HeatmapCanvasPoint[] }) => void;
    configure?: (config: Record<string, unknown>) => void;
    _renderer?: HeatmapRenderer;
    _store?: { _cfgRadius?: number };
};

type HeatmapCanvasPoint = {
    x: number;
    y: number;
    value: number;
    radius?: number;
};

export const getDefaultHeatmapOptions = (): IHeatmap.options => ({
    name: "Heatmap",
    position: [0, 0, 0],
    mode: "flat",
    size: {
        width: 10,
        height: 10,
    },
    resolution: {
        width: 512,
        height: 512,
    },
    material: {
        transparent: true,
        opacity: 1,
        depthWrite: false,
        depthTest: true,
        side: "double",
    },
    height: {
        scale: 2,
        segments: {
            width: 128,
            height: 128,
        },
    },
    heatmap: {
        radius: 40,
        blur: 0.85,
        maxOpacity: 1,
        minOpacity: 0,
        gradient: {
            "0.0": "#2c7bb6",
            "0.5": "#ffffbf",
            "1.0": "#d7191c",
        },
    },
    data: {
        max: 1,
        min: 0,
        points: [],
    },
});

const DEFAULT_HEATMAP_RADIUS = getDefaultHeatmapOptions().heatmap.radius ?? 0;

export default class Heatmap extends THREE.Mesh {
    type = "Heatmap";
    isHeatmap = true;

    options: IHeatmap.options = getDefaultHeatmapOptions();
    heatmap: HeatmapInstance;
    container: HTMLDivElement;
    canvas: HTMLCanvasElement;
    texture: THREE.CanvasTexture;

    constructor(options: Partial<IHeatmap.options> = {}) {
        super();

        if (typeof document === "undefined") {
            throw new Error("[Astral 3D]: Heatmap requires a DOM environment.");
        }

        deepAssign(this.options, options);
        if (options.heatmap?.gradient) {
            this.options.heatmap.gradient = { ...options.heatmap.gradient };
        }
        this.name = this.options.name;

        this.container = this.createContainer();
        this.heatmap = this.createHeatmap();
        this.canvas = this.resolveCanvas();
        this.texture = new THREE.CanvasTexture(this.canvas);
        this.texture.colorSpace = THREE.SRGBColorSpace;

        this.material = new THREE.MeshBasicMaterial({
            map: this.texture,
            transparent: this.options.material.transparent,
            opacity: this.options.material.opacity,
            depthWrite: this.options.material.depthWrite,
            depthTest: this.options.material.depthTest,
            side: this.resolveSide(this.options.material.side),
        });

        this.geometry = this.createGeometry();

        const position = this.options.position;
        this.position.set(position[0] || 0, position[1] || 0, position[2] || 0);
        // Default to horizontal (XZ) plane with Y-up.
        this.rotation.x = -Math.PI / 2;

        if (this.options.data?.points?.length) {
            this.setData(this.options.data);
        } else {
            this.heatmap.setData({
                max: this.options.data.max || 1,
                min: this.options.data.min || 0,
                data: [],
            });
        }
    }

    private createContainer() {
        const container = document.createElement("div");
        const { width, height } = this.options.resolution;
        container.style.position = "absolute";
        container.style.left = "-10000px";
        container.style.top = "-10000px";
        container.style.width = `${width}px`;
        container.style.height = `${height}px`;
        container.style.visibility = "hidden";
        container.style.pointerEvents = "none";

        const host = document.body || document.documentElement;
        host.appendChild(container);

        return container;
    }

    private createHeatmap() {
        const config = {
            container: this.container,
            radius: this.options.heatmap.radius,
            blur: this.options.heatmap.blur,
            maxOpacity: this.options.heatmap.maxOpacity,
            minOpacity: this.options.heatmap.minOpacity,
            gradient: this.options.heatmap.gradient,
        };

        return (h337 as { create: (cfg: Record<string, unknown>) => HeatmapInstance }).create(config);
    }

    private clearContainer() {
        while (this.container.firstChild) {
            this.container.removeChild(this.container.firstChild);
        }
    }

    private refreshCanvas() {
        const canvas = this.resolveCanvas();
        if (canvas !== this.canvas) {
            this.canvas = canvas;
            this.texture.image = this.canvas;
        }
        this.texture.needsUpdate = true;
    }

    private rebuildHeatmap() {
        this.clearContainer();
        this.heatmap = this.createHeatmap();
        this.refreshCanvas();
    }

    private syncHeatmapRuntimeConfig() {
        const heatmapAny = this.heatmap as HeatmapInstance | undefined;
        if (!heatmapAny) return;

        if (heatmapAny._store && typeof this.options.heatmap.radius === "number") {
            heatmapAny._store._cfgRadius = this.options.heatmap.radius;
        }

        const rendererAny = heatmapAny._renderer as { _templates?: Record<string, HTMLCanvasElement> } | undefined;
        if (rendererAny && rendererAny._templates) {
            rendererAny._templates = {};
        }
    }

    private getSegments() {
        if (this.options.mode !== "height") {
            return { width: 1, height: 1 };
        }

        const segments = this.options.height?.segments || { width: 1, height: 1 };
        return {
            width: Math.max(1, Math.floor(segments.width)),
            height: Math.max(1, Math.floor(segments.height)),
        };
    }

    private createGeometry() {
        const size = this.options.size;
        const segments = this.getSegments();
        return new THREE.PlaneGeometry(size.width, size.height, segments.width, segments.height);
    }

    private ensureGeometry() {
        const size = this.options.size;
        const segments = this.getSegments();
        const geometry = this.geometry as THREE.PlaneGeometry | undefined;
        const parameters = geometry?.parameters;

        const needsRebuild =
            !parameters ||
            parameters.width !== size.width ||
            parameters.height !== size.height ||
            parameters.widthSegments !== segments.width ||
            parameters.heightSegments !== segments.height;

        if (!needsRebuild) return;

        this.geometry.dispose();
        this.geometry = new THREE.PlaneGeometry(size.width, size.height, segments.width, segments.height);
    }

    private resolveCanvas() {
        const canvas = this.heatmap._renderer?.canvas || this.container.querySelector("canvas");
        if (!canvas) {
            const fallback = document.createElement("canvas");
            this.container.appendChild(fallback);
            return fallback;
        }

        canvas.width = this.options.resolution.width;
        canvas.height = this.options.resolution.height;

        return canvas;
    }

    private flattenGeometry() {
        const geometry = this.geometry as THREE.BufferGeometry;
        const position = geometry.getAttribute("position") as THREE.BufferAttribute | undefined;
        if (!position) return;

        for (let i = 0; i < position.count; i += 1) {
            position.setZ(i, 0);
        }

        position.needsUpdate = true;
        geometry.computeVertexNormals();
    }

    private updateRelief() {
        this.ensureGeometry();

        if (this.options.mode !== "height") {
            this.flattenGeometry();
            return;
        }

        const context = this.canvas.getContext("2d");
        if (!context) return;

        const canvasWidth = this.canvas.width;
        const canvasHeight = this.canvas.height;
        const imageData = context.getImageData(0, 0, canvasWidth, canvasHeight).data;
        const geometry = this.geometry as THREE.BufferGeometry;
        const position = geometry.getAttribute("position") as THREE.BufferAttribute | undefined;
        const uv = geometry.getAttribute("uv") as THREE.BufferAttribute | undefined;
        if (!position || !uv) return;

        const heightScale = this.options.height?.scale ?? 1;
        for (let i = 0; i < position.count; i += 1) {
            const u = uv.getX(i);
            const v = uv.getY(i);
            const x = Math.round(this.clamp(u * (canvasWidth - 1), 0, canvasWidth - 1));
            const y = Math.round(this.clamp((1 - v) * (canvasHeight - 1), 0, canvasHeight - 1));
            const index = (y * canvasWidth + x) * 4;
            const alpha = imageData[index + 3] / 255;
            position.setZ(i, alpha * heightScale);
        }

        position.needsUpdate = true;
        geometry.computeVertexNormals();
    }

    private resolveSide(side?: IHeatmap.Side) {
        switch (side) {
            case "front":
                return THREE.FrontSide;
            case "back":
                return THREE.BackSide;
            case "double":
            default:
                return THREE.DoubleSide;
        }
    }

    private clamp(value: number, min: number, max: number) {
        return Math.min(Math.max(value, min), max);
    }

    private toCanvasPoint(point: IHeatmap.Point): HeatmapCanvasPoint {
        const size = this.options.size;
        const resolution = this.options.resolution;
        const safeWidth = Math.max(size.width, 0.000001);
        const safeHeight = Math.max(size.height, 0.000001);
        const scaleX = resolution.width / safeWidth;
        const scaleY = resolution.height / safeHeight;

        // Map local plane coordinates (centered at origin) to heatmap pixels.
        const x = (point.x + size.width / 2) * scaleX;
        const y = (size.height / 2 - point.y) * scaleY;

        const canvasPoint: HeatmapCanvasPoint = {
            x: this.clamp(x, 0, resolution.width),
            y: this.clamp(y, 0, resolution.height),
            value: point.value,
        };

        if (typeof point.radius === "number") {
            const radiusScale = Math.min(scaleX, scaleY);
            const baseRadius = typeof this.options.heatmap.radius === "number" ? this.options.heatmap.radius : DEFAULT_HEATMAP_RADIUS;
            const radiusFactor = DEFAULT_HEATMAP_RADIUS > 0 ? Math.max(0, baseRadius) / DEFAULT_HEATMAP_RADIUS : 1;
            canvasPoint.radius = point.radius * radiusScale * radiusFactor;
        }

        return canvasPoint;
    }

    private buildHeatmapData(data: IHeatmap.Data) {
        const points = data.points || [];
        const values = points.map((point) => point.value).filter((value) => Number.isFinite(value));
        const max = typeof data.max === "number" ? data.max : values.length ? Math.max(...values) : 1;
        const min = typeof data.min === "number" ? data.min : values.length ? Math.min(...values) : 0;

        return {
            max,
            min,
            data: points.map((point) => this.toCanvasPoint(point)),
        };
    }

    setData(data: IHeatmap.Data) {
        this.options.data = {
            max: data.max,
            min: data.min,
            points: Array.isArray(data.points) ? data.points.slice() : [],
        };

        const heatmapData = this.buildHeatmapData(this.options.data);
        this.heatmap.setData(heatmapData);
        this.texture.needsUpdate = true;
        this.updateRelief();
    }

    addData(points: IHeatmap.Point | IHeatmap.Point[]) {
        const list = Array.isArray(points) ? points : [points];
        const next = this.options.data.points.concat(list);
        this.setData({
            ...this.options.data,
            points: next,
        });
    }

    clear() {
        this.setData({
            max: 1,
            min: 0,
            points: [],
        });

        return this;
    }

    setSize(size: IHeatmap.Size) {
        this.options.size = { width: size.width, height: size.height };
        this.ensureGeometry();
        this.setData(this.options.data);
    }

    setMode(mode: IHeatmap.Mode) {
        if (this.options.mode === mode) return this;

        this.options.mode = mode;
        this.ensureGeometry();
        this.updateRelief();
        return this;
    }

    updateHeatmapConfig(config: Partial<IHeatmap.HeatmapConfig>) {
        const previousGradient = this.options.heatmap.gradient;
        const { gradient, ...rest } = config;
        deepAssign(this.options.heatmap, rest);
        if (gradient && typeof gradient === "object" && !Array.isArray(gradient)) {
            this.options.heatmap.gradient = { ...gradient };
        }
        const gradientChanged = !deepEqual(previousGradient, this.options.heatmap.gradient);

        this.syncHeatmapRuntimeConfig();
        if (gradientChanged) {
            this.rebuildHeatmap();
        } else if (this.heatmap.configure) {
            this.heatmap.configure({
                radius: this.options.heatmap.radius,
                blur: this.options.heatmap.blur,
                maxOpacity: this.options.heatmap.maxOpacity,
                minOpacity: this.options.heatmap.minOpacity,
                gradient: this.options.heatmap.gradient,
            });
            this.refreshCanvas();
        } else {
            this.rebuildHeatmap();
        }

        this.setData(this.options.data);
    }

    toJSON(meta?: THREE.JSONMeta) {
        const options = JSON.parse(JSON.stringify(this.options));
        options.position = this.position.toArray();

        // 同步材质属性
        options.material = {
            transparent: (this.material as THREE.Material).transparent,
            opacity: (this.material as THREE.Material).opacity,
            depthWrite: (this.material as THREE.Material).depthWrite,
            depthTest: (this.material as THREE.Material).depthTest,
            side: this.options.material.side,
        };

        const superJSON = super.toJSON(meta);
        superJSON.object.type = this.type;
        superJSON.object.options = options;

        return superJSON;
    }

    static fromJSON(json: { options: IHeatmap.options }) {
        return new Heatmap(json.options);
    }

    dispose() {
        this.geometry.dispose();
        this.texture.dispose();
        (this.material as THREE.Material).dispose();

        if (this.container.parentElement) {
            this.container.parentElement.removeChild(this.container);
        }
    }
}