import * as THREE from 'three';

import {AI} from '../../ai/AI';
import {AIOptions} from '../../ai/AIOptions';
import {
  detectDeviceCameraTarget,
  getCameraParametersSnapshot,
  type CameraParametersSnapshot,
} from '../../camera/CameraUtils';
import {XRDeviceCamera} from '../../camera/XRDeviceCamera';
import {
  RendererHolder,
  type WebGLOrWebGPURenderer,
} from '../../core/RendererTypes';
import {Script} from '../../core/Script';
import {Depth} from '../../depth/Depth';
import {
  disposeMaterial,
  disposeObjectChildren,
  disposeObjectTree,
} from '../../utils/ThreeDisposal';
import {WorldOptions} from '../WorldOptions';
import {DetectedObject} from './DetectedObject';
import {
  BaseDetectorBackend,
  DetectorBackendContext,
} from './ObjectDetectorBackend';
import {GeminiDetectorBackend} from './backends/GeminiDetectorBackend';
import {MediaPipeDetectorBackend} from './backends/MediaPipeDetectorBackend';

/**
 * Represents a detected object in a normalized format, independent of the specific detector backend used.
 * Coordinates are normalized typically in the range [0, 1].
 *
 * T - The type of additional data associated with the detected object.
 */
export interface NormalizedDetectedObject<T> {
  ymin: number;
  xmin: number;
  ymax: number;
  xmax: number;
  objectName: string;
  additionalData?: T;
}

/**
 * Represents a snapshot taken from the device camera.
 * Can contain either a base64 encoded image string or raw ImageData.
 */
export interface CameraSnapshot {
  base64?: string;
  imageData?: ImageData;
}

/** Inputs for one object detection run, without changing shared options. */
export interface ObjectDetectionOptions {
  /** Uses this backend for this run only; otherwise uses the configured backend. */
  backend?: WorldOptions['objects']['backendConfig']['activeBackend'];
  /**
   * Uses this frame instead of capturing another one. Gemini requires
   * `base64`; MediaPipe requires `imageData`. Both may be supplied for the
   * same frame. Capture and submit in the same frame so the camera pose and
   * depth mesh match, and do not modify the pixels until the run completes.
   */
  snapshot?: CameraSnapshot;
}

interface ObjectDetectionFrame {
  cameraParametersSnapshot: CameraParametersSnapshot;
  depthMeshSnapshot: THREE.Mesh;
}

export interface SimulatorDetectedObjectInput<T = unknown> {
  label: string;
  position: THREE.Vector3;
  boundingBox: THREE.Box2;
  data?: T;
}

export interface SimulatorObjectDetectionSource {
  detect(): SimulatorDetectedObjectInput[];
}

/**
 * Detects objects in the user's environment using a specified backend.
 * It queries an AI model with the device camera feed and returns located
 * objects with 2D and 3D positioning data.
 */
export class ObjectDetector extends Script {
  static dependencies = {
    options: WorldOptions,
    ai: AI,
    aiOptions: AIOptions,
    deviceCamera: XRDeviceCamera,
    depth: Depth,
    camera: THREE.Camera,
    rendererHolder: RendererHolder,
  };

  /**
   * A map from the object's UUID to our custom `DetectedObject` instance.
   */
  private _detectedObjects = new Map<string, DetectedObject<unknown>>();
  private _detectorBackends = new Map<
    string,
    Promise<BaseDetectorBackend<unknown>>
  >();
  private activeClients = new Set<object>();
  private currentDetectionPromise: Promise<DetectedObject<unknown>[]> | null =
    null;
  private pendingDetectionPromise: Promise<DetectedObject<unknown>[]> | null =
    null;
  private lastContinuousDetectionStartedAtMs = -Infinity;
  private initialized = false;
  private disposed = false;
  private simulatorSource?: SimulatorObjectDetectionSource;

  private _debugVisualsGroup?: THREE.Group;

  /**
   * The latest detected objects.
   */
  public detectedObjects: DetectedObject<unknown>[] = [];

  // Injected dependencies
  private options!: WorldOptions;
  private ai!: AI;
  private aiOptions!: AIOptions;
  private deviceCamera!: XRDeviceCamera;
  private depth!: Depth;
  private camera!: THREE.PerspectiveCamera;
  private renderer!: WebGLOrWebGPURenderer;

  /**
   * Target device profile used to look up RGB camera intrinsics and pose
   * for converting detection bounding boxes into world space. Defaults to
   * `'galaxyxr'`; auto-overridden to `'quest3'` in {@link init} when the
   * Meta Quest browser is detected. Can be overridden manually before init.
   */
  targetDevice = 'galaxyxr';

