package com.ai.trackwalker.model;

import com.ai.trackwalker.config.RailProperties;
import org.springframework.boot.web.client.RestTemplateBuilder;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestTemplate;

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

@Component
public class ModelArtifactInstallerClient {
    private final RailProperties properties;
    private final RestTemplateBuilder builder;

    public ModelArtifactInstallerClient(RailProperties properties, RestTemplateBuilder builder) {
        this.properties = properties;
        this.builder = builder;
    }

    public Map<String, Object> install(
            String jobId,
            String modelGroup,
            String modelVersion,
            List<Map<String, Object>> components
    ) {
        return post("/api/v1/artifacts/install", Map.of(
                "job_id", jobId,
                "model_group", modelGroup,
                "model_version", modelVersion,
                "components", components
        ));
    }

    public Map<String, Object> uninstall(String jobId, String modelGroup, String modelVersion) {
        return post("/api/v1/artifacts/uninstall", Map.of(
                "job_id", jobId,
                "model_group", modelGroup,
                "model_version", modelVersion
        ));
    }

    @SuppressWarnings("unchecked")
    public Map<String, Object> status() {
        HttpHeaders headers = headers();
        return client().exchange(properties.getArtifactInstaller().getUrl() + "/api/v1/artifacts/status",
                org.springframework.http.HttpMethod.GET, new HttpEntity<>(headers), Map.class).getBody();
    }

    @SuppressWarnings("unchecked")
    private Map<String, Object> post(String path, Map<String, Object> body) {
        Map<String, Object> response = client().postForObject(
                properties.getArtifactInstaller().getUrl() + path,
                new HttpEntity<>(body, headers()), Map.class);
        return response == null ? Map.of() : response;
    }

    private HttpHeaders headers() {
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.set("X-Artifact-Token", properties.getArtifactInstaller().getToken());
        return headers;
    }

    private RestTemplate client() {
        int timeout = Math.max(30, properties.getArtifactInstaller().getTimeoutSeconds());
        return builder.setConnectTimeout(Duration.ofSeconds(10))
                .setReadTimeout(Duration.ofSeconds(timeout)).build();
    }
}
