frontend/src/store/agentStore.ts

/**
 * Agent store — manages UI state that is NOT handled by the Vercel AI SDK.
 *
 * Message state (messages, streaming, tool calls) is now managed by useChat().
 * This store only handles:
 *  - Connection / processing flags
 *  - Panel state (right panel tabs)
 *  - Plan state
 *  - User info / error banners
 *  - Edited scripts (for hf_jobs code editing)
 */
import { create } from 'zustand';
import type { User } from '@/types/agent';

export interface PlanItem {
  id: string;
  content: string;
  status: 'pending' | 'in_progress' | 'completed';
}

interface PanelTab {
  id: string;
  title: string;
  content: string;
  language?: string;
  parameters?: Record<string, unknown>;
}

export interface LLMHealthError {
  error: string;
  errorType: 'auth' | 'credits' | 'rate_limit' | 'network' | 'unknown';
  model: string;
}

export type ActivityStatus =
  | { type: 'idle' }
  | { type: 'thinking' }
  | { type: 'tool'; toolName: string }
  | { type: 'waiting-approval' }
  | { type: 'streaming' };

interface AgentStore {
  // Global UI flags
  isProcessing: boolean;
  isConnected: boolean;
  activityStatus: ActivityStatus;
  user: User | null;
  error: string | null;
  llmHealthError: LLMHealthError | null;

  // Right panel
  panelContent: { title: string; content: string; language?: string; parameters?: Record<string, unknown> } | null;
  panelTabs: PanelTab[];
  activePanelTab: string | null;

  // Plan
  plan: PlanItem[];

  // Edited scripts (tool_call_id -> edited content)
  editedScripts: Record<string, string>;

  // Job URLs (tool_call_id -> job URL) for HF jobs
  jobUrls: Record<string, string>;

  // Actions
  setProcessing: (isProcessing: boolean) => void;
  setConnected: (isConnected: boolean) => void;
  setActivityStatus: (status: ActivityStatus) => void;
  setUser: (user: User | null) => void;
  setError: (error: string | null) => void;
  setLlmHealthError: (error: LLMHealthError | null) => void;

  setPanelContent: (content: { title: string; content: string; language?: string; parameters?: Record<string, unknown> } | null) => void;
  setPanelTab: (tab: PanelTab) => void;
  updatePanelTabContent: (tabId: string, content: string) => void;
  setActivePanelTab: (tabId: string) => void;
  clearPanelTabs: () => void;
  removePanelTab: (tabId: string) => void;
  showToolOutput: (output: { tool: string; output?: string; args?: Record<string, unknown> }) => void;

  setPlan: (plan: PlanItem[]) => void;

  setEditedScript: (toolCallId: string, content: string) => void;
  getEditedScript: (toolCallId: string) => string | undefined;
  clearEditedScripts: () => void;

  setJobUrl: (toolCallId: string, jobUrl: string) => void;
  getJobUrl: (toolCallId: string) => string | undefined;
}

export const useAgentStore = create<AgentStore>()((set, get) => ({
  isProcessing: false,
  isConnected: false,
  activityStatus: { type: 'idle' },
  user: null,
  error: null,
  llmHealthError: null,

  panelContent: null,
  panelTabs: [],
  activePanelTab: null,

  plan: [],

  editedScripts: {},
  jobUrls: {},

  // ── Global flags ──────────────────────────────────────────────────

  setProcessing: (isProcessing) => set({ isProcessing, ...(!isProcessing ? { activityStatus: { type: 'idle' } } : {}) }),
  setConnected: (isConnected) => set({ isConnected }),
  setActivityStatus: (status) => set({ activityStatus: status }),
  setUser: (user) => set({ user }),
  setError: (error) => set({ error }),
  setLlmHealthError: (error) => set({ llmHealthError: error }),

  // ── Panel ─────────────────────────────────────────────────────────

  setPanelContent: (content) => set({ panelContent: content }),

  setPanelTab: (tab: PanelTab) => {
    set((state) => {
      const idx = state.panelTabs.findIndex(t => t.id === tab.id);
      let newTabs: PanelTab[];
      if (idx >= 0) {
        newTabs = [...state.panelTabs];
        newTabs[idx] = tab;
      } else {
        newTabs = [...state.panelTabs, tab];
      }
      return {
        panelTabs: newTabs,
        activePanelTab: state.activePanelTab || tab.id,
      };
    });
  },

  updatePanelTabContent: (tabId, content) => {
    set((state) => ({
      panelTabs: state.panelTabs.map(tab =>
        tab.id === tabId ? { ...tab, content } : tab
      ),
    }));
  },

  setActivePanelTab: (tabId) => set({ activePanelTab: tabId }),

  clearPanelTabs: () => set({ panelTabs: [], activePanelTab: null }),

  removePanelTab: (tabId) => {
    set((state) => {
      const newTabs = state.panelTabs.filter(t => t.id !== tabId);
      let newActiveTab = state.activePanelTab;
      if (state.activePanelTab === tabId) {
        newActiveTab = newTabs.length > 0 ? newTabs[newTabs.length - 1].id : null;
      }
      return { panelTabs: newTabs, activePanelTab: newActiveTab };
    });
  },

  showToolOutput: (info) => {
    const content = info.output || (info.args ? JSON.stringify(info.args, null, 2) : 'No output available');
    let language = 'text';
    if (content.trim().startsWith('{') || content.trim().startsWith('[')) language = 'json';
    else if (content.includes('```')) language = 'markdown';

    const state = get();
    const otherTabs = state.panelTabs.filter(t => t.id !== 'tool_output');
    set({
      panelTabs: [...otherTabs, { id: 'tool_output', title: info.tool, content, language }],
      activePanelTab: 'tool_output',
    });
  },

  // ── Plan ──────────────────────────────────────────────────────────

  setPlan: (plan) => set({ plan }),

  // ── Edited scripts ────────────────────────────────────────────────

  setEditedScript: (toolCallId, content) => {
    set((state) => ({
      editedScripts: { ...state.editedScripts, [toolCallId]: content },
    }));
  },

  getEditedScript: (toolCallId) => get().editedScripts[toolCallId],

  clearEditedScripts: () => set({ editedScripts: {} }),

  // ── Job URLs ────────────────────────────────────────────────────────

  setJobUrl: (toolCallId, jobUrl) => {
    set((state) => ({
      jobUrls: { ...state.jobUrls, [toolCallId]: jobUrl },
    }));
  },

  getJobUrl: (toolCallId) => get().jobUrls[toolCallId],
}));