proxy-playbook/code/upstream/WorkBuddyChatProxy.kt

package com.rainy.token.data.proxy

import com.rainy.token.data.repository.CredentialRepository
import com.rainy.token.domain.model.Credential
import com.rainy.token.domain.service.ServiceType
import java.util.UUID
import javax.inject.Inject
import javax.inject.Singleton
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import okhttp3.MediaType.Companion.toMediaType
import okhttp3.OkHttpClient
import okhttp3.Request
import okhttp3.RequestBody.Companion.toRequestBody

/**
 * WorkBuddy / CodeBuddy 上游 Chat 代理（国内版）。
 *
 * 上游协议为 OpenAI 兼容的 POST {base}/v2/chat/completions，Bearer accessToken。
 * 这里保留客户端原始 body（M1b 最小实现），并补齐桌面端请求头；完整归一化在后续迭代。
 */
@Singleton
class WorkBuddyChatProxy @Inject constructor(
    private val okHttpClient: OkHttpClient,
    private val credentialRepository: CredentialRepository,
) {

    suspend fun forwardChat(
        requestBody: String,
        accountId: String? = null,
        forcedRegion: ProviderRegion? = null,
    ): ProxyUpstreamResponse? =
        withContext(Dispatchers.IO) {
            val credential = credentialRepository.get(ServiceType.WORKBUDDY, accountId)
                ?: return@withContext null
            if (credential !is Credential.WorkBuddyCredential) return@withContext null
            val accessToken = credential.accessToken.trim().takeIf { it.isNotBlank() }
                ?: return@withContext null

            val effectiveRegion = forcedRegion ?: runCatching {
                ProviderRegion.valueOf(credential.region.uppercase())
            }.getOrNull()
            val base = if (effectiveRegion == ProviderRegion.INTL) {
                "https://www.workbuddy.ai"
            } else {
                "https://copilot.tencent.com"
            }
            val requestId = UUID.randomUUID().toString()
            val builder = Request.Builder()
                .url(base + "/v2/chat/completions")
                .addHeader("Content-Type", "application/json")
                .addHeader("Accept", "text/event-stream")
                .addHeader("Authorization", "Bearer $accessToken")
                .addHeader("User-Agent", "RainyToken/1.0")
                .addHeader("X-IDE-Type", "VSCode")
                .addHeader("X-IDE-Name", "CodeBuddy")
                .addHeader("X-IDE-Version", "3.0.0")
                .addHeader("X-Product", "CodeBuddy")
                .addHeader("X-Agent-Intent", "craft")
                .addHeader("X-Request-ID", requestId)
                .addHeader("X-Conv-Request-ID", requestId)
                .post(requestBody.toRequestBody("application/json".toMediaType()))

            val response = try {
                okHttpClient.newCall(builder.build()).execute()
            } catch (e: java.io.IOException) {
                throw e
            }
            val bytes = try { response.body?.bytes() ?: ByteArray(0) } catch (_: Throwable) { ByteArray(0) }
            val contentType = response.header("Content-Type") ?: "application/json"
            val status = response.code
            response.close()
            ProxyUpstreamResponse(status, contentType, bytes)
        }
    suspend fun openStreamingChat(
        requestBody: String,
        accountId: String? = null,
        forcedRegion: ProviderRegion? = null,
    ): ProxyUpstreamStream? = withContext(Dispatchers.IO) {
        val credential = credentialRepository.get(ServiceType.WORKBUDDY, accountId)
            ?: return@withContext null
        if (credential !is Credential.WorkBuddyCredential) return@withContext null
        val accessToken = credential.accessToken.trim().takeIf { it.isNotBlank() }
            ?: return@withContext null

        val effectiveRegion = forcedRegion ?: runCatching {
            ProviderRegion.valueOf(credential.region.uppercase())
        }.getOrNull()
        val base = if (effectiveRegion == ProviderRegion.INTL) {
            "https://www.workbuddy.ai"
        } else {
            "https://copilot.tencent.com"
        }
        val requestId = UUID.randomUUID().toString()
        val builder = Request.Builder()
            .url(base + "/v2/chat/completions")
            .addHeader("Content-Type", "application/json")
            .addHeader("Accept", "text/event-stream")
            .addHeader("Authorization", "Bearer $accessToken")
            .addHeader("User-Agent", "RainyToken/1.0")
            .addHeader("X-IDE-Type", "VSCode")
            .addHeader("X-IDE-Name", "CodeBuddy")
            .addHeader("X-IDE-Version", "3.0.0")
            .addHeader("X-Product", "CodeBuddy")
            .addHeader("X-Agent-Intent", "craft")
            .addHeader("X-Request-ID", requestId)
            .addHeader("X-Conv-Request-ID", requestId)
            .post(requestBody.toRequestBody("application/json".toMediaType()))

        val response = try {
            okHttpClient.newCall(builder.build()).execute()
        } catch (e: java.io.IOException) {
            throw e
        }
        val input = response.body?.byteStream() ?: run {
            response.close()
            return@withContext null
        }
        ProxyUpstreamStream(
            status = response.code,
            contentType = response.header("Content-Type") ?: "text/event-stream",
            input = input,
            close = { response.close() },
        )
    }

}