package com.agent.management.rag.bridge;
import com.agent.management.config.RagAiBridgeProperties;
import org.springframework.http.*;
import org.springframework.http.client.SimpleClientHttpRequestFactory;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestTemplate;
import java.util.*;
import java.time.Duration;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;

@Component
public class RagAiBridgeClient {
    private final RagAiBridgeProperties properties;
    private final RestTemplate rest;
    private final RestTemplate healthRest;
    private final Semaphore modelSlots = new Semaphore(1, true);
    public RagAiBridgeClient(RagAiBridgeProperties properties) {
        this.properties=properties;
        this.rest = rest(timeoutMillis(properties.getTimeout()));
        this.healthRest = rest(2000);
    }
    public boolean health(){ if(!properties.isEnabled()) return false; try{return Boolean.TRUE.equals(healthRest.getForEntity(url("/health"),Map.class).getBody().get("ok"));}catch(Exception e){return false;} }
    public Map<String,Object> textToSql(Map<String,Object> request){return post("/text2sql",request);}
    public Map<String,Object> textToCypher(Map<String,Object> request){return post("/text2cypher",request);}
    public Map<String,Object> answer(Map<String,Object> request){return post("/answer",request);}
    public Map<String,Object> suggestGovernance(Map<String,Object> request){return post("/governance/suggest",request);}
    public Map<String,Object> plan(Map<String,Object> request){return post("/plan",request);}
    public Map<String,Object> repair(Map<String,Object> request){return post("/repair",request);}
    public Map<String,Object> repair(Map<String,Object> request, Duration timeout){return post("/repair",request,timeout);}
    @SuppressWarnings("unchecked") private Map<String,Object> post(String path, Object body){
        return post(path, body, properties.getTimeout());
    }
    @SuppressWarnings("unchecked") private Map<String,Object> post(String path, Object body, Duration requestTimeout){
        if(!properties.isEnabled()) throw new IllegalStateException("rag-ai-bridge is disabled");
        boolean acquired = false;
        try {
            long timeoutMillis = effectiveTimeoutMillis(properties.getTimeout(), requestTimeout);
            if (timeoutMillis <= 0) {
                modelSlots.acquire();
                acquired = true;
            } else {
                acquired = modelSlots.tryAcquire(timeoutMillis, TimeUnit.MILLISECONDS);
                if (!acquired) throw new IllegalStateException("RAG model queue wait timed out");
            }
            RestTemplate client = timeoutMillis == timeoutMillis(properties.getTimeout()) ? rest : rest(timeoutMillis);
            var response=client.postForEntity(url(path),body,Map.class);
            return response.getBody()==null?Map.of():response.getBody();
        } catch (InterruptedException error) {
            Thread.currentThread().interrupt();
            throw new IllegalStateException("RAG model request was interrupted", error);
        } finally {
            if (acquired) modelSlots.release();
        }
    }
    private RestTemplate rest(long timeoutMillis) {
        var factory = new SimpleClientHttpRequestFactory();
        int timeout = timeoutMillis <= 0 ? 0 : (int) Math.min(Integer.MAX_VALUE, timeoutMillis);
        factory.setConnectTimeout(timeout);
        factory.setReadTimeout(timeout);
        return new RestTemplate(factory);
    }
    private static long effectiveTimeoutMillis(Duration globalTimeout, Duration requestTimeout) {
        long global = timeoutMillis(globalTimeout);
        long request = timeoutMillis(requestTimeout);
        if (global <= 0) return request;
        if (request <= 0) return global;
        return Math.min(global, request);
    }
    private static long timeoutMillis(Duration timeout) {
        return timeout == null || timeout.isZero() || timeout.isNegative() ? 0 : timeout.toMillis();
    }
    private String url(String path){return "http://"+properties.getHost()+":"+properties.getPort()+path;}
}
