ChatFM.swift

import SwiftUI
import Foundation
import FoundationModels
import AudioToolbox

enum Role: Hashable {
    case user
    case assistant
}

struct Message: Identifiable {
    var id: UUID = .init()
    var role: Role
    var text: String
}

struct ChatView: View {
    private static let instructions: String = """
    You are a helpful assistant. Answer the user's questions in the same language that the user uses.
    """

    @State private var messages: [Message] = []
    @State private var input: String = ""
    @State private var presentsClearAlert: Bool = false

    @State private var session: LanguageModelSession = .init(instructions: Self.instructions)
    @State private var currentRespondingTask: Task<Void, any Error>?

    var isResponding: Bool { currentRespondingTask != nil }
    var canSendMessage: Bool { !input.isEmpty && !isResponding }

    var body: some View {
        NavigationStack {
            GeometryReader { geometry in
                ScrollViewReader { scrollViewProxy in
                    ScrollView {
                        VStack {
                            ForEach(messages) { message in
                                MessageView(message: message)
                                    .frame(maxWidth: .infinity, alignment: message.role == .user ? .trailing : .leading)
                                    .padding(message.role == .user ? .leading : .trailing, 40)
                                    .padding(.horizontal)
                                    .id(message.id)
                            }
                            if isResponding {
                                ProgressView()
                                    .frame(maxWidth: .infinity, alignment: .leading)
                                    .padding(.horizontal)
                            }
                            Spacer()
                            TextField("Ask anything", text: $input)
                                .textFieldStyle(.roundedBorder)
                                .submitLabel(.send)
                                .onSubmit {
                                    sendMessage(scrollToBottom: {
                                        withAnimation {
                                            scrollViewProxy.scrollTo("messages", anchor: .bottom)
                                        }
                                    })
                                }
                                .padding(.horizontal)
                        }
                        .padding(.vertical)
                        .frame(minHeight: geometry.size.height)
                        .frame(width: geometry.size.width)
                        .id("messages")
                    }
                }
                .frame(height: geometry.size.height)
            }
            .toolbar {
                ToolbarItem(placement: .topBarTrailing) {
                    Button {
                        presentsClearAlert = true
                    } label: {
                        Image(systemName: "square.and.pencil")
                    }
                }
            }
            .alert("Confirmation", isPresented: $presentsClearAlert) {
                Button(role: .cancel) {} label: {
                    Text("Cancel")
                }
                Button(role: .destructive) {
                    clearChat()
                } label: {
                    Text("Clear")
                }
            } message: {
                Text("Are you sure you want to clear the current chat?")
            }
            .navigationTitle(Text("ChatFM"))
            .navigationBarTitleDisplayMode(.inline)
        }
    }

    func sendMessage(scrollToBottom: (() -> Void)?) {
        guard canSendMessage else { return }

        let text = input
        messages.append(.init(role: .user, text: text))
        input = ""
        playSendSound()
        if let scrollToBottom {
            Task {
                scrollToBottom()
            }
        }

        currentRespondingTask = Task {
            defer { currentRespondingTask = nil }

            do {
                let response = try await session.respond(to: text)
                try Task.checkCancellation()
                playReceiveSound()
                let newMessage: Message = .init(role: .assistant, text: response.content)
                messages.append(newMessage)
                if let scrollToBottom {
                    Task {
                        scrollToBottom()
                    }
                }
            } catch is CancellationError {
                // Do nothing
            } catch {
                // TODO: Error handling
                print(error)
            }
        }
    }

    func clearChat() {
        messages = []
        currentRespondingTask?.cancel()
        currentRespondingTask = nil
        session = LanguageModelSession(instructions: Self.instructions)
    }

    func playSendSound() {
        AudioServicesPlaySystemSound(1004)

    }

    func playReceiveSound() {
        AudioServicesPlaySystemSound(1003)
    }
}

struct MessageView: View {
    let message: Message

    var body: some View {
        Text(verbatim: message.text)
            .fixedSize(horizontal: false, vertical: true)
            .foregroundStyle(message.role == .user ? .white : .primary)
            .selectionDisabled(false)
            .padding()
            .background {
                let color: Color = switch message.role {
                case .user: .blue
                case .assistant: .init(uiColor: .systemGray6)
                }
                color.cornerRadius(24)
            }
    }
}