  /**
   * Initializes the ObjectDetector.
   * @override
   */
  init({
    options,
    ai,
    aiOptions,
    deviceCamera,
    depth,
    camera,
    rendererHolder,
  }: {
    options: WorldOptions;
    ai: AI;
    aiOptions: AIOptions;
    deviceCamera: XRDeviceCamera;
    depth: Depth;
    camera: THREE.PerspectiveCamera;
    rendererHolder: RendererHolder;
  }) {
    this.options = options;
    this.ai = ai;
    this.aiOptions = aiOptions;
    this.deviceCamera = deviceCamera;
    this.depth = depth;
    this.camera = camera;
    this.renderer = rendererHolder.renderer;
    this.initialized = true;
    this.disposed = false;

    if (this.targetDevice === 'galaxyxr') {
      this.targetDevice = detectDeviceCameraTarget();
    }

    if (this.options.objects.showDebugVisualizations) {
      this._debugVisualsGroup = new THREE.Group();
      // Disable raycasting for the debug group to prevent interaction errors.
      this._debugVisualsGroup.raycast = () => {};
      this.add(this._debugVisualsGroup);
    }
  }

  /**
   * Starts continuous object detection for the given client.
   * Detection starts on the next update after initialization.
   * @param client - The client object requesting object detection.
   */
  start(client: object): void {
    if (this.activeClients.has(client)) {
      return;
    }
    this.activeClients.add(client);
  }

  /**
   * Stops continuous object detection for the given client.
   * If this was the last client, stops the background detection loop.
   * @param client - The client object that no longer needs object detection.
   */
  stop(client: object): void {
    this.activeClients.delete(client);
  }

  /**
   * Called per frame by the engine. If there are active clients,
   * ensures the continuous object detection is running.
   */
  override update() {
    if (
      !this.initialized ||
      this.activeClients.size === 0 ||
      this.currentDetectionPromise ||
      this.pendingDetectionPromise
    ) {
      return;
    }

    const pollingIntervalMs = this.options.objects.pollingIntervalMs;
    if (
      pollingIntervalMs > 0 &&
      performance.now() - this.lastContinuousDetectionStartedAtMs <
        pollingIntervalMs
    ) {
      return;
    }

    this.runContinuousDetection();
  }

  private runContinuousDetection() {
    if (this.currentDetectionPromise || this.pendingDetectionPromise) {
      return;
    }
    this.lastContinuousDetectionStartedAtMs = performance.now();
    this.currentDetectionPromise = this.runDetectionInternal()
      .then((results) => {
        if (!this.disposed) {
          this.detectedObjects = results;
        }
        return results;
      })
      .finally(() => {
        this.currentDetectionPromise = null;
      });
  }

  /**
   * Runs object detection or returns the ongoing detection promise.
   *
   * - If continuous detection is started (has active clients), returns the
   *   promise for the next detection result.
   * - If continuous detection is not started, performs a one-off detection and
   *   returns the result. If a one-off detection is already in progress, returns
   *   the promise for that ongoing detection.
   * - Supplying a snapshot or backend queues a separate one-off run after
   *   pending detections, without changing the continuous detector's options.
   *   Supplied snapshots retain the camera pose and depth at submission time.
   *   Simulator ground truth, when installed, still takes precedence.
   *
   * @param options - Optional inputs for this run only.
   * @returns A promise that resolves with an
   * array of detected `DetectedObject` instances.
   * @throws If an explicit backend or snapshot is invalid, or the detector
   * is disposed before a queued request starts.
   */
  runDetection<T = null>(
    options: ObjectDetectionOptions = {}
  ): Promise<DetectedObject<T>[]> {
    if (this.disposed) {
      return Promise.reject(new Error('ObjectDetector has been disposed.'));
    }
    if (!this.initialized) {
      return Promise.reject(new Error('ObjectDetector is not initialized.'));
    }
    if (options.backend !== undefined || options.snapshot !== undefined) {
      return this.runDetectionWithOverrides<T>(options);
    }
    if (this.currentDetectionPromise) {
      return this.currentDetectionPromise as Promise<DetectedObject<T>[]>;
    }
    if (this.pendingDetectionPromise) {
      return this.pendingDetectionPromise as Promise<DetectedObject<T>[]>;
    }
    if (this.activeClients.size > 0) {
      this.runContinuousDetection();
      return this.currentDetectionPromise! as Promise<DetectedObject<T>[]>;
    }
    return this.runOneOffDetection() as Promise<DetectedObject<T>[]>;
  }

