package com.arms.api.sample.rereading;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.advisor.api.*;

import java.util.HashMap;
import java.util.Map;


@Slf4j
public class ReReadingAdvisor implements BaseAdvisor {

    private static final String DEFAULT_RE2_ADVISE_TEMPLATE = """
			{re2_input_query}
			Read the question again: {re2_input_query}
			""";


    private final String re2AdviseTemplate;

//    private int order = Ordered.LOWEST_PRECEDENCE;
    private int order = 0;

    // 기본생성자는 디폴트 템플릿을 사용
    public ReReadingAdvisor() {
        this(DEFAULT_RE2_ADVISE_TEMPLATE);
    }

    // 커스텀 템플릿을 사용
    public ReReadingAdvisor(String re2AdviseTemplate) {
        this.re2AdviseTemplate = re2AdviseTemplate;
    }

    @Override
    public AdvisedRequest before(AdvisedRequest advisedRequest) {
        log.info("ReReadingAdvisor::before - advisedRequest: {}", advisedRequest);

        String inputQuery = advisedRequest.userText();  // 실제 사용자가 입력한 message 값이 아닌 RAG의 결과값이 날라옴. RAG보다 먼저 실행 시킬 것인지 별도로 값을 받기
        log.info("ReReadingAdvisor::before - inputQuery: {}", inputQuery);

        Map<String, Object> params = new HashMap<>(advisedRequest.userParams());
        params.put("re2_input_query", inputQuery);
        log.info("ReReadingAdvisor::before - params: {}", params);

        AdvisedRequest advisedRequestBuild = AdvisedRequest.from(advisedRequest)
                .userText(DEFAULT_RE2_ADVISE_TEMPLATE)
                .userParams(params)
                .build();

        log.info("ReReadingAdvisor::before - advisedRequestBuild: {}", advisedRequestBuild);
        return advisedRequestBuild;

        // 스프링 공식 Template - https://docs.spring.io/spring-ai/reference/api/advisors.html
//        public ChatClientRequest before(ChatClientRequest chatClientRequest, AdvisorChain advisorChain) {
//            String augmentedUserText = PromptTemplate.builder()
//                    .template(this.re2AdviseTemplate)  //프롬프트 설정
//                    .variables(Map.of("re2_input_query", chatClientRequest.prompt().getUserMessage().getText()))  // 사용자가 입력한 문자
//                    .build()
//                    .render();
//
//            return chatClientRequest.mutate()  //chatClientRequest 변형
//                    .prompt(chatClientRequest.prompt().augmentUserMessage(augmentedUserText))  //변형된 프롬프트 설정
//                    .build();
//        }
    }

    @Override
    public AdvisedResponse after(AdvisedResponse advisedResponse) {
//        log.info("ReReadingAdvisor::after - advisedResponse: {}", advisedResponse);
        return advisedResponse;

        // 스프링 공식 Template
//        public ChatClientResponse after(ChatClientResponse chatClientResponse, AdvisorChain advisorChain) {
//            return chatClientResponse;
    }

//    @Override
//    public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
//        log.info("ReReadingAdvisor::aroundStream - advisedRequest: {}", advisedRequest);
//        return BaseAdvisor.super.aroundStream(advisedRequest, chain);
//    }

    @Override
    public int getOrder() {
//        log.info("ReReadingAdvisor::getOrder - order: {}", order);
        return this.order;
    }

    public ReReadingAdvisor withOrder(int order) {
        this.order = order;
//        log.info("ReReadingAdvisor::withOrder - order: {}, this.order: {}", order, this.order);
        return this;
    }

}

// 참고 소스: https://spring.io/blog/2024/10/02/supercharging-your-ai-applications-with-spring-ai-advisors