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

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.yaoai.ai.providers.config.ArkProperties;
import com.yaoai.ai.providers.config.ModelCapability;
import com.yaoai.ai.providers.config.PlatformModelConfigService;
import com.yaoai.ai.providers.config.ProviderRestClientFactory;
import com.yaoai.ai.providers.config.ResolvedModelConfig;
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.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestClientResponseException;

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

@Slf4j
@Service
@ConditionalOnProperty(name = "ai.local-stub.enabled", havingValue = "false", matchIfMissing = true)
public class SeedreamServiceImpl implements SeedreamService {

    private static final String API_KEY_HINT =
            "未配置图片模型 API Key，请在运营端模型配置中填写，或设置 VOLCENGINE_ARK_API_KEY";
    private static final ObjectMapper JSON = new ObjectMapper();

    private final ArkProperties properties;
    private final PlatformModelConfigService platformModelConfigService;
    private final ProviderRestClientFactory restClientFactory;

    public SeedreamServiceImpl(ArkProperties properties) {
        this(properties, null, new ProviderRestClientFactory());
    }

    @Autowired
    public SeedreamServiceImpl(ArkProperties properties,
                               PlatformModelConfigService platformModelConfigService,
                               ProviderRestClientFactory restClientFactory) {
        this.properties = properties;
        this.platformModelConfigService = platformModelConfigService;
        this.restClientFactory = restClientFactory;
    }

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

    @Override
    public String generateImage(String prompt, String size) {
        return generateImage(prompt, size, List.of());
    }

    @Override
    public String generateImage(String prompt, String size, List<String> referenceImageUrls) {
        ResolvedModelConfig config = resolveConfig();
        String apiKey = resolveApiKey(config);
        Map<String, Object> body = new LinkedHashMap<>();
        body.put("model", config.modelId());
        body.put("prompt", prompt);
        body.put("n", 1);
        body.put("size", size);
        if (!config.isAgnes()) {
            body.put("response_format", "url");
        }

        List<String> cleanReferences = referenceImageUrls == null
                ? List.of()
                : referenceImageUrls.stream()
                .filter(url -> url != null && !url.isBlank())
                .map(String::trim)
                .toList();
        if (!cleanReferences.isEmpty()) {
            if (config.isProtocol("openai") || config.isProtocol("agnes")) {
                body.put("image_urls", cleanReferences);
            } else {
                body.put("image", cleanReferences.size() == 1 ? cleanReferences.get(0) : cleanReferences);
            }
        }

        log.info("Image generate: provider={}, model={}, size={}, references={}",
                config.provider(), config.modelId(), size, cleanReferences.size());
        try {
            ImageResponse resp = restClientFactory.create(config, Duration.ofSeconds(10), Duration.ofMinutes(2)).post()
                    .uri(config.endpoint())
                    .header("Authorization", "Bearer " + apiKey)
                    .body(body)
                    .retrieve()
                    .body(ImageResponse.class);

            if (resp == null || resp.getData() == null || resp.getData().isEmpty()) {
                throw new BizException(ErrorCode.INTERNAL_ERROR, "图片模型返回空结果");
            }
            String url = resp.getData().get(0).getUrl();
            log.info("Image generated: url={}", url);
            return url;
        } catch (RestClientResponseException e) {
            if (e.getStatusCode().value() == 401) {
                throw new BizException(ErrorCode.INTERNAL_ERROR, "图片模型鉴权失败，请检查 API Key 是否有效");
            }
            log.error("Image generate failed: status={}, body={}", e.getStatusCode(), e.getResponseBodyAsString(), e);
            String providerMessage = extractProviderError(e.getResponseBodyAsString());
            throw new BizException(ErrorCode.INTERNAL_ERROR,
                    "图片生成失败: HTTP " + e.getStatusCode().value()
                            + (providerMessage.isBlank() ? "" : " - " + providerMessage));
        } catch (BizException e) {
            throw e;
        } catch (Exception e) {
            log.error("Image generate failed", e);
            throw new BizException(ErrorCode.INTERNAL_ERROR, "图片生成失败: " + e.getMessage());
        }
    }

    @Override
    public String getModelId() {
        return resolveConfig().modelId();
    }

    @Override
    public String getModelName() {
        return resolveConfig().modelName();
    }

    @Override
    public String getProvider() {
        return resolveConfig().provider();
    }

    private ResolvedModelConfig resolveConfig() {
        if (platformModelConfigService != null) {
            return platformModelConfigService.resolve(ModelCapability.IMAGE);
        }
        return new ResolvedModelConfig(
                "volcengine",
                "Volcengine ARK",
                "seedream",
                properties.getBaseUrl(),
                properties.getApiKey(),
                properties.getImageModel(),
                "Seedream Image",
                "/images/generations",
                null,
                null
        );
    }

    private String resolveApiKey(ResolvedModelConfig config) {
        String apiKey = config.apiKey();
        if (apiKey == null || apiKey.isBlank()) {
            throw new BizException(ErrorCode.INTERNAL_ERROR, API_KEY_HINT);
        }
        String normalized = apiKey.trim();
        if ("your-ark-api-key".equalsIgnoreCase(normalized)
                || "your_ark_api_key".equalsIgnoreCase(normalized)
                || PlatformModelConfigService.MASKED_API_KEY.equals(normalized)) {
            throw new BizException(ErrorCode.INTERNAL_ERROR, API_KEY_HINT);
        }
        return normalized;
    }

    private String extractProviderError(String body) {
        if (body == null || body.isBlank()) {
            return "";
        }
        try {
            JsonNode root = JSON.readTree(body);
            JsonNode message = root.path("error").path("message");
            if (message.isTextual() && !message.asText().isBlank()) {
                return message.asText().trim();
            }
        } catch (Exception ignore) {
            // Fall back to the raw body below when provider response is not JSON.
        }
        String normalized = body.replaceAll("\\s+", " ").trim();
        return normalized.length() > 300 ? normalized.substring(0, 300) + "..." : normalized;
    }

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

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