package kz.konturai.parser.service;

import com.google.auth.oauth2.GoogleCredentials;
import io.netty.resolver.DefaultAddressResolverGroup;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.core.io.ClassPathResource;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;
import org.springframework.web.reactive.function.client.ExchangeStrategies;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.reactive.function.client.WebClientResponseException;
import reactor.netty.http.client.HttpClient;
import org.springframework.http.client.reactive.ReactorClientHttpConnector;

import java.io.IOException;
import java.io.InputStream;
import java.net.URI;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.Semaphore;

/**
 * Сервис генерации изображений через Google Imagen 3.
 *
 * УЛУЧШЕНИЯ:
 * - Соотношение сторон зависит от платформы (4:5 для Instagram, 9:16 для TikTok/Reels)
 * - Негативный промпт на русском — единый язык с основным промптом
 * - Референс-изображение не перебивает основной промпт английскими инструкциями
 * - Метод generateImageForPlatform() для точного контроля формата
 */
@Service
public class NanoBananaImageGenerationService implements ImageGenerationService {

    private static final Logger logger = LoggerFactory.getLogger(NanoBananaImageGenerationService.class);

    private static final String VERTEX_API_TEMPLATE =
            "https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/publishers/google/models/%s:predict";
    private static final String CREDENTIALS_FILE_PATH = "keys/google-key.json";

    // Негативный промпт на русском — соответствует языку основных промптов
    private static final String NEGATIVE_PROMPT_RU =
            "размытое изображение, плохая анатомия, лишние конечности, шесть пальцев, три руки, " +
                    "деформированное тело, неправильные пропорции, мультяшный стиль, 3D рендер, " +
                    "искусственный вид, водяной знак, текстовые надписи на изображении, " +
                    "нечёткий фон, пересвет, плохое качество, низкое разрешение";

    // Один запрос в момент времени — Imagen 3 не любит параллельные запросы
    private final Semaphore semaphore = new Semaphore(1);
    private volatile long lastRequestTimestamp = 0;

    // Минимальный интервал между запросами (15 сек — требование rate limit Imagen 3)
    private static final long MIN_REQUEST_INTERVAL_MS = 15_000;

    private final WebClient webClient;

    @Value("${google.cloud.project-id}")
    private String projectId;

    @Value("${google.cloud.location:us-central1}")
    private String location;

    @Value("${google.gemini.image.model:imagen-3.0-generate-001}")
    private String model;

    public NanoBananaImageGenerationService() {
        HttpClient httpClient = HttpClient.create()
                .resolver(DefaultAddressResolverGroup.INSTANCE)
                .responseTimeout(Duration.ofMillis(180_000));

        ExchangeStrategies strategies = ExchangeStrategies.builder()
                .codecs(configurer -> configurer.defaultCodecs()
                        .maxInMemorySize(30 * 1024 * 1024))
                .build();

        this.webClient = WebClient.builder()
                .clientConnector(new ReactorClientHttpConnector(httpClient))
                .exchangeStrategies(strategies)
                .build();
    }

    // =====================================================================
    // PUBLIC API
    // =====================================================================

    /**
     * Базовая генерация — квадрат 1:1 (обратная совместимость).
     */
    @Override
    public byte[] generateImage(String prompt) {
        return generateImageInternal(prompt, null, "1:1");
    }

    /**
     * Генерация с логотипом — квадрат 1:1.
     */
    public byte[] generateImage(String prompt, MultipartFile referenceLogo) {
        if (referenceLogo == null || referenceLogo.isEmpty()) {
            return generateImage(prompt);
        }
        try {
            return generateImageWithReference(prompt, referenceLogo.getBytes());
        } catch (IOException e) {
            return generateImage(prompt);
        }
    }

    /**
     * Генерация с референсным изображением — квадрат 1:1.
     * Используется для обратной совместимости со старым кодом.
     */
    @Override
    public byte[] generateImageWithReference(String prompt, byte[] referenceImageBytes) {
        return generateImageInternal(prompt, referenceImageBytes, "1:1");
    }

    /**
     * ОСНОВНОЙ МЕТОД для генерации постов.
     * Соотношение сторон выбирается автоматически по платформе:
     * - Instagram      → 4:5  (портретный, оптимален для ленты)
     * - TikTok         → 9:16 (вертикальный, полный экран)
     * - Telegram       → 4:5  (портретный)
     * - LinkedIn       → 4:5  (портретный)
     * - По умолчанию  → 4:5
     *
     * @param prompt              Промпт на русском языке
     * @param referenceImageBytes Байты логотипа/референса (может быть null)
     * @param platform            Платформа: "instagram", "tiktok", "telegram", "linkedin"
     */
    public byte[] generateImageForPlatform(String prompt, byte[] referenceImageBytes, String platform) {
        String aspectRatio = resolveAspectRatio(platform);
        logger.info("[Imagen3] Platform='{}' → aspectRatio='{}'", platform, aspectRatio);
        return generateImageInternal(prompt, referenceImageBytes, aspectRatio);
    }

