package com.yaoai.ai.providers.service.impl;

import com.yaoai.ai.providers.config.ArkProperties;
import com.yaoai.ai.providers.service.SeedreamService;
import com.yaoai.common.exception.BizException;
import com.yaoai.common.exception.ErrorCode;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.client.SimpleClientHttpRequestFactory;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestClient;

import java.time.Duration;
import java.util.List;
import java.util.Map;

@Slf4j
@Service
public class SeedreamServiceImpl implements SeedreamService {

    private final ArkProperties properties;
    private final RestClient restClient;

    public SeedreamServiceImpl(ArkProperties properties) {
        this.properties = properties;
        SimpleClientHttpRequestFactory rf = new SimpleClientHttpRequestFactory();
        rf.setConnectTimeout(Duration.ofSeconds(10));
        rf.setReadTimeout(Duration.ofMinutes(2));
        this.restClient = RestClient.builder()
                .baseUrl(properties.getBaseUrl())
                .defaultHeader("Authorization", "Bearer " + properties.getApiKey())
                .defaultHeader("Content-Type", "application/json")
                .requestFactory(rf)
                .build();
    }

    @Override
    public String generateImage(String prompt) {
        return generateImage(prompt, "2048x2048");
    }

    @Override
    public String generateImage(String prompt, String size) {
        Map<String, Object> body = Map.of(
                "model", properties.getImageModel(),
                "prompt", prompt,
                "n", 1,
                "size", size,
                "response_format", "url"
        );

        log.info("Seedream generate image: model={}", properties.getImageModel());
        try {
            ImageResponse resp = restClient.post()
                    .uri("/images/generations")
                    .body(body)
                    .retrieve()
                    .body(ImageResponse.class);

            if (resp == null || resp.getData() == null || resp.getData().isEmpty()) {
                throw new BizException(ErrorCode.INTERNAL_ERROR, "Seedream 返回空结果");
            }
            String url = resp.getData().get(0).getUrl();
            log.info("Seedream image generated: url={}", url);
            return url;
        } catch (BizException e) {
            throw e;
        } catch (Exception e) {
            log.error("Seedream generate failed", e);
            throw new BizException(ErrorCode.INTERNAL_ERROR, "图片生成失败: " + e.getMessage());
        }
    }

    @Override
    public String getModelId() {
        return properties.getImageModel();
    }

    @Override
    public String getModelName() {
        return "Seedream 文生图";
    }

    // ---- internal response models ----

    @Data
    private static class ImageResponse {
        private Long created;
        private List<ImageData> data;
    }

    @Data
    private static class ImageData {
        private String url;
        private String b64Json;
    }
}
