package kr.re.etri.autoflow.controllers;

import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Mono;

import java.util.Collections;
import java.util.Map;

@Tag(name = "MLflow API", description = "MLflow Experiment 및 Run 조회 API")
@RestController
@RequestMapping("/api/mlflow")
public class MlflowController {

    private final WebClient webClient;

    public MlflowController() {
        this.webClient = WebClient.builder()
                .baseUrl("http://192.168.10.135:30128/api/2.0/mlflow")
                .defaultHeaders(headers -> headers.setBasicAuth("user", "LImQa2Me37nu"))
                .build();
    }

    @Operation(summary = "Experiment 조회", description = "Experiment 이름으로 MLflow Experiment 정보를 조회합니다.")
    @GetMapping(value = "/experiment", produces = MediaType.APPLICATION_JSON_VALUE)
    public Map<String, Object> getExperimentByName(@RequestParam String experimentName) {
        Map response = webClient.get()
                .uri(uriBuilder -> uriBuilder
                        .path("/experiments/get-by-name")
                        .queryParam("experiment_name", experimentName)
                        .build())
                .retrieve()
                .bodyToMono(Map.class)
                .block();

        if (response == null || !response.containsKey("experiment")) {
            throw new RuntimeException("Experiment not found: " + experimentName);
        }

        return (Map<String, Object>) response.get("experiment");
    }


    @Operation(summary = "Experiment의 Run 목록 조회", description = "주어진 Experiment ID의 Run을 최대 1000개까지 조회합니다.")
    @GetMapping(value = "/runs", produces = MediaType.APPLICATION_JSON_VALUE)
    public Mono<ResponseEntity<String>> getRuns(@RequestParam String experimentId) {
        Map<String, Object> body = Map.of(
                "experiment_ids", Collections.singletonList(experimentId),
                "order_by", Collections.singletonList("attribute.start_time DESC"),
                "max_results", 1000
        );

        return webClient.post()
                .uri("/runs/search")
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(body)
                .retrieve()
                .bodyToMono(String.class)
                .map(ResponseEntity::ok)
                .onErrorResume(e -> Mono.just(ResponseEntity.internalServerError().body(e.getMessage())));
    }

}