import {
    AbstractMesh,
    Color3,
    CubeTexture,
    CSG,
    HemisphericLight,
    ISoundOptions,
    Material,
    Matrix,
    Mesh,
    MeshBuilder,
    Node,
    Quaternion,
    Scene,
    SceneLoader,
    Sound,
    StandardMaterial,
    Texture,
    TransformNode,
    Vector2,
    Vector3,
} from '@babylonjs/core/Legacy/legacy';
import { Dataflow, IdentityMap, KeyedDictionary, Ref, Value, is, fromJS } from "@syndicate-lang/core";

import * as Shapes from './gen/shapes.js';
import { TurtleVM } from './turtle.js';

export class Environment {
    fields = new IdentityMap<symbol, Dataflow.Field<Value<Ref>>>();

    constructor(
        public spriteName: string,
        public formals: symbol[],
        public sceneDs: Ref,
    ) {
        formals.forEach(f => {
            field v: Value<Ref> = false;
            at this.sceneDs {
                on asserted Shapes.Variable({
                    "spriteName": this.spriteName,
                    "variable": f,
                    "value": $newValue,
                }) => {
                    v.value = newValue;
                }
            }
            this.fields.set(f, v);
        });
    }

    lookup(f: symbol, k: (v: Value<Ref>) => void) {
        dataflow {
            const v = this.fields.get(f);
            if (v === void 0) throw new Error(`Lookup of ${f.description} in ${this.spriteName} failed`);
            const w = v.value;
            if (w !== false) k(w);
        }
    }
}

export type ValueK = (v: Value<Ref>) => void;

export const activeFloorMeshes: Array<AbstractMesh> = [];
export const activeTouchableMeshes: Array<AbstractMesh> = [];

export class ShapeTree<N extends Node = Node> {
    shapePreserve: Value;
    cleanups: Array<() => void> = [];
    subnodes: Promise<N[]>;
    subtrees: ShapeTree[] = [];

    constructor (
        public scene: Scene,
        public shape: Shapes.Shape,
        public rootnode: N,
        subnodes?: Promise<N[]>,
    ) {
        this.shapePreserve = Shapes.fromShape(this.shape);
        const metadata = {};
        this.rootnode.metadata = metadata;
        this.subnodes = subnodes ?? Promise.resolve([]);
        this.subnodes.then(ns => ns.forEach(n => n.metadata = metadata));
    }

    get allnodes(): Promise<N[]> {
        return this.subnodes.then(ns => [this.rootnode, ... ns]);
    }

    reconcile(env: Environment, name: string, shape: Shapes.Shape): ShapeTree {
        if (is(Shapes.fromShape(shape), this.shapePreserve)) {
            return this;
        } else {
            this.remove();
            return build(env, name, this.scene, shape, {
                spriteName: async m => m.rootnode.metadata.spriteName = env.spriteName,
                collisions: async m => (await m.allnodes).forEach(n => n.checkCollisions = true),
            });
        }
    }

    set parent(t: ShapeTree) {
        t.subtrees.push(this);
        this.rootnode.parent = t.rootnode;
    }

    remove() {
        this.subtrees.forEach(t => t.remove());
        this.allnodes.then(ns => ns.forEach(n => n.dispose()));
        this.cleanups.forEach(c => c());
    }

    static empty(name: string, scene: Scene): ShapeTree {
        return ShapeTree.transform(name, scene, Shapes.Shape.many([]));
    }

    static transform(name: string, scene: Scene, shape: Shapes.Shape): ShapeTree<TransformNode> {
        return new ShapeTree(scene, shape, new TransformNode(name, scene));
    }
}

export function dv(env: Environment, v: Shapes.DoubleValue, k: (v: number) => void) {
    switch (v._variant) {
        case "immediate": k(v.value); break;
        case "reference": env.lookup(v.value, k as ValueK); break;
    }
}

export function u2(v: {x: number, y: number}): Shapes.ImmediateVector2 {
    return Shapes.ImmediateVector2({
        x: Shapes.DoubleValue.immediate(v.x),
        y: Shapes.DoubleValue.immediate(v.y),
    });
}