  private async runDetectionWithOverrides<T>(
    options: ObjectDetectionOptions
  ): Promise<DetectedObject<T>[]> {
    const backend =
      options.backend ?? this.options.objects.backendConfig.activeBackend;
    if (backend !== 'gemini' && backend !== 'mediapipe') {
      throw new Error(`ObjectDetector backend '${backend}' is not supported.`);
    }
    const snapshot =
      options.snapshot === undefined ? undefined : {...options.snapshot};
    const requiredField = backend === 'gemini' ? 'base64' : 'imageData';
    if (snapshot && !snapshot[requiredField]) {
      throw new Error(
        `ObjectDetector snapshot for '${backend}' must include ${requiredField}.`
      );
    }

    const frame =
      snapshot && !this.simulatorSource
        ? this.captureDetectionFrame()
        : undefined;
    const run = () => this.runOneOffDetection({backend, snapshot}, frame);
    const previous =
      this.pendingDetectionPromise ?? this.currentDetectionPromise;
    // A failed request rejects its own caller, but must not block later runs.
    const pending = (previous ? previous.then(run, run) : run()).finally(() => {
      if (this.pendingDetectionPromise === pending) {
        this.pendingDetectionPromise = null;
      }
    });
    this.pendingDetectionPromise = pending;
    return pending as Promise<DetectedObject<T>[]>;
  }

  private runOneOffDetection(
    options: ObjectDetectionOptions = {},
    frame?: ObjectDetectionFrame | null
  ): Promise<DetectedObject<unknown>[]> {
    const current = this.runDetectionInternal<unknown>(options, frame).finally(
      () => {
        if (this.currentDetectionPromise === current) {
          this.currentDetectionPromise = null;
        }
      }
    );
    this.currentDetectionPromise = current;
    return current;
  }

  /** Installs or removes the desktop simulator's ground-truth detector. */
  setSimulatorSource(source?: SimulatorObjectDetectionSource) {
    this.simulatorSource = source;
    return this;
  }

  private async runDetectionInternal<T = null>(
    options: ObjectDetectionOptions = {},
    frame?: ObjectDetectionFrame | null
  ): Promise<DetectedObject<T>[]> {
    let detectionFrame = frame;
    try {
      if (this.disposed) {
        throw new Error('ObjectDetector has been disposed.');
      }
      this.clearDetectedObjects(); // Clear previous scene results before starting a new detection.

      if (this.simulatorSource) {
        const detectedObjects = this.simulatorSource.detect().map((input) => {
          const object = new DetectedObject(
            input.label,
            null,
            input.boundingBox,
            input.data ?? {}
          );
          object.position.copy(input.position);
          return object;
        });
        for (const object of detectedObjects) {
          this._detectedObjects.set(object.uuid, object);
          this.add(object);
        }
        return detectedObjects as DetectedObject<T>[];
      }

      if (detectionFrame === undefined) {
        detectionFrame = this.captureDetectionFrame();
      }
      if (!detectionFrame) {
        // Device camera not ready yet (warming up); skip until it is available.
        return [];
      }

      const context = this.getDetectorContext();
      const activeBackend =
        options.backend ?? this.options.objects.backendConfig.activeBackend;
      const detectorBackendPromise = this.getOrCreateDetectorBackend<T>(
        activeBackend,
        context
      );

      let detectorBackend: BaseDetectorBackend<T>;
      try {
        detectorBackend = await detectorBackendPromise;
      } catch (error) {
        console.warn(
          `Failed to load or initialize ObjectDetector backend '${activeBackend}':`,
          error
        );
        return [];
      }
      if (this.disposed) {
        return [];
      }
      const detectedObjects = await detectorBackend.run(
        detectionFrame.depthMeshSnapshot,
        detectionFrame.cameraParametersSnapshot,
        options.snapshot
      );
      if (this.disposed) {
        return [];
      }
      for (const obj of detectedObjects) {
        this._detectedObjects.set(obj.uuid, obj);
        this.add(obj);
      }
      return detectedObjects;
    } finally {
      if (detectionFrame) {
        this.disposeDepthMeshSnapshot(detectionFrame.depthMeshSnapshot);
      }
    }
  }

