package io.synor.zk

import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import kotlinx.serialization.*
import kotlinx.serialization.json.*
import java.net.URI
import java.net.http.HttpClient
import java.net.http.HttpRequest
import java.net.http.HttpResponse
import java.time.Duration
import java.util.*
import java.util.concurrent.atomic.AtomicBoolean

/**
 * Synor ZK SDK for Kotlin
 *
 * Zero-Knowledge proof systems for ZK-Rollups and privacy.
 */

// Enums
enum class ProofSystem(val value: String) {
    @SerialName("groth16") GROTH16("groth16"),
    @SerialName("plonk") PLONK("plonk"),
    @SerialName("stark") STARK("stark");

    companion object {
        fun fromString(value: String) = entries.find { it.value == value } ?: GROTH16
    }
}

enum class CircuitType(val value: String) {
    @SerialName("transfer") TRANSFER("transfer"),
    @SerialName("batch") BATCH("batch"),
    @SerialName("deposit") DEPOSIT("deposit"),
    @SerialName("withdraw") WITHDRAW("withdraw");

    companion object {
        fun fromString(value: String) = entries.find { it.value == value } ?: TRANSFER
    }
}

enum class RollupState { ACTIVE, PAUSED, FINALIZING, FINALIZED }
enum class ProofStatus { GENERATING, COMPLETED, FAILED, VERIFIED }

// Data classes
@Serializable
data class Proof(
    val id: String,
    val system: ProofSystem,
    @SerialName("circuit_type") val circuitType: CircuitType,
    val data: String,
    @SerialName("public_inputs") val publicInputs: List<String>,
    val size: Int,
    @SerialName("generation_time_ms") val generationTimeMs: Long,
    @SerialName("created_at") val createdAt: Long
)

@Serializable
data class VerificationKey(
    @SerialName("circuit_id") val circuitId: String,
    @SerialName("circuit_type") val circuitType: CircuitType,
    val system: ProofSystem,
    val data: String,
    val size: Int
)

@Serializable
data class ProvingKey(
    @SerialName("circuit_id") val circuitId: String,
    @SerialName("circuit_type") val circuitType: CircuitType,
    val system: ProofSystem,
    val data: String,
    val size: Int
)

@Serializable
data class CircuitConfig(
    val type: CircuitType,
    @SerialName("max_batch_size") val maxBatchSize: Int? = null,
    @SerialName("tree_depth") val treeDepth: Int? = null,
    @SerialName("verify_signatures") val verifySignatures: Boolean? = null,
    @SerialName("custom_constraints") val customConstraints: JsonObject? = null
)

@Serializable
data class AccountState(
    val address: String,
    val balance: String,
    val nonce: Long,
    @SerialName("pubkey_hash") val pubkeyHash: String
)

@Serializable
data class Transfer(
    val from: String,
    val to: String,
    val amount: String,
    val fee: String,
    val nonce: Long,
    val signature: String? = null
)

@Serializable
data class Deposit(
    @SerialName("l1_tx_hash") val l1TxHash: String,
    val recipient: String,
    val amount: String,
    val token: String? = null
)

@Serializable
data class Withdrawal(
    val sender: String,
    val recipient: String,
    val amount: String,
    val token: String? = null,
    val nonce: Long,
    val signature: String? = null
)

@Serializable
data class Batch(
    @SerialName("batch_number") val batchNumber: Long,
    @SerialName("prev_state_root") val prevStateRoot: String,
    @SerialName("new_state_root") val newStateRoot: String,
    val transfers: List<Transfer>,
    val deposits: List<Deposit>,
    val withdrawals: List<Withdrawal>,
    val proof: Proof? = null,
    val timestamp: Long
)

@Serializable
data class RollupStats(
    @SerialName("current_batch") val currentBatch: Long,
    @SerialName("total_transactions") val totalTransactions: Long,
    @SerialName("total_proofs") val totalProofs: Long,
    @SerialName("avg_proof_time_ms") val avgProofTimeMs: Long,
    @SerialName("state_root") val stateRoot: String,
    @SerialName("account_count") val accountCount: Long,
    val tvl: String
)

