import { useMemo } from "react";
import { Color, DoubleSide, ShaderMaterial, Vector3, MeshStandardMaterial } from "three";
import useShaderReader from "../../../../utils/scene/useShaderReader";

import edgeVertex from "../../../../assets/shaders/polygon/edge-fade1.vert.glsl";
import edgeFragment from "../../../../assets/shaders/polygon/edge-fade1.frag.glsl";

type PolygonMaterialProps = {
    points: { position: [number, number, number] }[];
    color?: string;
    transparent?: boolean;
    depthWrite?: boolean;
    depthTest?: boolean;
    opacity?: number;
    polygonOffset?: boolean;
    polygonOffsetFactor?: number;
    polygonOffsetUnits?: number;
    variant?: "edge" | "none";
} & (
    | {
          variant?: "none";
      }
    | {
          variant: "edge";
          fadeDistance?: number;
      }
);

function PolygonMaterial({
    points,
    color = "white",
    transparent = false,
    depthWrite = true,
    depthTest = true,
    opacity = 1,
    polygonOffset = false,
    polygonOffsetFactor,
    polygonOffsetUnits,
    variant = "none",
    ...rest
}: Readonly<PolygonMaterialProps>) {
    const vertexShader = useShaderReader(edgeVertex);
    const fragmentShader = useShaderReader(edgeFragment);

    const edgeProps = variant === "edge" ? (rest as { fadeDistance?: number }) : null;
    const fadeDistance = edgeProps?.fadeDistance ?? 2.0;

    const edges = useMemo(() => {
        return points.map((p, i) => {
            const next = points[(i + 1) % points.length];
            return [new Vector3(...p.position), new Vector3(...next.position)];
        });
    }, [points]);

    const polygonOffsetProps = polygonOffset
        ? {
              polygonOffsetFactor: polygonOffsetFactor ?? -1,
              polygonOffsetUnits: polygonOffsetUnits ?? -1,
          }
        : {};

    const basicMaterial = useMemo(() => {
        if (variant !== "none") return null;
        return new MeshStandardMaterial({
            color: new Color(color),
            side: DoubleSide,
            transparent,
            depthWrite,
            depthTest,
            polygonOffset,
            ...polygonOffsetProps,
            opacity,
        });
    }, [color, transparent, depthWrite, depthTest, polygonOffset, polygonOffsetProps, opacity, variant]);

    const edgeMaterial = useMemo(() => {
        if (!vertexShader || !fragmentShader || variant !== "edge") return null;

        const edgeArray: number[] = [];
        edges.forEach(([a, b]) => {
            edgeArray.push(a.x, a.z, b.x, b.z);
        });

        return new ShaderMaterial({
            side: DoubleSide,
            vertexShader,
            fragmentShader,
            uniforms: {
                uColor: { value: new Color(color) },
                uFadeDistance: { value: fadeDistance },
                uEdges: { value: edgeArray },
                uEdgeCount: { value: edges.length },
            },
            transparent,
            depthWrite,
            depthTest,
            polygonOffset,
            ...polygonOffsetProps,
            opacity,
        });
    }, [color, fadeDistance, vertexShader, fragmentShader, edges, transparent, depthWrite, depthTest, polygonOffset, polygonOffsetProps, opacity, variant]);

    if (variant === "none" && basicMaterial) return <primitive object={basicMaterial} attach="material" />;
    if (variant === "edge" && edgeMaterial) return <primitive object={edgeMaterial} attach="material" />;

    return null;
}

export default PolygonMaterial;
