package io.synor.compute

import io.ktor.client.*
import io.ktor.client.call.*
import io.ktor.client.engine.cio.*
import io.ktor.client.plugins.*
import io.ktor.client.plugins.contentnegotiation.*
import io.ktor.client.request.*
import io.ktor.client.statement.*
import io.ktor.http.*
import io.ktor.serialization.kotlinx.json.*
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.flow
import kotlinx.serialization.json.*
import java.io.Closeable

/**
 * Synor Compute SDK - Kotlin Client
 *
 * Access distributed heterogeneous compute resources (CPU, GPU, TPU, NPU, LPU, FPGA, DSP)
 * for AI/ML workloads at 90% cost reduction compared to traditional cloud.
 *
 * ```kotlin
 * // Create client
 * val client = SynorCompute("your-api-key")
 *
 * // Matrix multiplication on GPU
 * val a = Tensor.rand(512, 512)
 * val b = Tensor.rand(512, 512)
 * val result = client.matmul(a, b, MatMulOptions(
 *     processor = ProcessorType.GPU,
 *     precision = Precision.FP16
 * ))
 *
 * if (result.isSuccess) {
 *     println("Result shape: ${result.result?.shape?.contentToString()}")
 *     println("Time: ${result.executionTimeMs}ms")
 * }
 *
 * // LLM inference
 * val response = client.inference("llama-3-70b", "Explain quantum computing")
 * println(response.result)
 *
 * // Streaming inference
 * client.inferenceStream("llama-3-70b", "Write a poem about AI").collect { token ->
 *     print(token)
 * }
 *
 * // Clean up
 * client.close()
 * ```
 */
