mimo-studio-android/app/src/main/java/com/mimo/studio/viewmodel/MainViewModel.kt

package com.mimo.studio.viewmodel

import android.app.Application
import android.content.Context
import androidx.lifecycle.AndroidViewModel
import androidx.lifecycle.viewModelScope
import com.google.gson.Gson
import com.google.gson.reflect.TypeToken
import com.mimo.studio.api.RetrofitClient
import com.mimo.studio.model.*
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext
import java.io.BufferedReader
import java.io.InputStreamReader

class MainViewModel(application: Application) : AndroidViewModel(application) {

    private val prefs = application.getSharedPreferences("mimo_prefs", Context.MODE_PRIVATE)
    private val gson = Gson()

    // State
    private val _models = MutableStateFlow<List<ModelInfo>>(emptyList())
    val models: StateFlow<List<ModelInfo>> = _models

    private val _conversations = MutableStateFlow<List<Conversation>>(emptyList())
    val conversations: StateFlow<List<Conversation>> = _conversations

    private val _currentConversation = MutableStateFlow<Conversation?>(null)
    val currentConversation: StateFlow<Conversation?> = _currentConversation

    private val _isLoading = MutableStateFlow(false)
    val isLoading: StateFlow<Boolean> = _isLoading

    private val _isStreaming = MutableStateFlow(false)
    val isStreaming: StateFlow<Boolean> = _isStreaming

    private val _streamingContent = MutableStateFlow("")
    val streamingContent: StateFlow<String> = _streamingContent

    private val _error = MutableStateFlow<String?>(null)
    val error: StateFlow<String?> = _error

    private val _selectedModel = MutableStateFlow("xiaomi/mimo-v2-flash")
    val selectedModel: StateFlow<String> = _selectedModel

    private val _useStreaming = MutableStateFlow(true)
    val useStreaming: StateFlow<Boolean> = _useStreaming

    private val _temperature = MutableStateFlow(1.0f)
    val temperature: StateFlow<Float> = _temperature

    private val _maxTokens = MutableStateFlow(4096)
    val maxTokens: StateFlow<Int> = _maxTokens

    init {
        loadSettings()
        loadConversations()
        if (conversations.value.isEmpty()) {
            newConversation()
        }
    }

    // ===== Settings =====
    fun getApiKey(): String = prefs.getString("api_key", "") ?: ""
    fun getBaseUrl(): String = prefs.getString("base_url", "https://api.xiaomimimo.com/") ?: "https://api.xiaomimimo.com/"

    fun saveSettings(apiKey: String, baseUrl: String) {
        prefs.edit()
            .putString("api_key", apiKey)
            .putString("base_url", baseUrl)
            .apply()
        RetrofitClient.configure(apiKey, baseUrl)
    }

    private fun loadSettings() {
        val key = getApiKey()
        val url = getBaseUrl()
        if (key.isNotBlank()) {
            RetrofitClient.configure(key, url)
        }
        _selectedModel.value = prefs.getString("selected_model", "xiaomi/mimo-v2-flash") ?: "xiaomi/mimo-v2-flash"
        _useStreaming.value = prefs.getBoolean("use_streaming", true)
        _temperature.value = prefs.getFloat("temperature", 1.0f)
        _maxTokens.value = prefs.getInt("max_tokens", 4096)
    }

    fun setSelectedModel(model: String) {
        _selectedModel.value = model
        prefs.edit().putString("selected_model", model).apply()
    }

    fun setStreaming(enabled: Boolean) {
        _useStreaming.value = enabled
        prefs.edit().putBoolean("use_streaming", enabled).apply()
    }

    fun setTemperature(temp: Float) {
        _temperature.value = temp
        prefs.edit().putFloat("temperature", temp).apply()
    }

    fun setMaxTokens(tokens: Int) {
        _maxTokens.value = tokens
        prefs.edit().putInt("max_tokens", tokens).apply()
    }

    // ===== Models =====
    fun fetchModels() {
        if (!RetrofitClient.isConfigured()) {
            _error.value = "请先配置 API Key"
            return
        }
        viewModelScope.launch {
            try {
                val response = withContext(Dispatchers.IO) { RetrofitClient.api.listModels() }
                _models.value = response.data
                _error.value = null
            } catch (e: Exception) {
                _error.value = "获取模型列表失败: ${e.message}"
            }
        }
    }

    // ===== Conversations =====
    fun newConversation() {
        val conv = Conversation(model = _selectedModel.value)
        val list = _conversations.value.toMutableList()
        list.add(0, conv)
        _conversations.value = list
        _currentConversation.value = conv
        saveConversations()
    }

    fun selectConversation(id: String) {
        _currentConversation.value = _conversations.value.find { it.id == id }
    }

