import { useEffect, useRef, useState } from 'react';

interface AudioSource {
  sourceNode: AudioNode;
  ctx: AudioContext;
}

/**
 * Creates an AudioContext and source node from either a MediaStream (live recording)
 * or an HTMLAudioElement (review playback).
 *
 * Important: `createMediaElementSource` can only be called once per element,
 * so we cache the source per element instance.
 */
export function useAudioSource(
  source: MediaStream | HTMLAudioElement | null,
): AudioSource | null {
  const [audioSource, setAudioSource] = useState<AudioSource | null>(null);
  const elementSourceCache = useRef<
    WeakMap<
      HTMLAudioElement,
      { sourceNode: MediaElementAudioSourceNode; ctx: AudioContext }
    >
  >(new WeakMap());

  useEffect(() => {
    // Build the source node (creating an AudioContext as needed) plus optional
    // teardown for contexts we own; the branches converge on one setState.
    let result: AudioSource | null = null;
    let cleanup: (() => void) | undefined;

    if (source instanceof MediaStream) {
      const ctx = new AudioContext();
      ctx.resume();
      const sourceNode = ctx.createMediaStreamSource(source);
      result = { sourceNode, ctx };
      cleanup = () => ctx.close();
    } else if (source) {
      // HTMLAudioElement — createMediaElementSource can only be called once per
      // element, so reuse a cached context/node when we have one.
      const cached = elementSourceCache.current.get(source);
      if (cached) {
        cached.ctx.resume();
        result = cached;
      } else {
        const ctx = new AudioContext();
        ctx.resume();
        const sourceNode = ctx.createMediaElementSource(source);
        // Connect element source to destination so audio is still audible
        sourceNode.connect(ctx.destination);
        elementSourceCache.current.set(source, { sourceNode, ctx });
        result = { sourceNode, ctx };
        cleanup = () => {
          ctx.close();
          elementSourceCache.current.delete(source);
        };
      }
    }

    // Publishing an imperatively-created Web Audio node — external-resource
    // sync, not a re-render cascade.
    // eslint-disable-next-line react-hooks/set-state-in-effect
    setAudioSource(result);

    return cleanup;
  }, [source]);

  return audioSource;
}