export function lv2(env: Environment, v: Shapes.ImmediateVector2, k: (v: Vector2) => void) {
    dv(env, v.x, x => dv(env, v.y, y => k(new Vector2(x, y))));
}

export function v2(env: Environment, v: Shapes.Vector2, k: (v: Vector2) => void) {
    switch (v._variant) {
        case "immediate": lv2(env, v.value, k); break;
        case "reference": env.lookup(v.value, v => lv2(env, Shapes.asImmediateVector2(v), k)); break;
    }
}

export function u3(v: {x: number, y: number, z: number}, scale = 1): Shapes.ImmediateVector3 {
    return Shapes.ImmediateVector3({
        x: Shapes.DoubleValue.immediate(v.x * scale),
        y: Shapes.DoubleValue.immediate(v.y * scale),
        z: Shapes.DoubleValue.immediate(v.z * scale),
    });
}

export function u3v(v: {x: number, y: number, z: number}, scale = 1): Shapes.Vector3 {
    return Shapes.Vector3.immediate(u3(v, scale));
}

export function lv3(env: Environment, v: Shapes.ImmediateVector3, k: (v: Vector3) => void) {
    dv(env, v.x, x => dv(env, v.y, y => dv(env, v.z, z => k(new Vector3(x, y, z)))));
}

export function v3(env: Environment, v: Shapes.Vector3, k: (v: Vector3) => void) {
    switch (v._variant) {
        case "immediate": lv3(env, v.value, k); break;
        case "reference": env.lookup(v.value, v => lv3(env, Shapes.asImmediateVector3(v), k)); break;
    }
}

export function lq(env: Environment, q: Shapes.ImmediateQuaternion, k: (v: Quaternion) => void) {
    dv(env, q.a, a => dv(env, q.b, b => dv(env, q.c, c => dv(env, q.d, d =>
        k(new Quaternion(a, b, c, d))))));
}

export function q(env: Environment, q: Shapes.Quaternion, k: (v: Quaternion) => void) {
    switch (q._variant) {
        case "immediate": lq(env, q.value, k); break;
        case "reference": env.lookup(q.value, q => lq(env, Shapes.asImmediateQuaternion(q), k)); break;
    }
}

export type MeshCustomizer = { [key: string]: ((m: ShapeTree<AbstractMesh>) => void) };

function applyCustomizer(m: ShapeTree<AbstractMesh>, c: MeshCustomizer) {
    Object.values(c).forEach(f => f(m));
}

type CachedTexture = {
    key: Value,
    referenceCount: number,
    material: Material,
};
const textureCache = new KeyedDictionary<Value, CachedTexture>();

function buildTexture(env: Environment, name: string, scene: Scene, spec: Shapes.TextureSpec): CachedTexture {
    const cacheKey = Shapes.fromTextureSpec(spec);
    const entry = textureCache.get(cacheKey);
    if (entry !== void 0) {
        entry.referenceCount++;
        return entry;
    }
    const mat = new StandardMaterial(name, scene);
    const tex = new Texture(spec.path, scene);
    mat.diffuseTexture = tex;
    switch (spec._variant) {
        case "simple":
            break;
        case "uvAlpha":
            mat.alpha = 0;
            dv(env, spec.alpha, a => mat.alpha = a);
            tex.hasAlpha = true;
            /* FALL THROUGH */
        case "uv":
            v2(env, spec.scale, scale => v2(env, spec.offset, offset => {
                tex.uScale = scale.x;
                tex.vScale = scale.y;
                tex.uOffset = offset.x;
                tex.vOffset = offset.y;
            }));
            break;
    }
    const newEntry = {
        key: cacheKey,
        referenceCount: 1,
        material: mat,
    };
    textureCache.set(cacheKey, newEntry);
    return newEntry;
}

function releaseTexture(entry: CachedTexture) {
    if (--entry.referenceCount === 0) {
        textureCache.delete(entry.key);
        entry.material.dispose();
    }
}

export type BuiltMesh = {
    rootnode: AbstractMesh,
    subnodes?: Promise<AbstractMesh[]>,
};

