'use server';

import { streamUI, tool } from 'ai/rsc';
import { createOpenAI as createGroq } from '@ai-sdk/openai';
import { z } from 'zod';

const groq = createGroq({
  baseURL: 'https://api.groq.com/openai/v1',
  apiKey: process.env.GROQ_API_KEY,
});

export async function streamComponent() {
  const result = await streamUI({
    model: groq('llama-3.1-70b-versatile'),
    prompt: 'Get the weather for San Francisco',
    text: ({ content }) => <div>{content}</div>,
    tools: {
      getWeather: tool({
        description: 'Get the weather for a location',
        parameters: z.object({ location: z.string() }),
        generate: async function* ({ location }) {
          yield <div>loading...</div>;
          const weather = '25c'; // await getWeather(location);
          return (
            <div>
              the weather in {location} is {weather}.
            </div>
          );
        },
      }),
    },
  });
  return result.value;
}

result = await streamUI({
      model: getModelFunction(model),
      initial: <SpinnerMessage model={model} />,
      maxTokens: getMaxTokens(model),
      system: getPrompt(systemPrompt),
      messages: [
        ...aiState.get().messages.map((message: any) => ({
          role: message.role,
          content: message.content,
          name: message.name,
          model: message.model
        }))
      ],
      text: ({ content, done, delta }) => {
        if (!textStream) {
          textStream = createStreamableValue('')
          textNode = <BotMessage model={model} content={textStream.value} />
        }

        if (done) {
          textStream.done()
          aiState.done({
            ...aiState.get(),
            messages: [
              ...aiState.get().messages,
              {
                id: nanoid(),
                role: 'assistant',
                content,
                model
              }
            ]
          })
        } else {
          textStream.update(delta)
        }
        return textNode
      },
      tools: {
        generateImage: {
          description: 'Generate an image based on the user message.',
          parameters: z.object({
            text: z.string().describe('A detailed description of the image to be generated.')
          }),
          generate: async function* ({ text }: { text: string }) {
            yield (
              <BotCard model={model}>
                {spinner}
              </BotCard>
            )

            await sleep(1000)
            const toolCallId = nanoid()

            let result: ImageResult[] = await generateImage({
              prompt: text,
            });


            aiState.done({
              ...aiState.get(),
              messages: [
                ...aiState.get().messages,
                {
                  id: nanoid(),
                  role: 'assistant',
                  content: [
                    {
                      type: 'tool-call',
                      toolName: 'generateImage',
                      toolCallId,
                      args: { text }
                    }
                  ],
                  model
                },
                {
                  id: nanoid(),
                  role: 'tool',
                  content: [
                    {
                      type: 'tool-result',
                      toolName: 'generateImage',
                      toolCallId,
                      result
                    }
                  ],
                  model
                }
              ]
            })
            return (
              <BotMessageImage result={result} model={model} />
            )
          },
        }
      }

    })
  } catch (error: any) {
    try {
      result = {
        value: error.responseBody ? JSON.parse(error.responseBody).error.message : error.message
      }
    } catch (err: any) {
      result = {
        value: (<h1>{error.message}</h1>)
      }
    }
  }
  return {
    id: nanoid(),
    display: result.value
  }
}