    fun deleteConversation(id: String) {
        val list = _conversations.value.filter { it.id != id }.toMutableList()
        _conversations.value = list
        if (_currentConversation.value?.id == id) {
            _currentConversation.value = list.firstOrNull()
        }
        saveConversations()
    }

    private fun saveConversations() {
        val json = gson.toJson(_conversations.value)
        prefs.edit().putString("conversations", json).apply()
    }

    private fun loadConversations() {
        val json = prefs.getString("conversations", null)
        if (json != null) {
            val type = object : TypeToken<List<Conversation>>() {}.type
            try {
                val list: List<Conversation> = gson.fromJson(json, type)
                _conversations.value = list
                _currentConversation.value = list.firstOrNull()
            } catch (_: Exception) {}
        }
    }

    // ===== Chat =====
    fun sendMessage(content: String) {
        val conv = _currentConversation.value ?: return
        if (!RetrofitClient.isConfigured()) {
            _error.value = "请先在设置中配置 API Key"
            return
        }

        // Add user message
        conv.messages.add(ChatMessage(role = "user", content = content))
        if (conv.messages.size == 1 && conv.title == "新对话") {
            // Auto-title from first message
            val newConv = conv.copy(title = content.take(30) + if (content.length > 30) "..." else "")
            updateConversation(newConv)
        } else {
            _currentConversation.value = conv.copy()
        }

        if (_useStreaming.value) {
            sendStreamingMessage(conv)
        } else {
            sendNormalMessage(conv)
        }
    }

    private fun sendNormalMessage(conv: Conversation) {
        _isLoading.value = true
        _error.value = null

        viewModelScope.launch {
            try {
                val messages = buildMessages(conv)
                val request = ChatRequest(
                    model = conv.model,
                    messages = messages,
                    temperature = _temperature.value,
                    maxTokens = _maxTokens.value
                )
                val response = withContext(Dispatchers.IO) {
                    RetrofitClient.api.chatCompletions(request)
                }
                val assistantMsg = response.choices?.firstOrNull()?.message
                if (assistantMsg != null) {
                    conv.messages.add(assistantMsg)
                    updateConversation(conv.copy())
                }
                _error.value = null
            } catch (e: Exception) {
                _error.value = "请求失败: ${e.message}"
            } finally {
                _isLoading.value = false
            }
        }
    }

    private fun sendStreamingMessage(conv: Conversation) {
        _isStreaming.value = true
        _streamingContent.value = ""
        _error.value = null

        viewModelScope.launch {
            try {
                val messages = buildMessages(conv)
                val request = ChatRequest(
                    model = conv.model,
                    messages = messages,
                    stream = true,
                    temperature = _temperature.value,
                    maxTokens = _maxTokens.value
                )
                val responseBody = withContext(Dispatchers.IO) {
                    RetrofitClient.api.chatCompletionsStream(request)
                }

                val reader = BufferedReader(InputStreamReader(responseBody.byteStream()))
                val fullContent = StringBuilder()

                try {
                    var line: String?
                    while (reader.readLine().also { line = it } != null) {
                        val l = line ?: continue
                        if (!l.startsWith("data: ")) continue
                        val data = l.removePrefix("data: ").trim()
                        if (data == "[DONE]") break

                        try {
                            val chunk = gson.fromJson(data, StreamChunk::class.java)
                            val delta = chunk.choices?.firstOrNull()?.delta
                            if (delta?.content != null) {
                                fullContent.append(delta.content)
                                _streamingContent.value = fullContent.toString()
                            }
                        } catch (_: Exception) {}
                    }
                } finally {
                    reader.close()
                    responseBody.close()
                }

                if (fullContent.isNotEmpty()) {
                    conv.messages.add(ChatMessage(role = "assistant", content = fullContent.toString()))
                    updateConversation(conv.copy())
                }
                _streamingContent.value = ""
            } catch (e: Exception) {
                _error.value = "流式请求失败: ${e.message}"
                _streamingContent.value = ""
            } finally {
                _isStreaming.value = false
            }
        }
    }

    private fun buildMessages(conv: Conversation): List<ChatMessage> {
        val messages = mutableListOf<ChatMessage>()
        if (conv.systemPrompt.isNotBlank()) {
            messages.add(ChatMessage(role = "system", content = conv.systemPrompt))
        }
        messages.addAll(conv.messages)
        return messages
    }

    private fun updateConversation(conv: Conversation) {
        _currentConversation.value = conv
        val list = _conversations.value.toMutableList()
        val idx = list.indexOfFirst { it.id == conv.id }
        if (idx >= 0) {
            list[idx] = conv
        }
        _conversations.value = list
        saveConversations()
    }

    fun clearError() {
        _error.value = null
    }
}