package com.qmth.ops.biz.ai.client;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.qmth.boot.core.rateLimit.service.RateLimiter;
import com.qmth.boot.core.rateLimit.service.impl.MemoryRateLimiter;
import com.qmth.ops.biz.ai.exception.OcrRateLimitExceeded;
import okhttp3.*;

import java.io.IOException;

/**
 * 大模型chat类接口基础实现
 */
public abstract class OcrApiClient {

    private OkHttpClient client;

    private OcrApiConfig config;

    private ObjectMapper mapper;

    private RateLimiter queryRateLimiter;

    public OcrApiClient(OcrApiConfig config) {
        this.client = new OkHttpClient.Builder().connectionPool(new ConnectionPool()).build();
        this.mapper = new ObjectMapper();
        this.config = config;
        if (config.getQps() > 0) {
            this.queryRateLimiter = new MemoryRateLimiter(config.getQps(), 1000);
        }
    }

    protected OcrApiConfig getConfig() {
        return config;
    }

    protected abstract String buildUrl() throws Exception;

    protected abstract String buildResult(byte[] data, ObjectMapper mapper) throws IOException;

    protected abstract String handleError(byte[] data, int statusCode, ObjectMapper mapper);

    public String call(byte[] image) throws Exception {
        if (queryRateLimiter != null && !queryRateLimiter.acquire()) {
            throw new OcrRateLimitExceeded(config.getQps());
        }
        RequestBody body = RequestBody.create(MediaType.parse("application/octet-stream"), image);
        Request httpRequest = new Request.Builder().url(buildUrl()).post(body).build();
        Response response = client.newCall(httpRequest).execute();
        byte[] data = response.body() != null ? response.body().bytes() : null;
        if (response.isSuccessful()) {
            return buildResult(data, mapper);
        } else {
            return handleError(data, response.code(), mapper);
        }
    }

}