export function buildMesh(
    name: string,
    scene: Scene | null,
    meshSpec: Shapes.Mesh,
): BuiltMesh {
    switch (meshSpec._variant) {
        case "Sphere": return { rootnode: MeshBuilder.CreateSphere(name, {}, scene) };
        case "Box": return { rootnode: MeshBuilder.CreateBox(name, {}, scene) };
        case "Ground": return { rootnode: MeshBuilder.CreateGround(name, {}, scene ?? void 0) };
        case "Plane": return { rootnode: MeshBuilder.CreatePlane(name, {}, scene) };
        case "External": {
            const rootnode = new Mesh(name, scene);
            return {
                rootnode,
                subnodes: new Promise(async resolve => {
                    const r =
                        await SceneLoader.ImportMeshAsync("", meshSpec.value.path, void 0, scene);
                    r.meshes.forEach(m => m.parent = rootnode);
                    resolve(r.meshes);
                }),
            };
        }
        case "turtle": {
            const t = new TurtleVM(name, scene, meshSpec.value.program.map(fromJS));
            try {
                t.exec();
            } catch (e) {
                console.error(e);
            }
            return {
                rootnode: t.container,
                subnodes: Promise.resolve(t.meshes),
            };
        }
    }
}

export function buildSound(name: string, scene: Scene, spec: Shapes.SoundSpec, spatial: boolean): Sound {
    const options: ISoundOptions = {
        loop: true,
        autoplay: true,
        skipCodecCheck: true,
    };
    if (spatial) {
        options.distanceModel = "inverse";
        options.rolloffFactor = 0.25;
    }
    switch (spec._variant) {
        case "stream":
            options.streaming = true;
            break;
        case "loop":
            break;
    }
    return new Sound(name, spec.url, scene, null, options);
}

