veridian/app/composables/useAgents.ts

import { nanoid } from "nanoid";
import { attempt } from "~~/types/result";
import * as schema from '~~/drizzle/schema';

export type Topic = typeof schema.topics.$inferSelect;
export type Agent = typeof schema.agents.$inferSelect;
export type AgentWithTopics = Agent & { topics: Topic[] };

export const useAgents = async () => {
    const agents = useState<AgentWithTopics[]>('agents_state', () => []);
    const loaded = useState('agents_loaded', () => false);

    const { refresh } = await useFetch<AgentWithTopics[]>('/api/agents', {
        key: 'agents_request',
        immediate: !loaded.value,
        onRequest() {
            loaded.value = true;
        },
        onResponse({ response }) {
            if (response.ok) {
                agents.value = response._data ?? [];
            }
        }
    });

    const createAgent = async (navigate: boolean = true) => {
        const { user } = useAuth();
        if (!user.value) {
            console.error('No user');
            return null;
        }

        const agentId = nanoid();

        const agent = {
            id: agentId,
            userId: user.value.id,
            name: 'New Agent',
            systemPrompt: 'You are a helpful assistant.',
            defaultModelId: null,
            imageUrl: null,
        } as AgentWithTopics;

        const router = useRouter()

        // @ts-ignore - stack depth
        const res = await attempt($fetch('/api/agent', {
            method: 'POST',
            body: agent,
            onRequest() {
                agents.value = [...(agents.value), { ...agent, topics: [] as Topic[], createdAt: new Date() }];
                if (navigate) {
                    router.push(`/agent/${agentId}`);
                }
            },
            onRequestError() {
                if (navigate) {
                    const route = useRoute();
                    if (route.params.id === agentId) {
                        router.push('/');
                    }
                }
                agents.value = agents.value.filter(a => a.id !== agentId);
            },
            onResponseError() {
                if (navigate) {
                    const route = useRoute();
                    if (route.params.id === agentId) {
                        router.push('/');
                    }
                }
                agents.value = agents.value.filter(a => a.id !== agentId);
            },
            async onResponse() {
                await refresh();
            }
        }));

        if (!res.ok) return null;

        return agent;
    };

    const patchAgentLocally = (id: string, updates: Partial<Agent>) => {
        if (!agents.value) return null;

        agents.value = agents.value.map(a =>
            a.id === id ? { ...a, ...updates } : a
        );
    };

    const patchTopicLocally = (id: string, updates: Partial<Topic>) => {
        if (!agents.value) return null;

        agents.value = agents.value.map(a =>
            a.topics.find(t => t.id === id) ? {
                ...a, topics: a.topics.map(t =>
                    t.id === id ? { ...t, ...updates } : t
                )
            } : a
        );
    };

    const updateAgent = async (id: string, updates: Partial<Agent>) => {
        const agent = agents.value.find(a => a.id === id);
        if (!agent) return;

        await $fetch(`/api/agent/${id}`, {
            method: 'PATCH',
            body: updates,
            onRequest() {
                agents.value = agents.value.map(a =>
                    a.id === id ? { ...a, ...updates } : a
                );
            },
            onRequestError() {
                agents.value = agents.value.map(a =>
                    a.id === id ? agent : a
                );
            },
            onResponseError() {
                agents.value = agents.value.map(a =>
                    a.id === id ? agent : a
                );
            },
            // async onResponse() {
            //     await refresh();
            // }
        });
    }

    const getAgent = (id: MaybeRef<string>) => {
        return computed(() => agents.value?.find((agent) => agent.id === unref(id)) || null);
    }

    const deleteAgent = async (id: string) => {
        let agent = agents.value.find(a => a.id === id);
        if (!agent) return;

        await $fetch(`/api/agent/${id}`, {
            method: 'DELETE',
            onRequest() {
                agents.value = agents.value.filter(a => a.id !== id);
            },
            onResponseError() {
                agents.value = [...agents.value.filter(a => a.id !== id), agent];
            },
            async onResponse() {
                await refresh();
            }
        });
    }

    // TODO: since topics are contained within each agent struct, all topic
    // actions must be done through the agents composable
    const createTopic = async (agentId: string) => {
        const topicId = nanoid();
        const topic = {
            id: topicId,
            name: 'New Topic',
            agentId,
        };

        await $fetch(`/api/topic`, {
            method: 'POST',
            body: topic,
            onRequest() {
                agents.value = agents.value.map(a =>
                    a.id === agentId ? { ...a, topics: [{ ...topic, createdAt: new Date() }, ...a.topics] } as AgentWithTopics : a
                );
            },
            onResponseError() {
                agents.value = agents.value.map(a =>
                    a.id === agentId ? { ...a, topics: a.topics.filter(t => t.id !== topicId) } : a
                );
            },
        });

        return topic;
    }

    const deleteTopic = async (agentId: string, topicId: string) => {
        const targetTopic = agents.value.flatMap(agent => agent.topics).find(topic => topic.id === topicId);
        if (!targetTopic) return;

        await $fetch(`/api/topic/${topicId}`, {
            method: 'DELETE',
            onRequest() {
                agents.value = agents.value.map(a =>
                    a.id === agentId ? { ...a, topics: a.topics.filter(t => t.id !== topicId) } : a
                );
            },
            onResponseError() {
                agents.value = agents.value.map(a =>
                    a.id === agentId ? { ...a, topics: [...a.topics.filter(t => t.id !== topicId), targetTopic] } : a
                );
            },
        });
    }

    const forkTopic = async (topicId: string, messageId: string): Promise<string | null> => {
        try {
            const result = await $fetch<{ ok: boolean; topicId: string; name: string }>(`/api/topic/${topicId}/fork`, {
                method: 'POST',
                body: { messageId },
            });

            if (result.ok) {
                const agent = agents.value.find(a => a.topics.some(t => t.id === topicId));
                if (agent) {
                    const newTopic: Topic = {
                        id: result.topicId,
                        name: result.name,
                        agentId: agent.id,
                        userId: agent.userId,
                        renaming: false,
                        createdAt: new Date(),
                    };
                    agents.value = agents.value.map(a =>
                        a.id === agent.id ? { ...a, topics: [newTopic, ...a.topics] } as AgentWithTopics : a
                    );
                }
                return result.topicId;
            }
            return null;
        } catch (error) {
            console.error('Failed to fork topic:', error);
            return null;
        }
    };

    return {
        agents,
        refresh,
        createAgent,
        createTopic,
        forkTopic,
        getAgent,
        patchAgentLocally,
        patchTopicLocally,
        updateAgent,
        deleteAgent,
        deleteTopic
    };
}