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(() => {
    if (!source) {
      setAudioSource(null);
      return;
    }

    if (source instanceof MediaStream) {
      const ctx = new AudioContext();
      ctx.resume();
      const sourceNode = ctx.createMediaStreamSource(source);
      setAudioSource({ sourceNode, ctx });

      return () => {
        ctx.close();
      };
    }

    // HTMLAudioElement — createMediaElementSource can only be called once per element
    const cached = elementSourceCache.current.get(source);
    if (cached) {
      cached.ctx.resume();
      setAudioSource(cached);
      return;
    }

    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 });
    setAudioSource({ sourceNode, ctx });

    return () => {
      ctx.close();
      elementSourceCache.current.delete(source);
    };
  }, [source]);

  return audioSource;
}