@Serializable
data class ProofRequest(
    @SerialName("circuit_type") val circuitType: CircuitType,
    @SerialName("public_inputs") val publicInputs: List<String>,
    @SerialName("private_inputs") val privateInputs: List<String>,
    val system: ProofSystem? = null
)

@Serializable
data class VerificationResult(
    val valid: Boolean,
    @SerialName("verification_time_ms") val verificationTimeMs: Long,
    val error: String? = null
)

@Serializable
data class MerkleProof(
    val leaf: String,
    val path: List<String>,
    val indices: List<Int>,
    val root: String
)

@Serializable
data class CeremonyContribution(
    @SerialName("contributor_id") val contributorId: String,
    val hash: String,
    val timestamp: Long,
    val verified: Boolean
)

@Serializable
data class TrustedSetup(
    val id: String,
    @SerialName("circuit_type") val circuitType: CircuitType,
    val system: ProofSystem,
    @SerialName("contribution_count") val contributionCount: Int,
    @SerialName("latest_hash") val latestHash: String,
    val complete: Boolean,
    val contributions: List<CeremonyContribution>
)

data class ZkConfig(
    val apiKey: String,
    val endpoint: String = "https://zk.synor.io/v1",
    val wsEndpoint: String = "wss://zk.synor.io/v1/ws",
    val timeout: Long = 60000,
    val retries: Int = 3,
    val defaultProofSystem: ProofSystem = ProofSystem.GROTH16,
    val debug: Boolean = false
)

class ZkException(
    message: String,
    val code: String? = null,
    val status: Int? = null
) : RuntimeException(message)

data class ProofSystemInfo(
    val name: String,
    val proofSize: Int,
    val verificationTimeMs: Int,
    val trustedSetup: Boolean,
    val universal: Boolean
)

fun getProofSystemInfo(system: ProofSystem): ProofSystemInfo = when (system) {
    ProofSystem.GROTH16 -> ProofSystemInfo("Groth16", 192, 10, true, false)
    ProofSystem.PLONK -> ProofSystemInfo("PLONK", 512, 15, true, true)
    ProofSystem.STARK -> ProofSystemInfo("STARK", 50000, 30, false, true)
}

const val MAX_BATCH_SIZE = 1000
const val STATE_TREE_DEPTH = 32
const val PROOF_EXPIRY_BLOCKS = 100

/**
 * Main ZK SDK client
 */
class SynorZk(private val config: ZkConfig) {
    private val httpClient = HttpClient.newBuilder()
        .connectTimeout(Duration.ofMillis(config.timeout))
        .build()
    private val json = Json { ignoreUnknownKeys = true }
    private val closed = AtomicBoolean(false)

    val proofs = ProofsClient(this)
    val circuits = CircuitsClient(this)
    val rollup = RollupClient(this)
    val state = StateClient(this)
    val ceremony = CeremonyClient(this)

    val defaultProofSystem: ProofSystem get() = config.defaultProofSystem

    suspend fun healthCheck(): Boolean = try {
        val result = get<JsonObject>("/health")
        result["status"]?.jsonPrimitive?.content == "healthy"
    } catch (e: Exception) {
        false
    }

    suspend fun getInfo(): JsonObject = get("/info")

    fun close() { closed.set(true) }
    val isClosed: Boolean get() = closed.get()

    internal suspend inline fun <reified T> get(path: String): T = request("GET", path, null)
    internal suspend inline fun <reified T> post(path: String, body: Any?): T = request("POST", path, body)

