swift-apple-intelligence-grpc/Sources/AppleIntelligenceApp/ViewModels/ChatViewModel.swift

import Foundation
import AppleIntelligenceCore

@MainActor
@Observable
final class ChatViewModel {
    var messages: [ChatMessage] = []
    var inputText: String = ""
    var isLoading: Bool = false
    var errorMessage: String?

    private var service: AppleIntelligenceService?
    private var currentTask: Task<Void, Never>?

    func initialize() async {
        service = await AppleIntelligenceService()
    }

    var isServiceAvailable: Bool {
        get async {
            await service?.isAvailable ?? false
        }
    }

    func sendMessage() {
        let text = inputText.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !text.isEmpty else { return }
        guard !isLoading else { return }

        // Add user message
        let userMessage = ChatMessage(role: .user, content: text)
        messages.append(userMessage)
        inputText = ""
        errorMessage = nil

        // Add placeholder for assistant response
        var assistantMessage = ChatMessage(role: .assistant, content: "", isStreaming: true)
        messages.append(assistantMessage)

        isLoading = true

        currentTask = Task {
            do {
                guard let service = service else {
                    throw AppleIntelligenceError.modelNotAvailable
                }

                let stream = await service.streamComplete(
                    prompt: text,
                    temperature: nil,
                    maxTokens: nil
                )

                var fullResponse = ""
                for try await partialResponse in stream {
                    fullResponse = partialResponse
                    // Update the last message (assistant's response)
                    if let index = messages.lastIndex(where: { $0.role == .assistant }) {
                        messages[index].content = fullResponse
                    }
                }

                // Mark streaming as complete
                if let index = messages.lastIndex(where: { $0.role == .assistant }) {
                    messages[index].isStreaming = false
                }

            } catch {
                errorMessage = error.localizedDescription
                // Remove the empty assistant message on error
                if let index = messages.lastIndex(where: { $0.role == .assistant && $0.content.isEmpty }) {
                    messages.remove(at: index)
                }
            }

            isLoading = false
        }
    }

    func stopGeneration() {
        currentTask?.cancel()
        currentTask = nil
        isLoading = false

        // Mark any streaming message as complete
        if let index = messages.lastIndex(where: { $0.isStreaming }) {
            messages[index].isStreaming = false
        }
    }

    func clearChat() {
        stopGeneration()
        messages.removeAll()
        errorMessage = nil
    }
}