"use client";
import React, {
  createContext,
  useContext,
  useState,
  useEffect,
  useRef,
  ReactNode,
} from "react";
import removeMarkdown from "remove-markdown";
import { toast } from "sonner";

// More robust sentence splitter using Intl.Segmenter for better accuracy.
function splitIntoSentences(text: string): string[] {
  if (typeof Intl !== "undefined" && Intl.Segmenter) {
    const segmenter = new Intl.Segmenter("en", { granularity: "sentence" });
    return Array.from(segmenter.segment(text)).map(
      (segment) => segment.segment
    );
  }
  // Fallback to regex-based splitting if Intl.Segmenter is unavailable.
  return text.match(/[^\.!\?]+[\.!\?]+/g) || [text];
}

interface TTSContextType {
  sentences: string[];
  currentSentence: number;
  ttsBuffer: (string | null)[];
  voices: any[];
  selectedSpeaker: string;
  status: "ready" | "running" | null;
  setSelectedSpeaker: (speaker: string) => void;
  setCurrentSentence: (index: number) => void;
  playSentence: (index: number) => void;
  skipToSentence: (index: number) => void;
  playInOrder: (index: number) => void;
  pause: () => void;
  resume: () => void;
  stop: () => void;
}

const TTSContext = createContext<TTSContextType | undefined>(undefined);

