import { ref, computed, type Ref, type ComputedRef } from "vue";
import { useAgentChat, type AgentMessage, type UseAgentChatOptions } from "./useAgentChat";
import { useAgentRateLimit } from "./useAgentRateLimit";

export interface AgentChatInstance {
  messages: Ref<AgentMessage[]>;
  isLoading: Ref<boolean>;
  hasPendingApproval: ComputedRef<boolean>;
  errorMessage: Ref<string>;
  isStopped: Ref<boolean>;
  rateLimit: ReturnType<typeof useAgentRateLimit>;
  send: ReturnType<typeof useAgentChat>["send"];
  stopGeneration: ReturnType<typeof useAgentChat>["stopGeneration"];
  clear: ReturnType<typeof useAgentChat>["clear"];
  loadMessages: ReturnType<typeof useAgentChat>["loadMessages"];
  respondToApproval: ReturnType<typeof useAgentChat>["respondToApproval"];
  sendFeedback: ReturnType<typeof useAgentChat>["sendFeedback"];
}

export interface UseAgentChatStoreOptions {
  modelId: () => number | null;
  enableThinking: () => boolean;
  enableTools: () => boolean;
  onStreamComplete?: (sessionId: string) => void;
}

export function useAgentChatStore(storeOptions: UseAgentChatStoreOptions) {
  const instances = new Map<string, AgentChatInstance>();
  const currentSessionId = ref<string | null>(null);

  const loadedSessionIds = ref<Set<string>>(new Set());

  function createInstance(sid: string): AgentChatInstance {
    const options: UseAgentChatOptions = {
      sessionId: () => sid,
      modelId: storeOptions.modelId,
      enableThinking: storeOptions.enableThinking,
      enableTools: storeOptions.enableTools,
      onStreamComplete: () => {
        storeOptions.onStreamComplete?.(sid);
      },
    };
    const chat = useAgentChat(options);
    return chat;
  }

  function getInstance(sid: string): AgentChatInstance {
    let inst = instances.get(sid);
    if (!inst) {
      inst = createInstance(sid);
      instances.set(sid, inst);
    }
    return inst;
  }

  function removeInstance(sid: string) {
    const inst = instances.get(sid);
    if (inst) {
      inst.stopGeneration();
      inst.clear();
      instances.delete(sid);
    }
    loadedSessionIds.value.delete(sid);
  }

  const currentInstance = computed<AgentChatInstance | null>(() => {
    const sid = currentSessionId.value;
    if (!sid) return null;
    return getInstance(sid);
  });

  const messages = computed(() => currentInstance.value?.messages.value ?? []);
  const isLoading = computed(() => currentInstance.value?.isLoading.value ?? false);
  const hasPendingApproval = computed(() => currentInstance.value?.hasPendingApproval.value ?? false);
  const errorMessage = computed(() => currentInstance.value?.errorMessage.value ?? "");
  const isStopped = computed(() => currentInstance.value?.isStopped.value ?? false);
  const fallbackRateLimit = useAgentRateLimit();
  const rateLimit = computed(() => currentInstance.value?.rateLimit ?? fallbackRateLimit);

  const anyLoading = computed(() => {
    for (const inst of instances.values()) {
      if (inst.isLoading.value) return true;
    }
    return false;
  });

  function getLoadingSessionIds(): string[] {
    const ids: string[] = [];
    for (const [sid, inst] of instances) {
      if (inst.isLoading.value) ids.push(sid);
    }
    return ids;
  }

  async function loadMessages(sid: string) {
    const inst = getInstance(sid);
    loadedSessionIds.value.add(sid);
    await inst.loadMessages(sid);
  }

  async function send(content: string, opts?: { editMessageId?: string; regenerate?: boolean }) {
    const sid = currentSessionId.value;
    if (!sid) return;
    const inst = getInstance(sid);
    await inst.send(content, opts);
  }

  async function respondToApproval(toolCallId: string, approved: boolean, reason?: string) {
    const sid = currentSessionId.value;
    if (!sid) return;
    const inst = getInstance(sid);
    await inst.respondToApproval(toolCallId, approved, reason);
  }

  async function sendFeedback(messageId: string, feedback: "like" | "dislike") {
    const sid = currentSessionId.value;
    if (!sid) return;
    const inst = getInstance(sid);
    await inst.sendFeedback(messageId, feedback);
  }

  function stopGeneration() {
    const sid = currentSessionId.value;
    if (!sid) return;
    const inst = getInstance(sid);
    inst.stopGeneration();
  }

  function stopAll() {
    for (const inst of instances.values()) {
      inst.stopGeneration();
    }
  }

  function clear() {
    const sid = currentSessionId.value;
    if (!sid) return;
    const inst = getInstance(sid);
    inst.clear();
  }

  function setCurrentSessionId(sid: string | null) {
    currentSessionId.value = sid;
  }

  function isSessionLoaded(sid: string): boolean {
    return loadedSessionIds.value.has(sid);
  }

  return {
    currentSessionId,
    currentInstance,
    messages,
    isLoading,
    hasPendingApproval,
    errorMessage,
    isStopped,
    rateLimit,
    anyLoading,
    getLoadingSessionIds,
    loadMessages,
    send,
    respondToApproval,
    sendFeedback,
    stopGeneration,
    stopAll,
    clear,
    setCurrentSessionId,
    isSessionLoaded,
    removeInstance,
    getInstance,
  };
}