export function build(env: Environment, name: string, scene: Scene, shape: Shapes.Shape, customize: MeshCustomizer): ShapeTree {
    switch (shape._variant) {
        case "Mesh": {
            const m = buildMesh(name, scene, shape.value);
            const t = new ShapeTree<AbstractMesh>(scene, shape, m.rootnode, m.subnodes);
            applyCustomizer(t, customize);
            return t;
        }

        case "Light": {
            const light = new HemisphericLight(name, new Vector3(0, 1, 0), scene);
            v3(env, shape.value.v, v => light.direction = v);
            return new ShapeTree(scene, shape, light);
        }

        case "Scale": {
            const t = ShapeTree.transform(name, scene, shape);
            v3(env, shape.value.v, v => t.rootnode.scaling = v);
            build(env, name + '.inner', scene, shape.value.shape, customize).parent = t;
            return t;
        }

        case "Move": {
            const t = ShapeTree.transform(name, scene, shape);
            v3(env, shape.value.v, v => t.rootnode.position = v);
            build(env, name + '.inner', scene, shape.value.shape, customize).parent = t;
            return t;
        }

        case "Rotate": {
            const t = ShapeTree.transform(name, scene, shape);
            switch (shape.value._variant) {
                case "euler":
                    v3(env, shape.value.v, v => {
                        t.rootnode.rotation = v;
                        t.rootnode.rotation.scaleInPlace(2 * Math.PI);
                    });
                    break;
                case "quaternion":
                    q(env, shape.value.q, q => t.rootnode.rotationQuaternion = q);
                    break;
            }
            build(env, name + '.inner', scene, shape.value.shape, customize).parent = t;
            return t;
        }

        case "many": {
            const t = ShapeTree.transform(name, scene, shape);
            shape.value.forEach((s, i) => {
                build(env, name + '[' + i + ']', scene, s, customize).parent = t;
            });
            return t;
        }

        case "Texture": {
            const entry = buildTexture(env, name + '.texture', scene, shape.value.spec);
            const t = build(env, name + '.inner', scene, shape.value.shape, {
                ... customize,
                material: async m => (await m.allnodes).forEach(n => n.material = entry.material),
            });
            t.cleanups.push(() => releaseTexture(entry));
            return t;
        }

        case "Color": {
            const mat = new StandardMaterial(name + '.texture', scene);
            dv(env, shape.value.r, r =>
                dv(env, shape.value.g, g =>
                    dv(env, shape.value.b, b => mat.diffuseColor = new Color3(r, g, b))));
            if (shape.value._variant === "transparent") {
                dv(env, shape.value.alpha, a => mat.alpha = a);
            }
            const t = build(env, name + '.inner', scene, shape.value.shape, {
                ... customize,
                material: async m => (await m.allnodes).forEach(n => n.material = mat),
            });
            t.cleanups.push(() => mat.dispose());
            return t;
        }

        case "Sound": {
            const sound = buildSound(name + ".sound", scene, shape.value.spec, true);
            const t = build(env, name + ".inner", scene, shape.value.shape, {
                ... customize,
                sound: async m => sound.attachToMesh(m.rootnode),
            });
            t.cleanups.push(() => sound.dispose());
            return t;
        }

        case "Name":
            return build(env, name + '.' + shape.value.base, scene, shape.value.shape, customize);

        case "Floor":
            return build(env, name, scene, shape.value.shape, {
                ... customize,
                floor: async m => {
                    const nodes = await m.allnodes;
                    activeFloorMeshes.push(... nodes);
                    m.cleanups.push(() => {
                        const i = activeFloorMeshes.indexOf(nodes[0]);
                        if (i !== -1) activeFloorMeshes.splice(i, nodes.length);
                    });
                },
            });

        case "Nonphysical":
            return build(env, name, scene, shape.value.shape, {
                ... customize,
                collisions: async m => (await m.allnodes).forEach(n => n.checkCollisions = false),
            });

        case "Touchable":
            return build(env, name, scene, shape.value.shape, {
                ... customize,
                touchable: async m => {
                    const nodes = await m.allnodes;
                    m.rootnode.metadata.touchable = true;
                    activeTouchableMeshes.push(... nodes);
                    m.cleanups.push(() => {
                        const i = activeTouchableMeshes.indexOf(nodes[0]);
                        if (i !== -1) activeTouchableMeshes.splice(i, nodes.length);
                    });
                }
            });

        case "CSG": {
            const m = buildCSG(name, scene, shape.value.expr);
            const t = new ShapeTree<AbstractMesh>(scene, shape, m.rootnode, m.subnodes);
            applyCustomizer(t, customize);
            return t;
        }

        case "Skybox": {
            const t = new ShapeTree<AbstractMesh>(scene, shape, MeshBuilder.CreateBox(name, { size: 2000 }, scene));
            const mat = new StandardMaterial(name, scene);
            mat.backFaceCulling = false;
            mat.reflectionTexture = new CubeTexture(shape.value.path, scene);
            mat.reflectionTexture.coordinatesMode = Texture.SKYBOX_MODE;
            mat.diffuseColor = new Color3(0, 0, 0);
            mat.specularColor = new Color3(0, 0, 0);
            t.rootnode.material = mat;
            applyCustomizer(t, customize);
            return t;
        }

        default:
            ((_shape: never) => {
                console.error('Unsupported shape variant', shape);
                throw new Error("Unsupported shape variant");
            })(shape);
    }
}