export const TTSProvider = ({
  pages,
  children,
}: {
  pages: string[];
  children: ReactNode;
}) => {
  // Combine pages and split into sentences.
  const fullText = pages.join("\n");
  const sentences = splitIntoSentences(fullText).filter(
    (sentence) => sentence.trim() !== "\\n" && sentence.trim() !== ""
  );

  const [currentSentence, setCurrentSentence] = useState(0);
  const [ttsBuffer, setTtsBuffer] = useState<(string | null)[]>(
    Array(sentences.length).fill(null)
  );
  const audioRef = useRef<HTMLAudioElement>(null);

  // Create a reference to the worker object.
  const worker = useRef<Worker>(null);

  const [selectedSpeaker, setSelectedSpeaker] = useState("af_heart");

  const [playing, setPlaying] = useState(false);
  const [sentence, setSentence] = useState<number>();
  const [voices, setVoices] = useState<any[]>([]);
  const [status, setStatus] = useState<"ready" | "running" | null>(null);
  const [error, setError] = useState(null);
  const [loadingMessage, setLoadingMessage] = useState("Loading...");

  const [results, setResults] = useState<{ text: string; src: string }[]>([]);

  async function generateTTSForIndex(
    sentence: string,
    index: number
  ): Promise<string> {
    const key = `tts-${index}`;
    const cached = localStorage.getItem(key);
    if (cached) {
      return cached;
    }

    return new Promise((resolve, reject) => {
      const handleMessage = (e: MessageEvent) => {
        if (e.data.index !== index) return; // Ignore messages for other indices

        if (e.data.status === "complete") {
          localStorage.setItem(key, e.data.audio);
          worker.current!.removeEventListener("message", handleMessage); // Clean up listener
          resolve(e.data.audio);
        } else if (e.data.status === "error") {
          worker.current!.removeEventListener("message", handleMessage); // Clean up listener
          toast.error(`Error generating audio: ${e.data.error}`);
          reject(e.data.error);
        }
      };

      worker.current!.addEventListener("message", handleMessage);

      worker.current!.postMessage({
        type: "generate",
        index,
        text: sentence,
        voice: selectedSpeaker,
      });
    });
  }

  // We use the `useEffect` hook to setup the worker as soon as the `App` component is mounted.
  useEffect(() => {
    console.log("Initializing worker...");
    worker.current ??= new Worker("/workers/kokoro-worker.js", {
      type: "module",
    });

    console.log("Worker initialized");

    const onMessageReceived = (e: any) => {
      switch (e.data.status) {
        case "device":
          setLoadingMessage(`Loading model (device="${e.data.device}")`);
          break;
        case "ready":
          setStatus("ready");
          setVoices(e.data.voices);
          break;
        case "error":
          setError(e.data.data);
          break;
        case "complete":
          const { audio, text } = e.data;
          setResults((prev) => [{ text, src: audio }, ...prev]);
          setStatus("ready");
          break;
      }
    };

    const onErrorReceived = (e: any) => {
      console.error("Worker error:", e);
      setError(e.message);
    };

    worker.current.addEventListener("message", onMessageReceived);
    worker.current.addEventListener("error", onErrorReceived);

    return () => {
      worker.current!.removeEventListener("message", onMessageReceived);
      worker.current!.removeEventListener("error", onErrorReceived);
    };
  }, []);

  // Pre-buffer current and next 5 sentences.
  useEffect(() => {
    let isCancelled = false;

    async function preloadBuffer() {
      const newBuffer = [...ttsBuffer];
      const end = Math.min(sentences.length, currentSentence + 5); // Preload 5 sentences ahead

      for (let i = currentSentence; i < end; i++) {
        if (isCancelled) break;
        if (!newBuffer[i]) {
          console.log("Preloading TTS for sentence:", i, sentences[i]);
          try {
            newBuffer[i] = await generateTTSForIndex(
              removeMarkdown(sentences[i]),
              i
            );
          } catch (error) {
            console.error("Error preloading TTS:", error);
          }
        }
      }

      if (!isCancelled) {
        setTtsBuffer((prev) => {
          // Only update state if the buffer has changed
          if (JSON.stringify(prev) !== JSON.stringify(newBuffer)) {
            return newBuffer;
          }
          return prev;
        });
      }
    }

    preloadBuffer();

    return () => {
      isCancelled = true; // Cancel preloading if the component unmounts or dependencies change
    };
  }, [currentSentence, sentences]);

  const playSentence = async (index: number) => {
    if (index === currentSentence) return; // Prevent redundant updates
    setCurrentSentence(index);

    let audioUrl = ttsBuffer[index];
    if (!audioUrl) {
      audioUrl = await generateTTSForIndex(
        removeMarkdown(sentences[index]),
        index
      );
      setTtsBuffer((prev) => {
        const updated = [...prev];
        updated[index] = audioUrl;
        return updated;
      });
    }

    if (audioRef.current) {
      audioRef.current.src = audioUrl;
      await new Promise((res) => {
        audioRef.current!.play();

        audioRef.current!.onended = () => res(true);
      });
    }
  };

  const skipToSentence = (index: number) => {
    if (index < 0 || index >= sentences.length) return;
    playSentence(index);
  };

  const playInOrder = async (index: number) => {
    if (index < 0 || index >= sentences.length) return;
    if (index === currentSentence && playing) return; // Prevent redundant playback
    setCurrentSentence(index);
    setPlaying(true);

    for (let i = index; i < sentences.length; i++) {
      console.log("Playing sentence:", i, sentences[i]);
      try {
        await playSentence(i);
      } catch (error) {
        console.error("Error playing sentence:", error);
        break; // Stop playback on error
      }
    }

    setPlaying(false);
  };

  const pause = () => {
    if (audioRef.current) {
      audioRef.current.pause();
    }
  };

  const resume = () => {
    if (audioRef.current) {
      audioRef.current.play();
    }
  };

  const stop = () => {
    if (audioRef.current) {
      audioRef.current.pause();
      audioRef.current.currentTime = 0;
    }
  };

  const value: TTSContextType = {
    sentences,
    currentSentence,
    ttsBuffer,
    voices,
    playSentence,
    skipToSentence,
    selectedSpeaker,
    setSelectedSpeaker,
    setCurrentSentence,
    playInOrder,
    pause,
    resume,
    stop,
    status,
  };

  return (
    <TTSContext.Provider value={value}>
      {children}
      {/* Hidden audio element used for playback */}
      <audio ref={audioRef} style={{ display: "none" }} />
    </TTSContext.Provider>
  );
};

export const useTTS = (): TTSContextType => {
  const context = useContext(TTSContext);
  if (!context) {
    throw new Error("useTTS must be used within a TTSProvider");
  }
  return context;
};