import { useCallback, useEffect, useRef, useState } from "react";
import { ChatMessage, Conversation } from "./types";

const STORAGE_KEY = "litellm_chat_history_v1";
const MAX_CONVERSATIONS = 100;
const TITLE_MAX_LENGTH = 40;

function generateTitle(firstUserMessage: string): string {
  const trimmed = firstUserMessage.trim();
  if (trimmed.length <= TITLE_MAX_LENGTH) {
    return trimmed;
  }
  return trimmed.slice(0, TITLE_MAX_LENGTH) + "…";
}

function loadFromStorage(): { conversations: Conversation[]; storageUnavailable: boolean } {
  try {
    const raw = localStorage.getItem(STORAGE_KEY);
    if (!raw) {
      return { conversations: [], storageUnavailable: false };
    }
    const parsed = JSON.parse(raw) as Conversation[];
    return { conversations: parsed, storageUnavailable: false };
  } catch {
    return { conversations: [], storageUnavailable: true };
  }
}

function saveToStorage(conversations: Conversation[]): boolean {
  try {
    localStorage.setItem(STORAGE_KEY, JSON.stringify(conversations));
    return true;
  } catch {
    return false;
  }
}

function trimConversations(conversations: Conversation[]): Conversation[] {
  if (conversations.length <= MAX_CONVERSATIONS) {
    return conversations;
  }
  return [...conversations].sort((a, b) => b.updatedAt - a.updatedAt).slice(0, MAX_CONVERSATIONS);
}

export function useChatHistory(activeConversationId: string | null): {
  conversations: Conversation[];
  activeConversation: Conversation | null;
  storageUnavailable: boolean;
  staleId: boolean;
  createConversation: (model: string) => string;
  appendMessage: (conversationId: string, message: Omit<ChatMessage, "id" | "timestamp">) => void;
  updateLastAssistantMessage: (
    conversationId: string,
    updates: Partial<Pick<ChatMessage, "content" | "reasoningContent" | "mcpEvents">>,
  ) => void;
  /** Remove the message with `messageId` and all subsequent messages from the conversation. */
  truncateFromMessage: (conversationId: string, messageId: string) => void;
  deleteConversation: (id: string) => void;
  renameConversation: (id: string, newTitle: string) => void;
  setActiveConversationId: (id: string | null) => void;
} {
  const [conversations, setConversations] = useState<Conversation[]>([]);
  const [storageUnavailable, setStorageUnavailable] = useState(false);
  const [staleId, setStaleId] = useState(false);
  const [currentActiveId, setCurrentActiveId] = useState<string | null>(activeConversationId);
  // Ref so updater functions stay pure (no state setter calls inside setConversations)
  const storageUnavailableRef = useRef(false);
  const initializedRef = useRef(false);

  // Sync internal active id whenever the URL-derived prop changes (e.g. "New chat" → null)
  useEffect(() => {
    setCurrentActiveId(activeConversationId);
    setStaleId(false);
  }, [activeConversationId]);

  useEffect(() => {
    const { conversations: loaded, storageUnavailable: unavailable } = loadFromStorage();
    storageUnavailableRef.current = unavailable;
    setConversations(loaded);
    setStorageUnavailable(unavailable);
    initializedRef.current = true;

    if (activeConversationId !== null) {
      const found = loaded.some((c) => c.id === activeConversationId);
      if (!found) {
        setStaleId(true);
      }
    }
  }, []); // eslint-disable-line react-hooks/exhaustive-deps

  // Persist to localStorage after every conversations change (pure effect, no setState inside updaters)
  useEffect(() => {
    if (!initializedRef.current) return;
    if (storageUnavailableRef.current) return;
    const success = saveToStorage(conversations);
    if (!success) {
      storageUnavailableRef.current = true;
      setStorageUnavailable(true);
    }
  }, [conversations]);

  const createConversation = useCallback((model: string): string => {
    const id = crypto.randomUUID();
    const now = Date.now();
    const newConversation: Conversation = {
      id,
      title: "New conversation",
      model,
      messages: [],
      mcpServerNames: [],
      createdAt: now,
      updatedAt: now,
    };
    setConversations((prev) => trimConversations([newConversation, ...prev]));
    setCurrentActiveId(id);
    return id;
  }, []);

  const appendMessage = useCallback((conversationId: string, message: Omit<ChatMessage, "id" | "timestamp">) => {
    const newMessage: ChatMessage = {
      ...message,
      id: crypto.randomUUID(),
      timestamp: Date.now(),
    };

    setConversations((prev) => {
      const updated = prev.map((conv) => {
        if (conv.id !== conversationId) return conv;
        const updatedMessages = [...conv.messages, newMessage];
        let title = conv.title;
        if (
          title === "New conversation" &&
          newMessage.role === "user" &&
          conv.messages.filter((m) => m.role === "user").length === 0
        ) {
          title = generateTitle(newMessage.content);
        }
        return { ...conv, title, messages: updatedMessages, updatedAt: Date.now() };
      });
      return trimConversations(updated);
    });
  }, []);

  const updateLastAssistantMessage = useCallback(
    (conversationId: string, updates: Partial<Pick<ChatMessage, "content" | "reasoningContent" | "mcpEvents">>) => {
      setConversations((prev) => {
        const updated = prev.map((conv) => {
          if (conv.id !== conversationId) return conv;
          const messages = [...conv.messages];
          const lastAssistantIndex = messages.reduceRight((found, msg, idx) => {
            if (found !== -1) return found;
            return msg.role === "assistant" ? idx : -1;
          }, -1);
          if (lastAssistantIndex === -1) return conv;
          messages[lastAssistantIndex] = { ...messages[lastAssistantIndex], ...updates };
          return { ...conv, messages, updatedAt: Date.now() };
        });
        return trimConversations(updated);
      });
    },
    [],
  );

  const truncateFromMessage = useCallback((conversationId: string, messageId: string) => {
    setConversations((prev) => {
      const updated = prev.map((conv) => {
        if (conv.id !== conversationId) return conv;
        const idx = conv.messages.findIndex((m) => m.id === messageId);
        if (idx === -1) return conv;
        return { ...conv, messages: conv.messages.slice(0, idx), updatedAt: Date.now() };
      });
      return trimConversations(updated);
    });
  }, []);

  const deleteConversation = useCallback(
    (id: string) => {
      setConversations((prev) => trimConversations(prev.filter((c) => c.id !== id)));
      if (currentActiveId === id) setCurrentActiveId(null);
    },
    [currentActiveId],
  );

  const renameConversation = useCallback((id: string, newTitle: string) => {
    setConversations((prev) =>
      trimConversations(
        prev.map((conv) => (conv.id === id ? { ...conv, title: newTitle, updatedAt: Date.now() } : conv)),
      ),
    );
  }, []);

  const setActiveConversationId = useCallback((id: string | null) => {
    setCurrentActiveId(id);
    setStaleId(false);
  }, []);

  const activeConversation =
    currentActiveId !== null ? conversations.find((c) => c.id === currentActiveId) ?? null : null;

  return {
    conversations,
    activeConversation,
    storageUnavailable,
    staleId,
    createConversation,
    appendMessage,
    updateLastAssistantMessage,
    truncateFromMessage,
    deleteConversation,
    renameConversation,
    setActiveConversationId,
  };
}
