package com.mistral.chat.api

import com.google.gson.JsonObject
import com.google.gson.Gson
import com.google.gson.JsonArray
import com.mistral.chat.data.Message
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.suspendCancellableCoroutine
import kotlinx.coroutines.withContext
import okhttp3.Call
import okhttp3.MediaType.Companion.toMediaType
import okhttp3.OkHttpClient
import okhttp3.Request
import okhttp3.RequestBody.Companion.toRequestBody
import okhttp3.Response
import java.util.concurrent.TimeUnit
import kotlin.coroutines.Continuation
import kotlin.coroutines.resume
import kotlin.coroutines.resumeWithException

class MistralClient(private val apiKey: String) {

    private var client = createNewClient()

    private fun createNewClient(): OkHttpClient {
        return OkHttpClient.Builder()
            .connectTimeout(15, TimeUnit.SECONDS)
            .readTimeout(25, TimeUnit.SECONDS)
            .writeTimeout(15, TimeUnit.SECONDS)
            .retryOnConnectionFailure(false)
            .build()
    }

    private val gson = Gson()
    private val jsonMediaType = "application/json".toMediaType()
    
    private var currentCall: Call? = null
    private var currentContinuation: Continuation<Result<Pair<String, String>>>? = null

    companion object {
        private const val BASE_URL = "https://api.mistral.ai/v1"
        
        private val SUPPORTED_MODELS = setOf(
            "mistral-small-latest",
            "mistral-medium-latest", 
            "mistral-large-latest",
            "mistral-small",
            "mistral-medium",
            "mistral-large",
            "codestral-latest",
            "codestral",
            "pixtral-large-latest",
            "pixtral-12b-2409"
        )
        
        val AVAILABLE_MODELS = listOf(
            "mistral-small-latest" to "Mistral Small",
            "mistral-medium-latest" to "Mistral Medium", 
            "mistral-large-latest" to "Mistral Large",
            "codestral-latest" to "Codestral",
            "pixtral-large-latest" to "Pixtral Large"
        )
    }

    suspend fun getModels(): Result<List<Pair<String, String>>> = withContext(Dispatchers.IO) {
        try {
            val request = Request.Builder()
                .url("$BASE_URL/models")
                .addHeader("Authorization", "Bearer $apiKey")
                .get()
                .build()

            client.newCall(request).execute().use { response ->
                if (!response.isSuccessful) {
                    return@withContext Result.failure(Exception("API error: ${response.code}"))
                }

                val responseBody = response.body?.string() ?: ""
                val responseJson = gson.fromJson(responseBody, JsonObject::class.java)
                
                val models = responseJson
                    .getAsJsonArray("data")
                    ?.mapNotNull { obj ->
                        try {
                            val jsonObj = obj.asJsonObject
                            val id = jsonObj.get("id")?.asString ?: return@mapNotNull null
                            if (id in SUPPORTED_MODELS) {
                                val displayName = id
                                    .replace("-latest", "")
                                    .replace("-12b-2409", "")
                                    .replace("-", " ")
                                    .split(" ")
                                    .joinToString(" ") { it.replaceFirstChar { c -> c.uppercase() } }
                                id to displayName
                            } else null
                        } catch (e: Exception) {
                            null
                        }
                    }
                    ?.distinctBy { it.first }
                    ?: emptyList()

                Result.success(models)
            }
        } catch (e: Exception) {
            if (e is java.io.IOException && !e.message.orEmpty().contains("cancel", ignoreCase = true)) {
                Result.failure(e)
            } else {
                Result.failure(Exception("Request cancelled"))
            }
        }
    }

    fun cancelRequest() {
        val call = currentCall
        val continuation = currentContinuation
        
        call?.cancel()
        currentCall = null
        
        continuation?.resume(Result.failure(Exception("Request cancelled")))
        currentContinuation = null
        
        client = createNewClient()
    }

    suspend fun chat(
        model: String,
        messages: List<Message>,
        onChunk: ((String) -> Unit)? = null
    ): Result<Pair<String, String>> = withContext(Dispatchers.IO) {
        try {
            val jsonObject = JsonObject()
            jsonObject.addProperty("model", model)
            jsonObject.addProperty("temperature", 0.7)
            jsonObject.addProperty("stream", false)

            val messagesArray = JsonArray()
            messages.forEach { msg ->
                val msgObj = JsonObject()
                msgObj.addProperty("role", if (msg.isUser) "user" else "assistant")
                msgObj.addProperty("content", msg.content)
                messagesArray.add(msgObj)
            }
            jsonObject.add("messages", messagesArray)

            val json = gson.toJson(jsonObject)
            val body = json.toRequestBody(jsonMediaType)

            val request = Request.Builder()
                .url("$BASE_URL/chat/completions")
                .addHeader("Authorization", "Bearer $apiKey")
                .addHeader("Content-Type", "application/json")
                .post(body)
                .build()

            currentCall = client.newCall(request)
            
            suspendCancellableCoroutine { continuation ->
                currentContinuation = continuation
                
                currentCall?.enqueue(object : okhttp3.Callback {
                    override fun onFailure(call: okhttp3.Call, e: java.io.IOException) {
                        val cont = currentContinuation
                        currentCall = null
                        currentContinuation = null
                        if (cont != null) {
                            if (call.isCanceled()) {
                                cont.resume(Result.failure(Exception("Request cancelled")))
                            } else {
                                cont.resume(Result.failure(e))
                            }
                        }
                    }

                    override fun onResponse(call: okhttp3.Call, response: Response) {
                        val cont = currentContinuation
                        
                        if (cont == null) {
                            response.close()
                            currentCall = null
                            return
                        }
                        
                        try {
                            if (!response.isSuccessful) {
                                val errorBody = response.body?.string() ?: "Unknown error"
                                currentCall = null
                                currentContinuation = null
                                cont.resume(Result.failure(Exception("API error: ${response.code} - $errorBody")))
                                return
                            }

                            val responseBody = response.body?.string() ?: ""
                            
                            if (onChunk != null) {
                                onChunk(responseBody)
                            }

                            val responseJson = try {
                                gson.fromJson(responseBody, JsonObject::class.java)
                            } catch (e: Exception) {
                                currentCall = null
                                currentContinuation = null
                                cont.resume(Result.failure(Exception("Invalid JSON response")))
                                return
                            }
                            
                            val choices = responseJson.getAsJsonArray("choices")
                            if (choices == null || choices.size() == 0) {
                                currentCall = null
                                currentContinuation = null
                                cont.resume(Result.failure(Exception("No response from API")))
                                return
                            }
                            
                            val firstChoice = choices.get(0)?.asJsonObject
                            if (firstChoice == null) {
                                currentCall = null
                                currentContinuation = null
                                cont.resume(Result.failure(Exception("Empty choice")))
                                return
                            }
                            
                            val message = firstChoice.getAsJsonObject("message")
                            val content = message?.get("content")?.asString ?: ""
                            
                            val usedModel = responseJson.get("model")?.asString ?: model
                            
                            currentCall = null
                            currentContinuation = null
                            cont.resume(Result.success(content to usedModel))
                        } catch (e: Exception) {
                            currentCall = null
                            currentContinuation = null
                            cont.resume(Result.failure(e))
                        }
                    }
                })
                
                continuation.invokeOnCancellation {
                    currentCall?.cancel()
                    currentCall = null
                    currentContinuation = null
                    client = createNewClient()
                }
            }
        } catch (e: Exception) {
            Result.failure(e)
        }
    }
}