package com.yaoai.pipeline.service.impl;

import cn.dev33.satoken.stp.StpUtil;
import com.yaoai.ai.providers.service.SeedreamService;
import com.yaoai.billing.dto.BillingChargeRequest;
import com.yaoai.billing.service.BillingService;
import com.yaoai.common.exception.BizException;
import com.yaoai.common.exception.ErrorCode;
import com.yaoai.pipeline.service.ImageGenPipelineService;
import com.yaoai.storage.service.TosService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.io.ByteArrayInputStream;
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.Map;

@Slf4j
@Service
@RequiredArgsConstructor
public class ImageGenPipelineServiceImpl implements ImageGenPipelineService {

    private final SeedreamService seedreamService;
    private final TosService tosService;
    private final BillingService billingService;
    private final HttpClient httpClient = HttpClient.newBuilder()
            .connectTimeout(Duration.ofSeconds(10))
            .build();

    @Override
    public String generateAndStore(Long tenantId, Long projectId, String prompt) {
        return generateAndStore(tenantId, StpUtil.getLoginIdAsLong(), projectId, prompt);
    }

    @Override
    public String generateAndStore(Long tenantId, Long userId, Long projectId, String prompt) {
        return generateAndStore(tenantId, userId, projectId, "scene", prompt);
    }

    @Override
    public String generateAndStore(Long tenantId, Long userId, Long projectId, String assetType, String prompt) {
        billingService.checkBalance(BillingChargeRequest.builder()
                .tenantId(tenantId)
                .userId(userId)
                .projectId(projectId)
                .operation("image_generate")
                .modality("IMAGE")
                .modelProvider(seedreamService.getProvider())
                .modelId(seedreamService.getModelId())
                .modelName(seedreamService.getModelName())
                .unitCount(1)
                .meta(Map.of("size", "2048x2048"))
                .build());

        // 1. 调用 Seedream 生成图片，获取临时 URL
        String imageUrl = seedreamService.generateImage(prompt);
        log.info("Seedream image URL: {}", imageUrl);

        // 2. 下载图片字节
        byte[] imageBytes = downloadImage(imageUrl);

        // 3. 上传到 TOS
        String normalizedAssetType = assetType != null && !assetType.isBlank() ? assetType : "image";
        String key = TosService.buildKey(tenantId, projectId, normalizedAssetType, normalizedAssetType + ".jpg");
        tosService.upload(key, new ByteArrayInputStream(imageBytes), imageBytes.length, "image/jpeg");
        log.info("Generated image stored: key={}", key);
        billingService.charge(BillingChargeRequest.builder()
                .tenantId(tenantId)
                .userId(userId)
                .projectId(projectId)
                .operation("image_generate")
                .modality("IMAGE")
                .modelProvider(seedreamService.getProvider())
                .modelId(seedreamService.getModelId())
                .modelName(seedreamService.getModelName())
                .unitCount(1)
                .meta(Map.of("size", "2048x2048"))
                .refId(key)
                .build());
        return key;
    }

    private byte[] downloadImage(String url) {
        try {
            HttpRequest request = HttpRequest.newBuilder()
                    .uri(URI.create(url))
                    .timeout(Duration.ofMinutes(2))
                    .GET()
                    .build();
            HttpResponse<byte[]> response = httpClient.send(request, HttpResponse.BodyHandlers.ofByteArray());
            if (response.statusCode() != 200) {
                throw new BizException(ErrorCode.INTERNAL_ERROR,
                        "下载图片 HTTP " + response.statusCode());
            }
            return response.body();
        } catch (BizException e) {
            throw e;
        } catch (Exception e) {
            log.error("Failed to download image from {}", url, e);
            throw new BizException(ErrorCode.INTERNAL_ERROR, "下载生成图片失败: " + e.getMessage());
        }
    }
}
