Spanish/Conjuga/Conjuga/Views/Practice/Chat/ChatView.swift

import SwiftUI
import SharedModels
import SwiftData
import FoundationModels

@Generable
private struct ChatWordInfo {
    @Guide(description: "Dictionary base form") var baseForm: String
    @Guide(description: "English translation") var english: String
    @Guide(description: "Part of speech") var partOfSpeech: String
}

struct ChatView: View {
    let conversation: Conversation

    @Environment(\.cloudModelContextProvider) private var cloudModelContextProvider
    @Environment(DictionaryService.self) private var dictionary
    @State private var service = ConversationService()
    @State private var messages: [ChatMessage] = []
    @State private var inputText = ""
    @State private var errorMessage: String?
    @State private var hasStarted = false
    @State private var selectedWord: WordAnnotation?
    @State private var lookupCache: [String: WordAnnotation] = [:]

    private var cloudContext: ModelContext { cloudModelContextProvider() }

    var body: some View {
        VStack(spacing: 0) {
            // Messages
            ScrollViewReader { proxy in
                ScrollView {
                    LazyVStack(spacing: 12) {
                        ForEach(messages) { message in
                            ChatBubble(message: message, dictionary: dictionary, lookupCache: $lookupCache) { word in
                                selectedWord = word
                            }
                            .id(message.id)
                        }

                        if service.isResponding {
                            HStack {
                                ProgressView()
                                    .padding(.horizontal, 16)
                                    .padding(.vertical, 10)
                                    .background(.fill.quaternary, in: RoundedRectangle(cornerRadius: 16))
                                Spacer()
                            }
                            .padding(.horizontal)
                        }
                    }
                    .padding(.vertical)
                }
                .onChange(of: messages.count) {
                    if let last = messages.last {
                        withAnimation { proxy.scrollTo(last.id, anchor: .bottom) }
                    }
                }
            }

            Divider()

            // Input
            HStack(spacing: 8) {
                TextField("Type in Spanish...", text: $inputText)
                    .textFieldStyle(.roundedBorder)
                    .autocorrectionDisabled()
                    .textInputAutocapitalization(.sentences)
                    .onSubmit { sendMessage() }

                Button {
                    sendMessage()
                } label: {
                    Image(systemName: "arrow.up.circle.fill")
                        .font(.title2)
                }
                .disabled(inputText.trimmingCharacters(in: .whitespaces).isEmpty || service.isResponding)
                .tint(.green)
            }
            .padding()
        }
        .navigationTitle(conversation.scenario)
        .navigationBarTitleDisplayMode(.inline)
        .sheet(item: $selectedWord) { word in
            ChatWordDetailSheet(word: word)
                .presentationDetents([.height(200)])
        }
        .alert("Error", isPresented: .init(
            get: { errorMessage != nil },
            set: { if !$0 { errorMessage = nil } }
        )) {
            Button("OK") { errorMessage = nil }
        } message: {
            Text(errorMessage ?? "")
        }
        .onAppear {
            messages = conversation.decodedMessages
            if !hasStarted && messages.isEmpty {
                startConversation()
            }
            hasStarted = true
        }
    }

    private func startConversation() {
        let opening = service.startConversation(scenario: conversation.scenario, level: conversation.level)
        let msg = ChatMessage(role: "assistant", content: opening)
        conversation.appendMessage(msg)
        messages = conversation.decodedMessages
        try? cloudContext.save()
    }

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

        let userMsg = ChatMessage(role: "user", content: text)
        conversation.appendMessage(userMsg)
        messages = conversation.decodedMessages
        inputText = ""
        try? cloudContext.save()

        Task {
            do {
                let response = try await service.respond(to: text)
                let assistantMsg = ChatMessage(role: "assistant", content: response)
                conversation.appendMessage(assistantMsg)
                messages = conversation.decodedMessages
                try? cloudContext.save()
            } catch {
                errorMessage = "Failed to get response: \(error.localizedDescription)"
            }
        }
    }
}

// MARK: - Chat Bubble

private struct ChatBubble: View {
    let message: ChatMessage
    let dictionary: DictionaryService
    @Binding var lookupCache: [String: WordAnnotation]
    let onWordTap: (WordAnnotation) -> Void

    private var isUser: Bool { message.role == "user" }

    var body: some View {
        HStack {
            if isUser { Spacer(minLength: 60) }

            VStack(alignment: isUser ? .trailing : .leading, spacing: 4) {
                if isUser {
                    Text(message.content)
                        .font(.body)
                        .padding(.horizontal, 14)
                        .padding(.vertical, 10)
                        .background(.green.opacity(0.2), in: RoundedRectangle(cornerRadius: 16))
                } else {
                    tappableBubble
                }

                if let correction = message.correction, !correction.isEmpty {
                    Text(correction)
                        .font(.caption)
                        .foregroundStyle(.orange)
                        .padding(.horizontal, 4)
                }
            }

            if !isUser { Spacer(minLength: 60) }
        }
        .padding(.horizontal)
    }

    private var tappableBubble: some View {
        let words = message.content.components(separatedBy: " ")
        return ChatFlowLayout(spacing: 0) {
            ForEach(Array(words.enumerated()), id: \.offset) { _, word in
                ChatWordButton(word: word, dictionary: dictionary, cache: lookupCache) { annotation in
                    if annotation.english.isEmpty {
                        lookupWordAsync(annotation.word)
                    } else {
                        onWordTap(annotation)
                    }
                }
            }
        }
        .padding(.horizontal, 14)
        .padding(.vertical, 10)
        .background(.fill.quaternary, in: RoundedRectangle(cornerRadius: 16))
    }