    internal suspend inline fun <reified T> request(method: String, path: String, body: Any?): T {
        if (closed.get()) throw ZkException("Client has been closed", "CLIENT_CLOSED")

        return withContext(Dispatchers.IO) {
            val url = "${config.endpoint}$path"
            val requestBuilder = HttpRequest.newBuilder()
                .uri(URI.create(url))
                .header("Content-Type", "application/json")
                .header("Authorization", "Bearer ${config.apiKey}")
                .header("X-SDK-Version", "kotlin/0.1.0")
                .timeout(Duration.ofMillis(config.timeout))

            when (method) {
                "GET" -> requestBuilder.GET()
                "POST" -> {
                    val jsonBody = if (body != null) json.encodeToString(body as? JsonElement ?: buildJsonObject {
                        (body as? Map<*, *>)?.forEach { (k, v) ->
                            put(k.toString(), when (v) {
                                is String -> JsonPrimitive(v)
                                is Number -> JsonPrimitive(v)
                                is Boolean -> JsonPrimitive(v)
                                is List<*> -> JsonArray(v.map { JsonPrimitive(it.toString()) })
                                else -> JsonPrimitive(v.toString())
                            })
                        }
                    }) else "{}"
                    requestBuilder.POST(HttpRequest.BodyPublishers.ofString(jsonBody))
                }
                "DELETE" -> requestBuilder.DELETE()
            }

            val response = httpClient.send(requestBuilder.build(), HttpResponse.BodyHandlers.ofString())

            if (response.statusCode() >= 400) {
                val error = json.parseToJsonElement(response.body()).jsonObject
                throw ZkException(
                    error["message"]?.jsonPrimitive?.content ?: "HTTP ${response.statusCode()}",
                    error["code"]?.jsonPrimitive?.content,
                    response.statusCode()
                )
            }

            json.decodeFromString(response.body())
        }
    }
}

class ProofsClient(private val zk: SynorZk) {
    suspend fun generate(request: ProofRequest): Proof {
        val body = buildJsonObject {
            put("circuit_type", request.circuitType.value)
            put("public_inputs", JsonArray(request.publicInputs.map { JsonPrimitive(it) }))
            put("private_inputs", JsonArray(request.privateInputs.map { JsonPrimitive(it) }))
            put("system", (request.system ?: zk.defaultProofSystem).value)
        }
        return zk.post("/proofs/generate", body)
    }

    suspend fun verify(proof: Proof): VerificationResult {
        val body = buildJsonObject {
            put("proof_id", proof.id)
            put("data", proof.data)
            put("public_inputs", JsonArray(proof.publicInputs.map { JsonPrimitive(it) }))
            put("system", proof.system.value)
            put("circuit_type", proof.circuitType.value)
        }
        return zk.post("/proofs/verify", body)
    }

    suspend fun get(proofId: String): Proof = zk.get("/proofs/$proofId")

    suspend fun list(limit: Int? = null, offset: Int? = null): List<Proof> {
        val params = listOfNotNull(
            limit?.let { "limit=$it" },
            offset?.let { "offset=$it" }
        ).joinToString("&")
        val path = if (params.isNotEmpty()) "/proofs?$params" else "/proofs"
        val result: JsonObject = zk.get(path)
        return result["proofs"]?.jsonArray?.map { Json.decodeFromJsonElement(it) } ?: emptyList()
    }

    suspend fun serialize(proof: Proof): ByteArray {
        val body = buildJsonObject { put("proof_id", proof.id) }
        val result: JsonObject = zk.post("/proofs/serialize", body)
        return Base64.getDecoder().decode(result["data"]?.jsonPrimitive?.content)
    }

    suspend fun deserialize(data: ByteArray, system: ProofSystem): Proof {
        val body = buildJsonObject {
            put("data", Base64.getEncoder().encodeToString(data))
            put("system", system.value)
        }
        return zk.post("/proofs/deserialize", body)
    }
}

class CircuitsClient(private val zk: SynorZk) {
    suspend fun getVerificationKey(circuitType: CircuitType, system: ProofSystem? = null): VerificationKey {
        val sys = system ?: zk.defaultProofSystem
        return zk.get("/circuits/${circuitType.value}/vk?system=${sys.value}")
    }

    suspend fun getProvingKey(circuitType: CircuitType, system: ProofSystem? = null): ProvingKey {
        val sys = system ?: zk.defaultProofSystem
        return zk.get("/circuits/${circuitType.value}/pk?system=${sys.value}")
    }

    suspend fun list(): List<CircuitConfig> {
        val result: JsonObject = zk.get("/circuits")
        return result["circuits"]?.jsonArray?.map { Json.decodeFromJsonElement(it) } ?: emptyList()
    }

