import { create } from "zustand";
import type { ChatMessage, ConnectionConfig, ContextUsage, DbFlavor, Tab } from "@/types";

interface AppState {
  connections: ConnectionConfig[];
  activeConnectionId: string | null;
  currentDatabase: string | null;
  connectedIds: Set<string>;
  readOnlyMap: Record<string, boolean>;
  dbFlavors: Record<string, DbFlavor>;
  tabs: Tab[];
  activeTabId: string | null;
  sidebarWidth: number;
  pgVersion: string | null;
  chatThreads: Record<string, ChatMessage[]>;
  chatPending: Record<string, boolean>;
  chatUsage: Record<string, ContextUsage>;

  setConnections: (connections: ConnectionConfig[]) => void;
  setActiveConnectionId: (id: string | null) => void;
  setCurrentDatabase: (db: string | null) => void;
  addConnectedId: (id: string) => void;
  removeConnectedId: (id: string) => void;
  setReadOnly: (connectionId: string, readOnly: boolean) => void;
  setDbFlavor: (connectionId: string, flavor: DbFlavor) => void;
  setPgVersion: (version: string | null) => void;

  addTab: (tab: Tab) => void;
  closeTab: (id: string) => void;
  setActiveTabId: (id: string | null) => void;
  updateTab: (id: string, updates: Partial<Tab>) => void;
  setSidebarWidth: (width: number) => void;

  appendChatMessages: (tabId: string, messages: ChatMessage[]) => void;
  replaceChatThread: (tabId: string, messages: ChatMessage[]) => void;
  clearChatThread: (tabId: string) => void;
  setChatPending: (tabId: string, pending: boolean) => void;
  setChatUsage: (tabId: string, usage: ContextUsage) => void;
}

export const useAppStore = create<AppState>((set) => ({
  connections: [],
  activeConnectionId: null,
  currentDatabase: null,
  connectedIds: new Set(),
  readOnlyMap: {},
  dbFlavors: {},
  tabs: [],
  activeTabId: null,
  sidebarWidth: 260,
  pgVersion: null,
  chatThreads: {},
  chatPending: {},
  chatUsage: {},

  setConnections: (connections) => set({ connections }),
  setActiveConnectionId: (id) => set({ activeConnectionId: id }),
  setCurrentDatabase: (db) => set({ currentDatabase: db }),
  addConnectedId: (id) =>
    set((state) => ({
      connectedIds: new Set([...state.connectedIds, id]),
      readOnlyMap: { ...state.readOnlyMap, [id]: true },
    })),
  removeConnectedId: (id) =>
    set((state) => {
      const next = new Set(state.connectedIds);
      next.delete(id);
      const restRo = Object.fromEntries(
        Object.entries(state.readOnlyMap).filter(([k]) => k !== id)
      );
      const restFlavors = Object.fromEntries(
        Object.entries(state.dbFlavors).filter(([k]) => k !== id)
      );
      return { connectedIds: next, readOnlyMap: restRo, dbFlavors: restFlavors };
    }),
  setReadOnly: (connectionId, readOnly) =>
    set((state) => ({
      readOnlyMap: { ...state.readOnlyMap, [connectionId]: readOnly },
    })),
  setDbFlavor: (connectionId, flavor) =>
    set((state) => ({
      dbFlavors: { ...state.dbFlavors, [connectionId]: flavor },
    })),
  setPgVersion: (version) => set({ pgVersion: version }),

  addTab: (tab) =>
    set((state) => ({
      tabs: [...state.tabs, tab],
      activeTabId: tab.id,
    })),
  closeTab: (id) =>
    set((state) => {
      const tabs = state.tabs.filter((t) => t.id !== id);
      const activeTabId =
        state.activeTabId === id
          ? tabs.length > 0
            ? tabs[tabs.length - 1].id
            : null
          : state.activeTabId;
      const chatThreads = { ...state.chatThreads };
      delete chatThreads[id];
      const chatPending = { ...state.chatPending };
      delete chatPending[id];
      const chatUsage = { ...state.chatUsage };
      delete chatUsage[id];
      return { tabs, activeTabId, chatThreads, chatPending, chatUsage };
    }),
  setActiveTabId: (id) => set({ activeTabId: id }),
  updateTab: (id, updates) =>
    set((state) => ({
      tabs: state.tabs.map((t) => (t.id === id ? { ...t, ...updates } : t)),
    })),
  setSidebarWidth: (width) => set({ sidebarWidth: width }),

  appendChatMessages: (tabId, messages) =>
    set((state) => ({
      chatThreads: {
        ...state.chatThreads,
        [tabId]: [...(state.chatThreads[tabId] ?? []), ...messages],
      },
    })),
  replaceChatThread: (tabId, messages) =>
    set((state) => ({
      chatThreads: { ...state.chatThreads, [tabId]: messages },
    })),
  clearChatThread: (tabId) =>
    set((state) => {
      const chatThreads = { ...state.chatThreads, [tabId]: [] };
      const chatUsage = { ...state.chatUsage };
      delete chatUsage[tabId];
      return { chatThreads, chatUsage };
    }),
  setChatPending: (tabId, pending) =>
    set((state) => ({
      chatPending: { ...state.chatPending, [tabId]: pending },
    })),
  setChatUsage: (tabId, usage) =>
    set((state) => ({
      chatUsage: { ...state.chatUsage, [tabId]: usage },
    })),
}));