    // =====================================================================
    // CORE GENERATION
    // =====================================================================

    private byte[] generateImageInternal(String prompt, byte[] referenceImageBytes, String aspectRatio) {
        if (projectId == null || projectId.trim().isEmpty()) {
            logger.error("[Imagen3] Project ID not configured!");
            return null;
        }

        try {
            semaphore.acquire();

            // Rate limit — минимум 15 сек между запросами
            long timeSinceLast = System.currentTimeMillis() - lastRequestTimestamp;
            if (timeSinceLast < MIN_REQUEST_INTERVAL_MS) {
                long waitMs = MIN_REQUEST_INTERVAL_MS - timeSinceLast;
                logger.debug("[Imagen3] Rate limit wait: {}ms", waitMs);
                Thread.sleep(waitMs);
            }

            String accessToken = getAccessToken();
            if (accessToken == null) {
                logger.error("[Imagen3] Failed to get access token");
                return null;
            }

            String endpointUrl = String.format(VERTEX_API_TEMPLATE,
                    location, projectId, location, model);

            String finalPrompt = (prompt != null && !prompt.isBlank()) ? prompt : "Качественное изображение";

            logger.info("[Imagen3] Generating. AspectRatio={}, PromptLen={}",
                    aspectRatio, finalPrompt.length());

            try {
                Map<String, Object> requestBody =
                        buildRequestBody(finalPrompt, referenceImageBytes, aspectRatio);
                Map<String, Object> response =
                        executeRequest(endpointUrl, accessToken, requestBody);
                lastRequestTimestamp = System.currentTimeMillis();

                byte[] result = extractImage(response);
                if (result != null) {
                    logger.info("[Imagen3] ✅ Success. Bytes={}", result.length);
                } else {
                    logger.warn("[Imagen3] Empty result from API");
                }
                return result;

            } catch (WebClientResponseException.BadRequest badRequestEx) {
                logger.warn("[Imagen3] 400 BadRequest with reference — retrying without reference. Body: {}",
                        badRequestEx.getResponseBodyAsString());

                // Если референс вызвал ошибку — пробуем без него
                if (referenceImageBytes != null && referenceImageBytes.length > 0) {
                    Map<String, Object> fallbackBody =
                            buildRequestBody(finalPrompt, null, aspectRatio);
                    Map<String, Object> fallbackResponse =
                            executeRequest(endpointUrl, accessToken, fallbackBody);
                    lastRequestTimestamp = System.currentTimeMillis();
                    return extractImage(fallbackResponse);
                }
                return null;

            } catch (Exception ex) {
                logger.error("[Imagen3] Generation error: {}", ex.getMessage());
                return null;
            }

        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            logger.error("[Imagen3] Interrupted");
            return null;
        } catch (Exception e) {
            logger.error("[Imagen3] Fatal: {}", e.getMessage(), e);
            return null;
        } finally {
            semaphore.release();
        }
    }

    // =====================================================================
    // REQUEST BUILDING
    // =====================================================================

    private Map<String, Object> buildRequestBody(String prompt,
                                                 byte[] referenceImageBytes,
                                                 String aspectRatio) {
        Map<String, Object> body = new HashMap<>();
        List<Map<String, Object>> instances = new ArrayList<>();
        Map<String, Object> instance = new HashMap<>();

        // Промпт — только на русском, без добавления английских инструкций
        instance.put("prompt", prompt);

        // Референсное изображение (логотип клиента) — если передан
        if (referenceImageBytes != null && referenceImageBytes.length > 0) {
            String base64Image = Base64.getEncoder().encodeToString(referenceImageBytes);

            Map<String, Object> imageBytesMap = new HashMap<>();
            imageBytesMap.put("bytesBase64Encoded", base64Image);

            Map<String, Object> referenceImageParams = new HashMap<>();
            referenceImageParams.put("referenceImage", imageBytesMap);
            // STYLE — используем стиль лого, не SUBJECT (SUBJECT часто вызывает 400)
            referenceImageParams.put("referenceType", "STYLE");

            instance.put("referenceImages", List.of(referenceImageParams));

            logger.debug("[Imagen3] Using reference image ({} bytes)", referenceImageBytes.length);
        }

        instances.add(instance);
        body.put("instances", instances);

        // Параметры генерации
        Map<String, Object> parameters = new HashMap<>();
        parameters.put("sampleCount", 1);
        parameters.put("aspectRatio", aspectRatio);
        parameters.put("safetyFilterLevel", "block_some");
        parameters.put("personGeneration", "allow_adult");
        parameters.put("negativePrompt", NEGATIVE_PROMPT_RU);

        body.put("parameters", parameters);
        return body;
    }

