android/app/ChatScreen.kt

package com.minimind.mind2.ui

import androidx.compose.foundation.layout.*
import androidx.compose.foundation.lazy.LazyColumn
import androidx.compose.foundation.lazy.items
import androidx.compose.foundation.lazy.rememberLazyListState
import androidx.compose.foundation.shape.RoundedCornerShape
import androidx.compose.material.icons.Icons
import androidx.compose.material.icons.filled.Send
import androidx.compose.material.icons.filled.Stop
import androidx.compose.material3.*
import androidx.compose.runtime.*
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.graphics.Color
import androidx.compose.ui.text.font.FontWeight
import androidx.compose.ui.unit.dp
import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewModelScope
import com.minimind.mind2.Mind2Model
import kotlinx.coroutines.flow.catch
import kotlinx.coroutines.launch

/**
 * Chat ViewModel for MiniMind
 */
class ChatViewModel : ViewModel() {
    private val model = Mind2Model.getInstance()

    var messages = mutableStateListOf<ChatMessage>()
        private set

    var isGenerating by mutableStateOf(false)
        private set

    var isLoading by mutableStateOf(false)
        private set

    var error by mutableStateOf<String?>(null)
        private set

    var modelInfo by mutableStateOf("")
        private set

    private var currentResponse = StringBuilder()

    fun loadModel(context: android.content.Context, modelName: String = "mind2-lite.gguf") {
        viewModelScope.launch {
            isLoading = true
            error = null

            model.load(context, modelName)
                .onSuccess {
                    modelInfo = model.getInfo()
                }
                .onFailure {
                    error = "Failed to load model: ${it.message}"
                }

            isLoading = false
        }
    }

    fun sendMessage(content: String) {
        if (content.isBlank() || isGenerating) return

        // Add user message
        messages.add(ChatMessage("user", content))

        // Add placeholder for assistant
        currentResponse.clear()
        messages.add(ChatMessage("assistant", ""))

        isGenerating = true
        error = null

        val history = messages.dropLast(1).map {
            Mind2Model.ChatMessage(it.role, it.content)
        }

        viewModelScope.launch {
            model.chatStream(content, history)
                .catch { e ->
                    error = "Generation error: ${e.message}"
                    isGenerating = false
                }
                .collect { token ->
                    currentResponse.append(token)
                    // Update last message
                    val lastIndex = messages.lastIndex
                    messages[lastIndex] = ChatMessage("assistant", currentResponse.toString())
                }

            isGenerating = false
        }
    }

    fun stopGeneration() {
        model.stop()
        isGenerating = false
    }

    fun clearChat() {
        messages.clear()
        currentResponse.clear()
    }

    override fun onCleared() {
        super.onCleared()
        model.release()
    }
}

data class ChatMessage(
    val role: String,
    val content: String
)

/**
 * Chat Screen Composable
 */
@OptIn(ExperimentalMaterial3Api::class)
@Composable
fun ChatScreen(
    viewModel: ChatViewModel
) {
    var inputText by remember { mutableStateOf("") }
    val listState = rememberLazyListState()

    // Auto-scroll to bottom when new messages arrive
    LaunchedEffect(viewModel.messages.size) {
        if (viewModel.messages.isNotEmpty()) {
            listState.animateScrollToItem(viewModel.messages.lastIndex)
        }
    }

    Scaffold(
        topBar = {
            TopAppBar(
                title = {
                    Column {
                        Text("MiniMind", fontWeight = FontWeight.Bold)
                        if (viewModel.isLoading) {
                            Text(
                                "Loading model...",
                                style = MaterialTheme.typography.bodySmall,
                                color = MaterialTheme.colorScheme.onSurfaceVariant
                            )
                        }
                    }
                },
                colors = TopAppBarDefaults.topAppBarColors(
                    containerColor = MaterialTheme.colorScheme.primaryContainer
                )
            )
        }
    ) { padding ->
        Column(
            modifier = Modifier
                .fillMaxSize()
                .padding(padding)
        ) {
            // Error banner
            viewModel.error?.let { errorMsg ->
                Surface(
                    color = MaterialTheme.colorScheme.errorContainer,
                    modifier = Modifier.fillMaxWidth()
                ) {
                    Text(
                        text = errorMsg,
                        color = MaterialTheme.colorScheme.onErrorContainer,
                        modifier = Modifier.padding(16.dp)
                    )
                }
            }

            // Messages list
            LazyColumn(
                state = listState,
                modifier = Modifier
                    .weight(1f)
                    .fillMaxWidth(),
                contentPadding = PaddingValues(16.dp),
                verticalArrangement = Arrangement.spacedBy(12.dp)
            ) {
                items(viewModel.messages) { message ->
                    MessageBubble(message)
                }
            }

            // Input area
            Surface(
                tonalElevation = 3.dp,
                modifier = Modifier.fillMaxWidth()
            ) {
                Row(
                    modifier = Modifier
                        .padding(16.dp)
                        .fillMaxWidth(),
                    verticalAlignment = Alignment.CenterVertically
                ) {
                    OutlinedTextField(
                        value = inputText,
                        onValueChange = { inputText = it },
                        modifier = Modifier.weight(1f),
                        placeholder = { Text("Type a message...") },
                        shape = RoundedCornerShape(24.dp),
                        enabled = !viewModel.isLoading && !viewModel.isGenerating
                    )

                    Spacer(modifier = Modifier.width(8.dp))

                    if (viewModel.isGenerating) {
                        FilledIconButton(
                            onClick = { viewModel.stopGeneration() },
                            colors = IconButtonDefaults.filledIconButtonColors(
                                containerColor = MaterialTheme.colorScheme.error
                            )
                        ) {
                            Icon(Icons.Default.Stop, contentDescription = "Stop")
                        }
                    } else {
                        FilledIconButton(
                            onClick = {
                                viewModel.sendMessage(inputText)
                                inputText = ""
                            },
                            enabled = inputText.isNotBlank() && !viewModel.isLoading
                        ) {
                            Icon(Icons.Default.Send, contentDescription = "Send")
                        }
                    }
                }
            }
        }
    }
}

@Composable
fun MessageBubble(message: ChatMessage) {
    val isUser = message.role == "user"

    Row(
        modifier = Modifier.fillMaxWidth(),
        horizontalArrangement = if (isUser) Arrangement.End else Arrangement.Start
    ) {
        Surface(
            shape = RoundedCornerShape(
                topStart = 16.dp,
                topEnd = 16.dp,
                bottomStart = if (isUser) 16.dp else 4.dp,
                bottomEnd = if (isUser) 4.dp else 16.dp
            ),
            color = if (isUser)
                MaterialTheme.colorScheme.primary
            else
                MaterialTheme.colorScheme.surfaceVariant,
            modifier = Modifier.widthIn(max = 300.dp)
        ) {
            Text(
                text = message.content.ifEmpty { "..." },
                modifier = Modifier.padding(12.dp),
                color = if (isUser)
                    MaterialTheme.colorScheme.onPrimary
                else
                    MaterialTheme.colorScheme.onSurfaceVariant
            )
        }
    }
}