'use client';

import { createContext, useContext, useEffect, useRef, useState } from 'react';
import { useRoomContext } from '@livekit/components-react';

export interface GeneratedImage {
  id: string;
  imageUrl: string;
  mimeType: string;
  prompt: string;
  timestamp: number;
}

const GeneratedImagesContext = createContext<GeneratedImage[]>([]);

/**
 * Registers the "generated-image" byte stream handler once for the whole tree.
 * Wrap this around any subtree that contains components using useGeneratedImages().
 */
export function GeneratedImagesProvider({ children }: { children: React.ReactNode }) {
  const room = useRoomContext();
  const [images, setImages] = useState<GeneratedImage[]>([]);
  const registeredRef = useRef(false);

  useEffect(() => {
    if (registeredRef.current) return;
    registeredRef.current = true;

    room.registerByteStreamHandler('generated-image', async (reader) => {
      try {
        const chunks = await reader.readAll();
        const mimeType = reader.info.mimeType || 'image/png';
        const blobParts = chunks.map((chunk) => {
          const copy = new Uint8Array(chunk.byteLength);
          copy.set(chunk);
          return copy.buffer;
        });
        const blob = new Blob(blobParts, { type: mimeType });
        const imageUrl = URL.createObjectURL(blob);

        const image: GeneratedImage = {
          id: `${Date.now()}-${Math.random().toString(36).slice(2)}`,
          imageUrl,
          mimeType,
          prompt: reader.info.attributes?.prompt ?? '',
          timestamp: Date.now(),
        };

        setImages((prev) => [...prev, image]);
      } catch (err) {
        console.error('Failed to receive generated-image byte stream:', err);
      }
    });

    return () => {
      room.unregisterByteStreamHandler('generated-image');
      registeredRef.current = false;
    };
  }, [room]);

  return (
    <GeneratedImagesContext.Provider value={images}>{children}</GeneratedImagesContext.Provider>
  );
}

/**
 * Returns all images generated by the agent. Must be used inside GeneratedImagesProvider.
 */
export function useGeneratedImages(): GeneratedImage[] {
  return useContext(GeneratedImagesContext);
}