package com.seer.teach.ai.service;

import cn.hutool.core.bean.BeanUtil;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.seer.teach.ai.client.AiModelClientFactory;
import com.seer.teach.ai.client.model.AiModelClient;
import com.seer.teach.ai.client.model.config.ModelConfig;
import com.seer.teach.common.enums.ResultCodeEnum;
import com.seer.teach.common.exception.CommonException;
import com.seer.teach.common.utils.AssertUtils;
import com.seer.teach.teacher.module.entity.AiApiKeyEntity;
import com.seer.teach.teacher.module.entity.AiModelEntity;
import com.seer.teach.teacher.module.entity.AiScenarioConfigEntity;
import com.seer.teach.teacher.service.IAiApiKeyService;
import com.seer.teach.teacher.service.IAiModelService;
import com.seer.teach.teacher.service.IAiScenarioConfigService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Objects;
import java.util.Optional;

@Slf4j
@Service
@RequiredArgsConstructor
public class AiChatModelService {

    private final IAiScenarioConfigService aiScenarioConfigService;

    private final IAiApiKeyService aiApiKeyService;

    private final IAiModelService aiModelService;

    private final AiModelClientFactory aiModelClientFactory;

    public ChatResponse chatMessage(String scenarioCode, Object prompt){
        AiScenarioConfigEntity scenarioConfig = aiScenarioConfigService.getOneByScenarioCode(scenarioCode);
        log.info("scenarioConfig:{}", scenarioConfig);

        AssertUtils.notNull(scenarioConfig, ResultCodeEnum.SCENARIO_NOT_FOUND);

        AiModelClient client = aiModelClientFactory.getClient(scenarioConfig.getPlatform());

        ModelConfig modelConfig = getModelConfig(scenarioConfig);

        ChatModel chatModel = client.createChatModel(modelConfig).orElseThrow(() -> new CommonException(ResultCodeEnum.AI_MODEL_NOT_FOUND));

        Optional<ChatOptions> options = client.createChatOptions(BeanUtil.beanToMap(scenarioConfig));
        return chatModel.call(new Prompt(prompt.toString(), options.orElse(null)));
    }

    public Flux<ChatResponse> chatStream(String scenarioCode, Object params){
        AiScenarioConfigEntity scenarioConfig = aiScenarioConfigService.getOneByScenarioCode(scenarioCode);
        log.info("chatStream scenarioConfig:{}", scenarioConfig);

        AssertUtils.notNull(scenarioConfig, ResultCodeEnum.SCENARIO_NOT_FOUND);

        AiModelClient client = aiModelClientFactory.getClient(scenarioConfig.getPlatform());

        ModelConfig modelConfig = getModelConfig(scenarioConfig);

        ChatModel chatModel = client.createChatModel(modelConfig).orElseThrow(() -> new CommonException(ResultCodeEnum.AI_MODEL_NOT_FOUND));

        Optional<ChatOptions> options = client.createChatOptions(BeanUtil.beanToMap(scenarioConfig));
        Prompt prompt = new Prompt(params.toString(), options.orElse(null));
        return chatModel.stream(prompt);
    }


    private ModelConfig getModelConfig(AiScenarioConfigEntity scenarioConfig) {
        AiModelEntity module = aiModelService.getById(scenarioConfig.getModuleId());
        AiApiKeyEntity apiKey = getApiKey(scenarioConfig);
        return ModelConfig.builder()
                .url(module.getUrl()).apiKey(apiKey.getApiKey()).model(module.getModelIdentifier()).build();
    }

    private AiApiKeyEntity getApiKey(AiScenarioConfigEntity aiScenarioConfig) {
        List<AiApiKeyEntity> aiApiKeyList = aiApiKeyService.list(new LambdaQueryWrapper<>(AiApiKeyEntity.class).eq(AiApiKeyEntity::getPlatform, aiScenarioConfig.getPlatform()));
        AiApiKeyEntity apiKey = null;
        Optional<AiApiKeyEntity> any = aiApiKeyList.stream().filter(aiApiKey -> Objects.nonNull(aiScenarioConfig.getApiKeyId()) && aiScenarioConfig.getApiKeyId().intValue() == aiApiKey.getId()).findAny();
        if (any.isPresent()) {
            apiKey = any.get();
        }
        if (Objects.isNull(apiKey)) {
            Optional<AiApiKeyEntity> apiKeyOptional = aiApiKeyList.stream().filter(aiApiKey -> Objects.nonNull(aiApiKey.getIsDefault()) && aiApiKey.getIsDefault() == 1).findAny();
            if (apiKeyOptional.isPresent()) {
                apiKey = apiKeyOptional.get();
            }
        }
        return apiKey;
    }

}