    private func lookupWordAsync(_ word: String) {
        // Try dictionary first
        if let entry = dictionary.lookup(word) {
            let annotation = WordAnnotation(word: word, baseForm: entry.baseForm, english: entry.english, partOfSpeech: entry.partOfSpeech)
            lookupCache[word] = annotation
            onWordTap(annotation)
            return
        }

        // Show loading then AI lookup
        onWordTap(WordAnnotation(word: word, baseForm: word, english: "Looking up...", partOfSpeech: ""))

        Task {
            do {
                let session = LanguageModelSession(instructions: "You are a Spanish dictionary. Provide base form, English translation, and part of speech.")
                let response = try await session.respond(to: "Word: \"\(word)\"", generating: ChatWordInfo.self)
                let info = response.content
                let annotation = WordAnnotation(word: word, baseForm: info.baseForm, english: info.english, partOfSpeech: info.partOfSpeech)
                lookupCache[word] = annotation
                onWordTap(annotation)
            } catch {
                onWordTap(WordAnnotation(word: word, baseForm: word, english: "Lookup unavailable", partOfSpeech: ""))
            }
        }
    }
}

// MARK: - Chat Word Button

private struct ChatWordButton: View {
    let word: String
    let dictionary: DictionaryService
    let cache: [String: WordAnnotation]
    let onTap: (WordAnnotation) -> Void

    private var cleaned: String {
        word.lowercased()
            .trimmingCharacters(in: .punctuationCharacters)
            .trimmingCharacters(in: .whitespaces)
    }

    private var annotation: WordAnnotation? {
        if let cached = cache[cleaned] { return cached }
        if let entry = dictionary.lookup(cleaned) {
            return WordAnnotation(word: cleaned, baseForm: entry.baseForm, english: entry.english, partOfSpeech: entry.partOfSpeech)
        }
        return nil
    }

    var body: some View {
        Button {
            onTap(annotation ?? WordAnnotation(word: cleaned, baseForm: cleaned, english: "", partOfSpeech: ""))
        } label: {
            Text(word + " ")
                .font(.body)
                .foregroundStyle(.primary)
                .underline(annotation != nil, color: .teal.opacity(0.3))
        }
        .buttonStyle(.plain)
    }
}

// MARK: - Word Detail Sheet

private struct ChatWordDetailSheet: View {
    let word: WordAnnotation

    var body: some View {
        VStack(spacing: 16) {
            HStack {
                Text(word.word)
                    .font(.title2.bold())
                Spacer()
                if !word.partOfSpeech.isEmpty {
                    Text(word.partOfSpeech)
                        .font(.caption.weight(.medium))
                        .foregroundStyle(.secondary)
                        .padding(.horizontal, 8)
                        .padding(.vertical, 4)
                        .background(.fill.tertiary, in: Capsule())
                }
            }

            Divider()

            if word.english == "Looking up..." {
                HStack(spacing: 8) {
                    ProgressView()
                    Text("Looking up word...")
                        .font(.subheadline)
                        .foregroundStyle(.secondary)
                }
                .frame(maxWidth: .infinity)
            } else {
                VStack(alignment: .leading, spacing: 8) {
                    if !word.baseForm.isEmpty && word.baseForm != word.word {
                        HStack {
                            Text("Base form:")
                                .font(.subheadline)
                                .foregroundStyle(.secondary)
                            Text(word.baseForm)
                                .font(.subheadline.weight(.semibold))
                                .italic()
                        }
                    }
                    if !word.english.isEmpty {
                        HStack {
                            Text("English:")
                                .font(.subheadline)
                                .foregroundStyle(.secondary)
                            Text(word.english)
                                .font(.subheadline.weight(.semibold))
                        }
                    }
                }
                .frame(maxWidth: .infinity, alignment: .leading)
            }

            Spacer()
        }
        .padding()
    }
}

// MARK: - Chat Flow Layout

private struct ChatFlowLayout: Layout {
    var spacing: CGFloat = 0
    func sizeThatFits(proposal: ProposedViewSize, subviews: Subviews, cache: inout ()) -> CGSize {
        let rows = computeRows(proposal: proposal, subviews: subviews)
        var height: CGFloat = 0
        for row in rows { height += row.map { $0.height }.max() ?? 0 }
        height += CGFloat(max(0, rows.count - 1)) * spacing
        return CGSize(width: proposal.width ?? 0, height: height)
    }
    func placeSubviews(in bounds: CGRect, proposal: ProposedViewSize, subviews: Subviews, cache: inout ()) {
        let rows = computeRows(proposal: proposal, subviews: subviews)
        var y = bounds.minY; var idx = 0
        for row in rows {
            var x = bounds.minX; let rh = row.map { $0.height }.max() ?? 0
            for size in row {
                subviews[idx].place(at: CGPoint(x: x, y: y), proposal: ProposedViewSize(size))
                x += size.width; idx += 1
            }
            y += rh + spacing
        }
    }
    private func computeRows(proposal: ProposedViewSize, subviews: Subviews) -> [[CGSize]] {
        let mw = proposal.width ?? .infinity; var rows: [[CGSize]] = [[]]; var cw: CGFloat = 0
        for sv in subviews {
            let s = sv.sizeThatFits(.unspecified)
            if cw + s.width > mw && !rows[rows.count - 1].isEmpty { rows.append([]); cw = 0 }
            rows[rows.count - 1].append(s); cw += s.width
        }
        return rows
    }
}