    suspend fun getConfig(circuitType: CircuitType): CircuitConfig =
        zk.get("/circuits/${circuitType.value}/config")

    suspend fun compile(config: CircuitConfig): String {
        val result: JsonObject = zk.post("/circuits/compile", Json.encodeToJsonElement(config))
        return result["circuit_id"]?.jsonPrimitive?.content ?: ""
    }
}

class RollupClient(private val zk: SynorZk) {
    suspend fun getStats(): RollupStats = zk.get("/rollup/stats")
    suspend fun getCurrentBatch(): Batch = zk.get("/rollup/batch/current")
    suspend fun getBatch(batchNumber: Long): Batch = zk.get("/rollup/batch/$batchNumber")

    suspend fun submitTransfer(transfer: Transfer): String {
        val result: JsonObject = zk.post("/rollup/transfer", Json.encodeToJsonElement(transfer))
        return result["tx_id"]?.jsonPrimitive?.content ?: ""
    }

    suspend fun submitDeposit(deposit: Deposit): String {
        val result: JsonObject = zk.post("/rollup/deposit", Json.encodeToJsonElement(deposit))
        return result["tx_id"]?.jsonPrimitive?.content ?: ""
    }

    suspend fun submitWithdrawal(withdrawal: Withdrawal): String {
        val result: JsonObject = zk.post("/rollup/withdraw", Json.encodeToJsonElement(withdrawal))
        return result["tx_id"]?.jsonPrimitive?.content ?: ""
    }

    suspend fun finalizeBatch(): Batch = zk.post("/rollup/batch/finalize", buildJsonObject {})

    suspend fun getPendingTransactions(): List<Transfer> {
        val result: JsonObject = zk.get("/rollup/pending")
        return result["transactions"]?.jsonArray?.map { Json.decodeFromJsonElement(it) } ?: emptyList()
    }
}

class StateClient(private val zk: SynorZk) {
    suspend fun getRoot(): String {
        val result: JsonObject = zk.get("/state/root")
        return result["root"]?.jsonPrimitive?.content ?: ""
    }

    suspend fun getAccount(address: String): AccountState = zk.get("/state/account/$address")
    suspend fun getMerkleProof(address: String): MerkleProof = zk.get("/state/proof/$address")

    suspend fun verifyMerkleProof(proof: MerkleProof): Boolean {
        val result: JsonObject = zk.post("/state/proof/verify", Json.encodeToJsonElement(proof))
        return result["valid"]?.jsonPrimitive?.boolean ?: false
    }

    suspend fun getStateAtBatch(batchNumber: Long): JsonObject = zk.get("/state/batch/$batchNumber")
}

class CeremonyClient(private val zk: SynorZk) {
    suspend fun getStatus(circuitType: CircuitType, system: ProofSystem? = null): TrustedSetup {
        val sys = system ?: zk.defaultProofSystem
        return zk.get("/ceremony/${circuitType.value}?system=${sys.value}")
    }

    suspend fun contribute(circuitType: CircuitType, entropy: ByteArray, system: ProofSystem? = null): CeremonyContribution {
        val sys = system ?: zk.defaultProofSystem
        val body = buildJsonObject {
            put("circuit_type", circuitType.value)
            put("entropy", Base64.getEncoder().encodeToString(entropy))
            put("system", sys.value)
        }
        return zk.post("/ceremony/contribute", body)
    }

    suspend fun verifyContribution(circuitType: CircuitType, contributionHash: String): Boolean {
        val body = buildJsonObject {
            put("circuit_type", circuitType.value)
            put("contribution_hash", contributionHash)
        }
        val result: JsonObject = zk.post("/ceremony/verify", body)
        return result["valid"]?.jsonPrimitive?.boolean ?: false
    }

    suspend fun listContributions(circuitType: CircuitType): List<CeremonyContribution> {
        val result: JsonObject = zk.get("/ceremony/${circuitType.value}/contributions")
        return result["contributions"]?.jsonArray?.map { Json.decodeFromJsonElement(it) } ?: emptyList()
    }
}
