import { Color, Mesh, MeshBasicMaterial, Object3D } from "three";

import { GameObject } from "../../../Component.js";
import { SpriteRenderer } from "../../../SpriteRenderer.js";
import { $shadowDomOwner, BaseUIComponent } from "../../../ui/BaseUIComponent.js";
import { Canvas } from "../../../ui/Canvas.js";
import { RenderMode } from "../../../ui/Canvas.js";
import { CanvasGroup } from "../../../ui/CanvasGroup.js";
import { RectTransform } from "../../../ui/RectTransform.js";
import { Text } from "../../../ui/Text.js";
import type { IUSDExporterExtension } from "../Extension.js";
import { USDObject, USDZExporterContext } from "../ThreeUSDZExporter.js";
import { TextExtension } from "./USDZText.js";

export class USDZUIExtension implements IUSDExporterExtension {
    get extensionName(): string {
        return "tmui";
    }

    // TODO would probably be better to export each object instead of the entire Canvas
    // so that we don't export them twice (once as regular hierarchy, once as part of Canvas export)
    onExportObject(object: Object3D, model: USDObject, _context: USDZExporterContext) {
        
        const canvas = GameObject.getComponent(object, Canvas);

        if (canvas && canvas.enabled && canvas.renderMode === RenderMode.WorldSpace) {

            const textExt = new TextExtension();
            const rt = GameObject.getComponent(object, RectTransform);
            const canvasGroup = GameObject.getComponent(object, CanvasGroup);

            // we have to do some temporary changes (enable UI component so that they're building
            // their shadow hierarchy, then revert them back to the original state)
            const revertActions = new Array<() => void>();

            let width = 100;
            let height = 100;
            if (rt) {

                // Workaround: since UI components are only present in the shadow hierarchy when objects are on,
                // we need to enable them temporarily to export them including their potential child components.
                // For example, Text can have multiple child objects (e.g. rich text, panels, ...)
                if (!GameObject.isActiveSelf(object)) {
                    const wasActive = GameObject.isActiveSelf(object);
                    GameObject.setActive(object, true);
                    rt.onEnable();
                    rt.updateTransform();

                    revertActions.push(() => {
                        rt.onDisable();
                        GameObject.setActive(object, wasActive);
                    });
                }

                object.traverse((child) => {
                    if (!GameObject.isActiveInHierarchy(child)) {
                        const wasActive = GameObject.isActiveSelf(child);
                        GameObject.setActive(child, true);
                        const baseUIComponent = GameObject.getComponent(child, BaseUIComponent);
                        if (baseUIComponent) {
                            baseUIComponent.onEnable();
                            revertActions.push(() => {
                                baseUIComponent.onDisable();
                            });
                        }

                        const rectTransform = GameObject.getComponent(child, RectTransform);
                        if (rectTransform) {
                            rectTransform.onEnable();
                            rectTransform.updateTransform();
                            // This method bypasses the checks for whether the object is enabled etc.
                            // so that we can ensure even a disabled object has the correct layout.
                            rectTransform["onApplyTransform"]();
                            revertActions.push(() => {
                                rectTransform.onDisable();
                            });
                        }

                        const text = GameObject.getComponent(child, Text);
                        if (text) {
                            text.onEnable();
                            revertActions.push(() => {
                                text.onDisable();
                            });
                        }

                        revertActions.push(() => {
                            GameObject.setActive(child, wasActive);
                        });
                    }
                });

                width = rt.width;
                height = rt.height;

                const shadowRootModel = USDObject.createEmpty();
                const shadowComponent = rt.shadowComponent as unknown as Object3D;
                model.add(shadowRootModel);

                if (shadowComponent) {
                    const mat = shadowComponent.matrix;
                    shadowRootModel.setMatrix(mat);

                    const usdObjectMap = new Map<Object3D, USDObject>();
                    const opacityMap = new Map<Object3D, number>();
                    usdObjectMap.set(shadowComponent, shadowRootModel);
                    opacityMap.set(shadowComponent, canvasGroup ? canvasGroup.alpha : 1);

                    shadowComponent.traverse((child) => {
                        // console.log("traversing UI shadow components", shadowComponent.name + " ->" + child.name);
                        if (child === shadowComponent) return;

                        const childModel = USDObject.createEmpty();
                        childModel.setMatrix(child.matrix);

                        const childParent = child.parent;
                        const isText = !!childParent && typeof childParent["textContent"] === "string" && childParent["textContent"].length > 0;
                        let hierarchyOpacity = opacityMap.get(childParent!) || 1;
                        
                        // TODO CanvasGroup doesn't render something but modifies opacity
                        // get CanvasGroup and modify alpha here

                        const canvasGroup = GameObject.getComponent(child, CanvasGroup);
                        if (canvasGroup)
                            hierarchyOpacity *= canvasGroup.alpha;

                        if (child instanceof Mesh && isText) {
                            // get shadoDomOwner so we can export Text from the text extension directly
                            const shadowDomOwner = child[$shadowDomOwner];
                            if (!shadowDomOwner)
                                console.error("Error when exporting UI: shadow component owner not found. This is likely a bug.", child);
                            else
                                textExt.exportText(shadowDomOwner.gameObject, childModel, _context);
                        }

                        if (child instanceof Mesh && !isText)
                        {
                            // UI behaves weird: it seems it's always flipped right now,
                            // and three magically fixes it when rendering
                            // see https://github.com/mrdoob/three.js/pull/12720#issue-275923930
                            // https://github.com/mrdoob/three.js/issues/17361
                            // So we need to fix the winding order after applying the negative scale.
                            const clonedGeo = child.geometry.clone();
                            clonedGeo.scale(1, 1, -1);
                            this.flipWindingOrder(clonedGeo);
                            childModel.geometry = clonedGeo;

                            const color = new Color();
                            const ownOpacity = child.material.opacity;
                            color.copy(child.material.color);

                            // Calculate opacity from parent chain and own alpha
                            childModel.material = new MeshBasicMaterial({
                                color: color,
                                opacity: ownOpacity * hierarchyOpacity,
                                map: child.material.map,
                                transparent: true,
                            });
                        }

                        usdObjectMap.set(child, childModel);
                        opacityMap.set(child, hierarchyOpacity);

                        const parentUsdzObject = usdObjectMap.get(childParent!);
                        if (!parentUsdzObject) {
                            console.error("Error when exporting UI: shadow component parent not found!", child, child.parent);
                            return;
                        }
                        parentUsdzObject.add(childModel);
                    });
                }
            }

            // revert temporary changes that we did here
            for (const revert of revertActions) {
                revert();
            }
        }
    }
    
    private flipWindingOrder(geometry) {
        const index = geometry.index.array
        for (let i = 0, il = index.length / 3; i < il; i++) {
          const x = index[i * 3]
          index[i * 3] = index[i * 3 + 2]
          index[i * 3 + 2] = x
        }
        geometry.index.needsUpdate = true
    }
}