  private captureDetectionFrame(): ObjectDetectionFrame | null {
    const cameraParametersSnapshot = getCameraParametersSnapshot(
      this.camera,
      this.renderer.xr.getCamera() as THREE.WebXRArrayCamera,
      this.deviceCamera,
      this.targetDevice
    );
    return cameraParametersSnapshot
      ? {
          cameraParametersSnapshot,
          depthMeshSnapshot: this.getDepthMeshSnapshot(),
        }
      : null;
  }

  private getDetectorContext(): DetectorBackendContext {
    return {
      options: this.options,
      ai: this.ai,
      aiOptions: this.aiOptions,
      deviceCamera: this.deviceCamera,
      debugVisualsGroup: this._debugVisualsGroup,
    };
  }

  private getOrCreateDetectorBackend<T>(
    activeBackend: string,
    context: DetectorBackendContext
  ): Promise<BaseDetectorBackend<T>> {
    let detectorBackendPromise = this._detectorBackends.get(activeBackend) as
      | Promise<BaseDetectorBackend<T>>
      | undefined;

    if (!detectorBackendPromise) {
      detectorBackendPromise = (async () => {
        switch (activeBackend) {
          case 'gemini':
            return new GeminiDetectorBackend(
              context
            ) as unknown as BaseDetectorBackend<T>;
          case 'mediapipe':
            return new MediaPipeDetectorBackend(
              context
            ) as unknown as BaseDetectorBackend<T>;
          default:
            throw new Error(
              `ObjectDetector backend '${activeBackend}' is not supported.`
            );
        }
      })();
      this._detectorBackends.set(
        activeBackend,
        detectorBackendPromise as Promise<BaseDetectorBackend<unknown>>
      );
    }
    return detectorBackendPromise;
  }

  private getDepthMeshSnapshot() {
    const depthMesh = this.depth.depthMesh!;
    const geometry = this.depth.options.depthMesh.updateFullResolutionGeometry
      ? depthMesh.geometry
      : depthMesh.downsampledGeometry || depthMesh.geometry;
    const clonedGeometry = geometry.clone();
    clonedGeometry.computeBoundingSphere();
    clonedGeometry.computeBoundingBox();
    const depthMeshSnapshot = new THREE.Mesh(
      clonedGeometry,
      new THREE.MeshBasicMaterial()
    );
    depthMesh.getWorldPosition(depthMeshSnapshot.position);
    depthMesh.getWorldQuaternion(depthMeshSnapshot.quaternion);
    depthMesh.getWorldScale(depthMeshSnapshot.scale);
    depthMeshSnapshot.updateMatrixWorld(true);
    return depthMeshSnapshot;
  }

  /**
   * Retrieves a list of currently detected objects.
   *
   * @param label - The semantic label to filter by (e.g., 'chair'). If null,
   * all objects are returned.
   * @returns An array of `Object` instances.
   */
  get<T = null>(label = null): DetectedObject<T>[] {
    const allObjects = Array.from(this._detectedObjects.values());
    if (!label) {
      return allObjects as DetectedObject<T>[];
    }
    return allObjects.filter(
      (obj) => obj.label === label
    ) as DetectedObject<T>[];
  }

  /**
   * Removes all currently detected objects from the scene and internal
   * tracking.
   */
  clear() {
    this.clearDetectedObjects();
    this.detectedObjects = [];
    return this;
  }

  private clearDetectedObjects() {
    for (const obj of this._detectedObjects.values()) {
      disposeObjectTree(obj);
      this.remove(obj);
    }
    this._detectedObjects.clear();
    if (this._debugVisualsGroup) {
      disposeObjectChildren(this._debugVisualsGroup);
    }
  }

  private disposeDepthMeshSnapshot(depthMeshSnapshot: THREE.Mesh) {
    depthMeshSnapshot.geometry.dispose();
    disposeMaterial(depthMeshSnapshot.material);
  }

  /**
   * Toggles the visibility of all debug visualizations for detected objects.
   * @param visible - Whether the visualizations should be visible.
   */
  showDebugVisualizations(visible = true) {
    if (this._debugVisualsGroup) {
      this._debugVisualsGroup.visible = visible;
    }
  }

  override dispose() {
    this.initialized = false;
    this.disposed = true;
    this.activeClients.clear();
    disposeObjectChildren(this);
    this.clear(); // Unlinks children so needs to come last

    for (const backendPromise of this._detectorBackends.values()) {
      void backendPromise
        .then((backend) => backend.dispose?.())
        .catch(() => {});
    }
    this._detectorBackends.clear();
    this.simulatorSource = undefined;
  }
}