package com.sunyard.chsm.controller;

import com.sunyard.chsm.auth.AuthCode;
import com.sunyard.chsm.constant.AuthCodeConst;
import com.sunyard.chsm.constant.CryptoConst;
import com.sunyard.chsm.model.R;
import com.sunyard.chsm.param.HashReq;
import com.sunyard.chsm.param.HashResp;
import com.sunyard.chsm.pool.DeviceManager;
import com.sunyard.chsm.pool.TMKContext;
import com.sunyard.chsm.sdf.SdfApiService;
import com.sunyard.chsm.sdf.adapter.SdfApiAdapter;
import com.sunyard.chsm.sdf.context.AlgId;
import com.sunyard.chsm.sdf.model.EccPubKey;
import com.sunyard.chsm.utils.CodecUtils;
import org.springframework.util.Assert;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RestController;

import javax.annotation.Resource;
import javax.validation.Valid;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 杂凑运算类接口
 *
 * @author liulu
 * @since 2024/12/19
 */
@RestController
public class HashController {

    @Resource
    private SdfApiService sdfApiService;
    @Resource
    private DeviceManager deviceManager;

    /**
     * 计算Hash
     *
     * @param req
     * @return
     */
    @PostMapping("/hash")
    @AuthCode(AuthCodeConst.cal_hash)
    public R<HashResp> hash(@Valid @RequestBody HashReq req) {
        Assert.hasText(req.getPlainData(), "原文不能为空");
        byte[] bytes = CodecUtils.decodeBase64(req.getPlainData());
        Assert.isTrue(bytes.length <= 8192, "原文长度最大支持8192");
        byte[] pubkey = CodecUtils.decodeBase64(req.getPubKey());
        byte[] userId = CodecUtils.decodeBase64(req.getUserId());
        byte[] hash = sdfApiService.hash(bytes, pubkey, userId);
        HashResp resp = new HashResp();
        resp.setHash(CodecUtils.encodeBase64(hash));
        return R.data(resp);
    }

    private static final Map<String, TMKContext> HANDLE_MAP = new ConcurrentHashMap<>();


    /**
     * 多包HashInit
     *
     * @param req
     * @return
     */
    @PostMapping("/multi/package/hash/init")
    @AuthCode(AuthCodeConst.cal_hash)
    public R<HashResp> hashInit(@Valid @RequestBody HashReq req) {
        byte[] pubKey = CodecUtils.decodeBase64(req.getPubKey());
        byte[] userId = CodecUtils.decodeBase64(req.getUserId());
        TMKContext context = deviceManager.chooseOne();
        try {
            SdfApiAdapter sdf = context.getSdfApiAdapter();
            TMKContext newContext = new TMKContext();
            newContext.setDeviceHandle(context.getDeviceHandle());
            newContext.setSdfApiAdapter(sdf);
            String hs = sdf.openSession(context.getDeviceHandle());
            newContext.setSessionHandle(hs);
            EccPubKey eccPubKey = pubKey == null || pubKey.length == 0 ? null : EccPubKey.fromBytes(pubKey);
            byte[] finalUserId = Objects.nonNull(pubKey) && Objects.isNull(userId) ? CryptoConst.USER_ID : userId;
            sdf.hashInit(hs, AlgId.SGD_SM3, eccPubKey, finalUserId);
            String handle = UUID.randomUUID().toString();
            HANDLE_MAP.put(handle, newContext);

            HashResp resp = new HashResp();
            resp.setHandle(handle);
            return R.data(resp);
        } finally {
            Optional.ofNullable(context.getPool())
                    .ifPresent(it -> {
                        context.setPool(null);
                        it.returnObject(context);
                    });
        }
    }

    /**
     * 多包Hash update
     *
     * @param req
     * @return
     */
    @PostMapping("/multi/package/hash/update")
    @AuthCode(AuthCodeConst.cal_hash)
    public R<HashResp> hashUpdate(@Valid @RequestBody HashReq req) {
        Assert.hasText(req.getPlainData(), "原文不能为空");
        Assert.hasText(req.getHandle(), "句柄不能为空");
        byte[] bytes = CodecUtils.decodeBase64(req.getPlainData());
        Assert.isTrue(bytes.length <= 8192, "原文长度最大支持8192");
        TMKContext context = HANDLE_MAP.get(req.getHandle());
        Assert.notNull(context, "句柄引用不存在");
        SdfApiAdapter sdf = context.getSdfApiAdapter();
        sdf.hashUpdate(context.getSessionHandle(), bytes);
        HashResp resp = new HashResp();
        resp.setHandle(req.getHandle());
        return R.data(resp);
    }


    /**
     * 多包Hash finish
     *
     * @param req
     * @return
     */
    @PostMapping("/multi/package/hash/finish")
    @AuthCode(AuthCodeConst.cal_hash)
    public R<HashResp> hashFinish(@Valid @RequestBody HashReq req) {
        Assert.hasText(req.getHandle(), "句柄不能为空");
        TMKContext context = HANDLE_MAP.remove(req.getHandle());
        Assert.notNull(context, "句柄引用不存在");
        SdfApiAdapter sdf = context.getSdfApiAdapter();
        byte[] hash = sdf.hashFinish(context.getSessionHandle());
        sdf.closeSession(context.getSessionHandle());
        HashResp resp = new HashResp();
        resp.setHandle(req.getHandle());
        resp.setHash(CodecUtils.encodeBase64(hash));
        return R.data(resp);
    }


}