    // =====================================================================
    // HTTP
    // =====================================================================

    private Map<String, Object> executeRequest(String endpointUrl,
                                               String accessToken,
                                               Map<String, Object> requestBody) {
        try {
            return webClient.post()
                    .uri(URI.create(endpointUrl))
                    .header("Authorization", "Bearer " + accessToken)
                    .contentType(MediaType.APPLICATION_JSON)
                    .bodyValue(requestBody)
                    .retrieve()
                    .bodyToMono(new ParameterizedTypeReference<Map<String, Object>>() {})
                    // Retry только на 429 (rate limit)
                    .retryWhen(reactor.util.retry.Retry
                            .backoff(5, Duration.ofSeconds(10))
                            .maxBackoff(Duration.ofSeconds(60))
                            .filter(t -> t instanceof WebClientResponseException
                                    && ((WebClientResponseException) t).getStatusCode().value() == 429)
                            .doBeforeRetry(signal ->
                                    logger.warn("[Imagen3] 429 rate limit — retrying ({}/5)...",
                                            signal.totalRetries() + 1)))
                    .block(Duration.ofMillis(180_000));
        } catch (WebClientResponseException e) {
            if (e.getStatusCode().value() == 400) {
                throw e; // Пробрасываем для обработки в вызывающем коде (retry без референса)
            }
            logger.error("[Imagen3] HTTP {}: {}", e.getStatusCode(), e.getResponseBodyAsString());
            return null;
        } catch (Exception e) {
            logger.error("[Imagen3] Request error: {}", e.getMessage());
            return null;
        }
    }

    // =====================================================================
    // AUTH
    // =====================================================================

    private String getAccessToken() {
        try {
            ClassPathResource resource = new ClassPathResource(CREDENTIALS_FILE_PATH);
            if (!resource.exists()) {
                logger.error("[Imagen3] Credentials file not found: {}", CREDENTIALS_FILE_PATH);
                return null;
            }
            try (InputStream is = resource.getInputStream()) {
                GoogleCredentials credentials = GoogleCredentials.fromStream(is)
                        .createScoped("https://www.googleapis.com/auth/cloud-platform");
                credentials.refreshIfExpired();
                return credentials.getAccessToken().getTokenValue();
            }
        } catch (IOException e) {
            logger.error("[Imagen3] Failed to get access token: {}", e.getMessage());
            return null;
        }
    }

    // =====================================================================
    // RESPONSE PARSING
    // =====================================================================

    @SuppressWarnings("unchecked")
    private byte[] extractImage(Map<String, Object> response) {
        if (response == null) return null;
        try {
            List<Map<String, Object>> predictions =
                    (List<Map<String, Object>>) response.get("predictions");
            if (predictions != null && !predictions.isEmpty()) {
                Map<String, Object> first = predictions.get(0);
                // Imagen 3 возвращает bytesBase64Encoded
                String base64 = (String) first.get("bytesBase64Encoded");
                if (base64 == null) {
                    // Fallback — старый ключ
                    base64 = (String) first.get("b64");
                }
                if (base64 != null && !base64.isBlank()) {
                    return Base64.getDecoder().decode(base64);
                }
            }
        } catch (Exception e) {
            logger.error("[Imagen3] Failed to extract image from response: {}", e.getMessage());
        }
        return null;
    }

    // =====================================================================
    // UTILS
    // =====================================================================

    /**
     * Выбирает правильное соотношение сторон под платформу.
     *
     * Imagen 3 поддерживает: 1:1, 4:3, 3:4, 16:9, 9:16
     *
     * Instagram лента → 4:5 НЕ поддерживается Imagen 3, поэтому используем 3:4 (близко к 4:5)
     * TikTok / Reels  → 9:16 (вертикальный полный экран)
     */
    private String resolveAspectRatio(String platform) {
        if (platform == null) return "3:4";
        return switch (platform.toLowerCase().trim()) {
            case "tiktok"    -> "9:16"; // вертикальный для TikTok Reels
            case "instagram" -> "3:4";  // портретный для Instagram ленты (ближайший к 4:5)
            case "telegram"  -> "3:4";  // портретный
            case "linkedin"  -> "4:3";  // горизонтальный для LinkedIn
            case "youtube"   -> "16:9"; // горизонтальный для YouTube
            default          -> "3:4";  // по умолчанию портретный
        };
    }
}