export function buildCSG(name: string, scene: Scene, expr: Shapes.CSGExpr): BuiltMesh {
    async function walk(expr: Shapes.CSGExpr, matrix: Matrix): Promise<CSG> {
        switch (expr._variant) {
            case "mesh": {
                const m = buildMesh("", null, expr.shape);
                const nodes: Mesh[] =
                    ([m.rootnode, ... (m.subnodes ? await m.subnodes : [])]
                         .filter(n => n instanceof Mesh)) as Mesh[];
                nodes.forEach(n => {
                    n.freezeWorldMatrix(matrix, true);
                    n.unfreezeWorldMatrix();
                });
                const cs = nodes.flatMap(n => {
                    try {
                        return [CSG.FromMesh(n)];
                    } catch (_e) {
                        return [];
                    }
                });
                const c = cs[0];
                cs.slice(1).forEach(d => c.unionInPlace(d));
                nodes.forEach(n => n.dispose());
                return c;
            }
            case "scale":
                return walk(expr.shape, matrix.multiply(
                    Matrix.Scaling(expr.v.x, expr.v.y, expr.v.z)));
            case "move":
                return walk(expr.shape, matrix.multiply(
                    Matrix.Translation(expr.v.x, expr.v.y, expr.v.z)));
            case "rotate":
                return walk(expr.shape, matrix.multiply(
                    Matrix.RotationYawPitchRoll(expr.v.y * 2 * Math.PI,
                                                expr.v.x * 2 * Math.PI,
                                                expr.v.z * 2 * Math.PI)));
            case "subtract": {
                const c = await walk(expr.base, matrix);
                for (const d of expr.more) {
                    c.subtractInPlace(await walk(d, matrix));
                }
                return c;
            }
            case "union": {
                const c = await walk(expr.base, matrix);
                for (const d of expr.more) {
                    c.unionInPlace(await walk(d, matrix));
                }
                return c;
            }
            case "intersect": {
                const c = await walk(expr.base, matrix);
                for (const d of expr.more) {
                    c.intersectInPlace(await walk(d, matrix));
                }
                return c;
            }
            case "invert": {
                const c = await walk(expr.shape, matrix);
                c.inverseInPlace();
                return c;
            }
        }
    }
    const rootnode = new Mesh(name, scene);
    return {
        rootnode,
        subnodes: new Promise(async resolve => {
            const c = await walk(expr, Matrix.Identity());
            const m = c.toMesh(name + ".csg", null, scene);
            m.parent = rootnode;
            resolve([m]);
        }),
    };
}

export const builder = {
    sphere: () => Shapes.Shape.Mesh(Shapes.Mesh.Sphere(Shapes.Sphere())),
    box: () => Shapes.Shape.Mesh(Shapes.Mesh.Box(Shapes.Box())),
    plane: () => Shapes.Shape.Mesh(Shapes.Mesh.Plane(Shapes.Plane())),
    ground: () => Shapes.Shape.Mesh(Shapes.Mesh.Ground(Shapes.Ground())),
    light: (v: Shapes.Vector3) => Shapes.Shape.Light(Shapes.Light(v)),
    scale: (v: Shapes.Vector3, shape: Shapes.Shape) => Shapes.Shape.Scale(Shapes.Scale({ v, shape })),
    move: (v: Shapes.Vector3, shape: Shapes.Shape) => Shapes.Shape.Move(Shapes.Move({ v, shape })),
    rotate: (v: Shapes.Vector3, shape: Shapes.Shape) => Shapes.Shape.Rotate(Shapes.Rotate.euler({ v, shape })),
    many: (shapes: Shapes.Shape[]) => Shapes.Shape.many(shapes),
    texture: (spec: Shapes.TextureSpec, shape: Shapes.Shape) => Shapes.Shape.Texture(Shapes.Texture({ spec, shape })),
    color: (r0: number | symbol, g0: number | symbol, b0: number | symbol, shape: Shapes.Shape, alpha0: number | symbol = 1.0) => {
        function vd(x: number | symbol): Shapes.DoubleValue {
            return typeof x === 'number' ? Shapes.DoubleValue.immediate(x) : Shapes.DoubleValue.reference(x);
        }
        const r = vd(r0);
        const g = vd(g0);
        const b = vd(b0);
        const alpha = vd(alpha0);
        return Shapes.Shape.Color((alpha0 === 1.0)
            ? Shapes.Color.opaque({ r, g, b, shape })
            : Shapes.Color.transparent({ r, g, b, alpha, shape }));
    },
    name: (base: string, shape: Shapes.Shape) => Shapes.Shape.Name(Shapes.Name({ base, shape })),
    floor: (shape: Shapes.Shape) => Shapes.Shape.Floor(Shapes.Floor(shape)),
    nonphysical: (shape: Shapes.Shape) => Shapes.Shape.Nonphysical(Shapes.Nonphysical(shape)),
} satisfies { [key: string]: (... args: any[]) => Shapes.Shape };