/*
 * Copyright 2025 Google. This software is provided as-is, without warranty or representation for any use or purpose. 
 * Your use of it is subject to your agreement with Google.
 */

package com.example.controller;

import com.example.dto.llm.webhook.WebhookRequestDTO;
import com.example.dto.llm.webhook.SessionInfoDTO;
import com.example.dto.llm.webhook.WebhookResponseDTO;
import com.example.service.llm.LlmResponseTunerService;
import java.util.HashMap;
import java.util.Map;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.ExceptionHandler;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Mono;

@RestController
@RequestMapping("/api/v1/llm")
public class LlmResponseTunerController {

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

    private final LlmResponseTunerService llmResponseTunerService;

    public LlmResponseTunerController(LlmResponseTunerService llmResponseTunerService) {
        this.llmResponseTunerService = llmResponseTunerService;
    }

    @PostMapping("/tune-response")
    public Mono<WebhookResponseDTO> tuneResponse(@RequestBody WebhookRequestDTO request) {
        String uuid = (String) request.getSessionInfo().getParameters().get("uuid");
        return llmResponseTunerService
                .getValue(uuid)
                .map(
                        value -> {
                            Map<String, Object> parameters = new HashMap<>();
                            parameters.put("webhook_success", true);
                            parameters.put("response", value);
                            SessionInfoDTO sessionInfo = new SessionInfoDTO(parameters);
                            return new WebhookResponseDTO(sessionInfo);
                        })
                .defaultIfEmpty(createErrorResponse("No response found for the given UUID.", false))
                .onErrorResume(
                        e -> {
                            logger.error("Error tuning response: {}", e.getMessage());
                            return Mono.just(
                                    createErrorResponse("An internal error occurred.", true));
                        });
    }

    private WebhookResponseDTO createErrorResponse(String errorMessage, boolean isError) {
        Map<String, Object> parameters = new HashMap<>();
        parameters.put("webhook_success", false);
        parameters.put("error_message", errorMessage);
        SessionInfoDTO sessionInfo = new SessionInfoDTO(parameters);
        return new WebhookResponseDTO(sessionInfo);
    }

    @ExceptionHandler(Exception.class)
    public ResponseEntity<Map<String, String>> handleException(Exception e) {
        logger.error("An unexpected error occurred: {}", e.getMessage());
        Map<String, String> response = new HashMap<>();
        response.put("error", "Internal Server Error");
        response.put("message", "An unexpected error occurred. Please try again later.");
        return new ResponseEntity<>(response, HttpStatus.INTERNAL_SERVER_ERROR);
    }

    @ExceptionHandler(IllegalArgumentException.class)
    public ResponseEntity<Map<String, String>> handleIllegalArgumentException(
            IllegalArgumentException e) {
        logger.error("Bad request: {}", e.getMessage());
        Map<String, String> response = new HashMap<>();
        response.put("error", "Bad Request");
        response.put("message", e.getMessage());
        return new ResponseEntity<>(response, HttpStatus.BAD_REQUEST);
    }
}