components/AiChat.tsx

'use client'

import React, { useState, useEffect, useRef } from 'react'
import { Button } from "@/components/ui/button"
import { Input } from "@/components/ui/input"
import { Textarea } from "@/components/ui/textarea"
import { Card, CardContent } from "@/components/ui/card"
import { ChevronDown, ChevronUp, Send } from 'lucide-react'

export default function AIChat() {
  const [messages, setMessages] = useState([])
  const [input, setInput] = useState('')
  const [apiKey, setApiKey] = useState('')
  const [thinking, setThinking] = useState('')
  const [isThinkingVisible, setIsThinkingVisible] = useState(false)
  const chatEndRef = useRef(null)

  useEffect(() => {
    chatEndRef.current?.scrollIntoView({ behavior: 'smooth' })
  }, [messages])

  const handleSubmit = async (e) => {
    e.preventDefault()
    if (!input.trim()) return

    const newMessage = { role: 'user', content: input }
    setMessages([...messages, newMessage])
    setInput('')

    try {
      // Call /think endpoint
      const thinkResponse = await fetch('/api/think', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ messages: [...messages, newMessage], api_key: apiKey }),
      })

      if (!thinkResponse.ok) throw new Error('Think request failed')

      const reader = thinkResponse.body.getReader()
      let thinkingContent = ''

      while (true) {
        const { done, value } = await reader.read()
        if (done) break
        const chunk = new TextDecoder().decode(value)
        thinkingContent += chunk
        setThinking(thinkingContent)
      }

      // Call /chat endpoint with the thinking result
      const chatResponse = await fetch('/api/chat', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ messages: [...messages, newMessage, { role: 'assistant', content: thinkingContent }], api_key: apiKey }),
      })

      if (!chatResponse.ok) throw new Error('Chat request failed')

      const chatReader = chatResponse.body.getReader()
      let chatContent = ''

      while (true) {
        const { done, value } = await chatReader.read()
        if (done) break
        const chunk = new TextDecoder().decode(value)
        chatContent += chunk
        setMessages(prev => [...prev.slice(0, -1), { role: 'assistant', content: chatContent }])
      }
    } catch (error) {
      console.error('Error:', error)
      setMessages(prev => [...prev, { role: 'assistant', content: 'An error occurred. Please try again.' }])
    }
  }

  return (
    <div className="min-h-screen bg-gray-900 text-white p-4 flex flex-col">
      <Card className="flex-grow overflow-auto mb-4 bg-gray-800 bg-opacity-50 backdrop-filter backdrop-blur-lg">
        <CardContent className="p-4">
          {messages.map((message, index) => (
            <div key={index} className={`mb-4 ${message.role === 'user' ? 'text-right' : 'text-left'}`}>
              <div className={`inline-block p-2 rounded-lg ${message.role === 'user' ? 'bg-blue-600' : 'bg-gray-700'}`}>
                {message.content}
              </div>
              {message.role === 'assistant' && thinking && (
                <div className="mt-2">
                  <Button
                    variant="ghost"
                    size="sm"
                    onClick={() => setIsThinkingVisible(!isThinkingVisible)}
                    className="text-xs text-gray-400"
                  >
                    {isThinkingVisible ? <ChevronUp className="h-4 w-4" /> : <ChevronDown className="h-4 w-4" />}
                    Thinking Process
                  </Button>
                  {isThinkingVisible && (
                    <div className="mt-2 text-xs text-gray-400 bg-gray-800 p-2 rounded">
                      {thinking.split('\n').map((line, i) => {
                        if (line.startsWith('<thinking>') || line.startsWith('<step>')) {
                          return <p key={i}>{line.replace(/<\/?thinking>|<\/?step>/g, '')}</p>
                        }
                        return null
                      })}
                    </div>
                  )}
                </div>
              )}
            </div>
          ))}
          <div ref={chatEndRef} />
        </CardContent>
      </Card>
      <form onSubmit={handleSubmit} className="flex gap-2">
        <Input
          type="text"
          value={apiKey}
          onChange={(e) => setApiKey(e.target.value)}
          placeholder="Enter API Key (optional)"
          className="flex-grow bg-gray-800 text-white"
        />
        <Textarea
          value={input}
          onChange={(e) => setInput(e.target.value)}
          placeholder="Type your message..."
          className="flex-grow bg-gray-800 text-white"
        />
        <Button type="submit" className="bg-blue-600 hover:bg-blue-700">
          <Send className="h-4 w-4" />
        </Button>
      </form>
    </div>
  )
}