Luna-Powered Research Widget

// hooks/useLunaChat.ts
import { useState, useCallback, useRef } from "react";

interface Message {
  role: "user" | "assistant";
  content: string;
  sources?: { id: string; title: string; relevance: number }[];
}

const API_KEY = process.env.NEXT_PUBLIC_MOONLIT_API_KEY!;
const BASE = "https://api.moonlit.ai/v1.1";

export function useLunaChat(jurisdictions: string[] = ["European Union", "Netherlands"]) {
  const [messages, setMessages] = useState<Message[]>([]);
  const [isStreaming, setIsStreaming] = useState(false);
  const chatIdRef = useRef<string | null>(null);

  const initSession = useCallback(async () => {
    const res = await fetch(`${BASE}/luna/initialize_chat`, {
      method: "POST",
      headers: {
        "Ocp-Apim-Subscription-Key": API_KEY,
        "Content-Type": "application/json",
      },
      body: JSON.stringify({ jurisdictions }),
    });
    const data = await res.json();
    chatIdRef.current = data.chat_id;
    return data.chat_id;
  }, [jurisdictions]);

  const ask = useCallback(async (input: string) => {
    if (!chatIdRef.current) await initSession();

    setMessages((prev) => [...prev, { role: "user", content: input }]);
    setIsStreaming(true);

    const res = await fetch(`${BASE}/luna/ask_question`, {
      method: "POST",
      headers: {
        "Ocp-Apim-Subscription-Key": API_KEY,
        "Content-Type": "application/json",
      },
      body: JSON.stringify({
        chat_id: chatIdRef.current,
        input,
        jurisdictions,
      }),
    });

    const reader = res.body!.getReader();
    const decoder = new TextDecoder();
    let answer = "";
    let sources: Message["sources"] = [];

    setMessages((prev) => [...prev, { role: "assistant", content: "" }]);

    while (true) {
      const { done, value } = await reader.read();
      if (done) break;

      for (const line of decoder.decode(value).split("\n")) {
        if (line.startsWith("data: ")) {
          const payload = line.slice(6);
          if (payload === "[DONE]") break;

          try {
            const parsed = JSON.parse(payload);
            if (parsed.sources) { sources = parsed.sources; continue; }
          } catch { /* plain text chunk */ }

          answer += payload;
          setMessages((prev) => {
            const updated = [...prev];
            updated[updated.length - 1] = { role: "assistant", content: answer, sources };
            return updated;
          });
        }
      }
    }
    setIsStreaming(false);
  }, [initSession, jurisdictions]);

  const reset = useCallback(() => {
    chatIdRef.current = null;
    setMessages([]);
  }, []);

  return { messages, ask, reset, isStreaming };
}

// components/LunaWidget.tsx
"use client";

import { useState, useRef, useEffect } from "react";
import { useLunaChat } from "@/hooks/useLunaChat";

export function LunaWidget() {
  const { messages, ask, reset, isStreaming } = useLunaChat(["European Union", "Netherlands"]);
  const [input, setInput] = useState("");
  const bottomRef = useRef<HTMLDivElement>(null);

  useEffect(() => {
    bottomRef.current?.scrollIntoView({ behavior: "smooth" });
  }, [messages]);

  const handleSubmit = (e: React.FormEvent) => {
    e.preventDefault();
    if (!input.trim() || isStreaming) return;
    ask(input.trim());
    setInput("");
  };

  return (
    <div className="w-full max-w-2xl border rounded-lg shadow-lg bg-white">
      <div className="p-4 border-b flex justify-between items-center">
        <h3 className="font-semibold">Luna Legal Research</h3>
        <button onClick={reset} className="text-sm text-gray-500 hover:text-gray-800">
          New session
        </button>
      </div>

      <div className="h-96 overflow-y-auto p-4 space-y-4">
        {messages.map((msg, i) => (
          <div key={i} className={`flex ${msg.role === "user" ? "justify-end" : "justify-start"}`}>
            <div className={`max-w-[80%] p-3 rounded-lg ${
              msg.role === "user" ? "bg-blue-600 text-white" : "bg-gray-100"
            }`}>
              <p className="whitespace-pre-wrap text-sm">{msg.content}</p>
              {msg.sources && msg.sources.length > 0 && (
                <div className="mt-2 pt-2 border-t border-gray-200">
                  <p className="text-xs font-semibold mb-1">Sources:</p>
                  {msg.sources.map((s) => (
                    <p key={s.id} className="text-xs text-blue-700">{s.id}</p>
                  ))}
                </div>
              )}
            </div>
          </div>
        ))}
        <div ref={bottomRef} />
      </div>

      <form onSubmit={handleSubmit} className="p-4 border-t flex gap-2">
        <input
          value={input}
          onChange={(e) => setInput(e.target.value)}
          placeholder="Ask a legal question..."
          className="flex-1 border rounded px-3 py-2 text-sm"
          disabled={isStreaming}
        />
        <button type="submit" disabled={isStreaming}
          className="px-4 py-2 bg-blue-600 text-white rounded text-sm disabled:opacity-50">
          Ask
        </button>
      </form>
    </div>
  );
}

// Enhanced ask function with session recovery
const askWithRetry = useCallback(async (input: string) => {
  try {
    await ask(input);
  } catch (error: any) {
    if (error?.status === 404 || error?.message?.includes("session")) {
      console.log("Session expired, reinitializing...");
      chatIdRef.current = null;
      await initSession();
      await ask(input);  // Retry with new session
    } else {
      throw error;
    }
  }
}, [ask, initSession]);