class SynorCompute(
    apiKey: String,
    baseUrl: String = "https://api.synor.io/compute/v1",
    defaultProcessor: ProcessorType = ProcessorType.AUTO,
    defaultPrecision: Precision = Precision.FP32,
    defaultPriority: Priority = Priority.NORMAL,
    timeoutMs: Long = 30000,
    debug: Boolean = false
) : Closeable {

    constructor(config: SynorConfig) : this(
        apiKey = config.apiKey,
        baseUrl = config.baseUrl,
        defaultProcessor = config.defaultProcessor,
        defaultPrecision = config.defaultPrecision,
        defaultPriority = config.defaultPriority,
        timeoutMs = config.timeoutMs,
        debug = config.debug
    )

    private val config = SynorConfig(
        apiKey, baseUrl, defaultProcessor, defaultPrecision, defaultPriority, timeoutMs, debug
    )

    private val json = Json {
        ignoreUnknownKeys = true
        isLenient = true
    }

    private val httpClient = HttpClient(CIO) {
        install(ContentNegotiation) {
            json(this@SynorCompute.json)
        }
        install(HttpTimeout) {
            requestTimeoutMillis = timeoutMs
            connectTimeoutMillis = timeoutMs
            socketTimeoutMillis = timeoutMs
        }
        defaultRequest {
            header("Authorization", "Bearer ${config.apiKey}")
            header("X-SDK-Version", "kotlin/$VERSION")
            contentType(ContentType.Application.Json)
        }
    }

    @Volatile
    private var closed = false

    // ==================== Matrix Operations ====================

    /** Perform matrix multiplication */
    suspend fun matmul(
        a: Tensor,
        b: Tensor,
        options: MatMulOptions = MatMulOptions()
    ): JobResult<Tensor> {
        checkNotClosed()

        val body = buildJsonObject {
            put("operation", "matmul")
            put("a", tensorToJson(a))
            put("b", tensorToJson(b))
            put("precision", options.precision.value)
            put("processor", options.processor.value)
            put("priority", options.priority.value)
        }

        return post("/compute", body)
    }

    /** Perform 2D convolution */
    suspend fun conv2d(
        input: Tensor,
        kernel: Tensor,
        options: Conv2dOptions = Conv2dOptions()
    ): JobResult<Tensor> {
        checkNotClosed()

        val body = buildJsonObject {
            put("operation", "conv2d")
            put("input", tensorToJson(input))
            put("kernel", tensorToJson(kernel))
            putJsonArray("stride") {
                add(options.stride.first)
                add(options.stride.second)
            }
            putJsonArray("padding") {
                add(options.padding.first)
                add(options.padding.second)
            }
            put("precision", options.precision.value)
        }

        return post("/compute", body)
    }

    /** Perform attention computation */
    suspend fun attention(
        query: Tensor,
        key: Tensor,
        value: Tensor,
        options: AttentionOptions = AttentionOptions()
    ): JobResult<Tensor> {
        checkNotClosed()

        val body = buildJsonObject {
            put("operation", "attention")
            put("query", tensorToJson(query))
            put("key", tensorToJson(key))
            put("value", tensorToJson(value))
            put("num_heads", options.numHeads)
            put("flash", options.flash)
            put("precision", options.precision.value)
        }

        return post("/compute", body)
    }

    // ==================== LLM Inference ====================

    /** Run inference on a model */
    suspend fun inference(
        model: String,
        prompt: String,
        options: InferenceOptions = InferenceOptions()
    ): JobResult<String> {
        checkNotClosed()

        val body = buildJsonObject {
            put("operation", "inference")
            put("model", model)
            put("prompt", prompt)
            put("max_tokens", options.maxTokens)
            put("temperature", options.temperature)
            put("top_p", options.topP)
            put("top_k", options.topK)
            options.processor?.let { put("processor", it.value) }
        }

        return postString("/inference", body)
    }

    /** Run streaming inference */
    fun inferenceStream(
        model: String,
        prompt: String,
        options: InferenceOptions = InferenceOptions()
    ): Flow<String> = flow {
        checkNotClosed()

        val body = buildJsonObject {
            put("operation", "inference")
            put("model", model)
            put("prompt", prompt)
            put("max_tokens", options.maxTokens)
            put("temperature", options.temperature)
            put("stream", true)
        }

        httpClient.preparePost("${config.baseUrl}/inference/stream") {
            setBody(body)
        }.execute { response ->
            val channel = response.bodyAsChannel()
            val buffer = StringBuilder()

            while (!channel.isClosedForRead) {
                val line = channel.readUTF8Line() ?: break
                if (line.startsWith("data: ")) {
                    val data = line.removePrefix("data: ")
                    if (data == "[DONE]") break
                    try {
                        val jsonData = json.parseToJsonElement(data).jsonObject
                        jsonData["token"]?.jsonPrimitive?.contentOrNull?.let { emit(it) }
                    } catch (e: Exception) {
                        // Skip malformed JSON
                    }
                }
            }
        }
    }

    // ==================== Model Registry ====================

    /** List available models */
    suspend fun listModels(category: ModelCategory? = null): List<ModelInfo> {
        checkNotClosed()
        val url = category?.let { "/models?category=${it.value}" } ?: "/models"
        val response: JsonObject = httpClient.get("${config.baseUrl}$url").body()
        return response["models"]?.jsonArray?.map {
            json.decodeFromJsonElement(it)
        } ?: emptyList()
    }

    /** Get model by ID */
    suspend fun getModel(modelId: String): ModelInfo {
        checkNotClosed()
        return httpClient.get("${config.baseUrl}/models/$modelId").body()
    }

    /** Search models */
    suspend fun searchModels(query: String): List<ModelInfo> {
        checkNotClosed()
        val response: JsonObject = httpClient.get("${config.baseUrl}/models/search?q=$query").body()
        return response["models"]?.jsonArray?.map {
            json.decodeFromJsonElement(it)
        } ?: emptyList()
    }

    // ==================== Pricing & Usage ====================

    /** Get current pricing information */
    suspend fun getPricing(): List<PricingInfo> {
        checkNotClosed()
        val response: JsonObject = httpClient.get("${config.baseUrl}/pricing").body()
        return response["pricing"]?.jsonArray?.map {
            json.decodeFromJsonElement(it)
        } ?: emptyList()
    }

    /** Get usage statistics */
    suspend fun getUsage(): UsageStats {
        checkNotClosed()
        return httpClient.get("${config.baseUrl}/usage").body()
    }

    // ==================== Health Check ====================

    /** Check service health */
    suspend fun healthCheck(): Boolean {
        return try {
            val response: JsonObject = httpClient.get("${config.baseUrl}/health").body()
            response["status"]?.jsonPrimitive?.content == "healthy"
        } catch (e: Exception) {
            false
        }
    }

    // ==================== Internal Methods ====================

    private suspend inline fun <reified T> post(path: String, body: JsonObject): JobResult<T> {
        val response: JsonObject = httpClient.post("${config.baseUrl}$path") {
            setBody(body)
        }.body()
        return parseJobResult(response)
    }

    private suspend fun postString(path: String, body: JsonObject): JobResult<String> {
        val response: JsonObject = httpClient.post("${config.baseUrl}$path") {
            setBody(body)
        }.body()
        return JobResult(
            jobId = response["job_id"]?.jsonPrimitive?.contentOrNull,
            status = response["status"]?.jsonPrimitive?.contentOrNull?.let {
                JobStatus.entries.find { s -> s.value == it }
            } ?: JobStatus.PENDING,
            result = response["result"]?.jsonPrimitive?.contentOrNull,
            error = response["error"]?.jsonPrimitive?.contentOrNull,
            executionTimeMs = response["execution_time_ms"]?.jsonPrimitive?.longOrNull,
            processor = response["processor"]?.jsonPrimitive?.contentOrNull?.let {
                ProcessorType.fromValue(it)
            },
            cost = response["cost"]?.jsonPrimitive?.doubleOrNull
        )
    }

    private inline fun <reified T> parseJobResult(response: JsonObject): JobResult<T> {
        val result = response["result"]?.let { json.decodeFromJsonElement<T>(it) }
        return JobResult(
            jobId = response["job_id"]?.jsonPrimitive?.contentOrNull,
            status = response["status"]?.jsonPrimitive?.contentOrNull?.let {
                JobStatus.entries.find { s -> s.value == it }
            } ?: JobStatus.PENDING,
            result = result,
            error = response["error"]?.jsonPrimitive?.contentOrNull,
            executionTimeMs = response["execution_time_ms"]?.jsonPrimitive?.longOrNull,
            processor = response["processor"]?.jsonPrimitive?.contentOrNull?.let {
                ProcessorType.fromValue(it)
            },
            cost = response["cost"]?.jsonPrimitive?.doubleOrNull
        )
    }

    private fun tensorToJson(tensor: Tensor): JsonObject = buildJsonObject {
        putJsonArray("shape") { tensor.shape.forEach { add(it) } }
        putJsonArray("data") { tensor.data.forEach { add(it) } }
        put("dtype", tensor.dtype.value)
    }

    private fun checkNotClosed() {
        check(!closed) { "Client has been closed" }
    }

    override fun close() {
        closed = true
        httpClient.close()
    }

    companion object {
        const val VERSION = "0.1.0"
    }
}
