src/main/java/com/dalab/policyengine/service/PolicyEvaluationService.java

package com.dalab.policyengine.service;

import java.time.Instant;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.stream.Collectors;

import org.jeasy.rules.api.Facts;
import org.jeasy.rules.api.Rule;
import org.jeasy.rules.api.Rules;
import org.jeasy.rules.api.RulesEngine;
import org.jeasy.rules.core.DefaultRulesEngine;
import org.jeasy.rules.mvel.MVELRule;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.Pageable;
import org.springframework.data.jpa.domain.Specification;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.util.StringUtils;

import com.dalab.common.event.AssetChangeEvent;
import com.dalab.policyengine.common.ResourceNotFoundException;
import com.dalab.policyengine.dto.PolicyEvaluationOutputDTO;
import com.dalab.policyengine.dto.PolicyEvaluationRequestDTO;
import com.dalab.policyengine.dto.PolicyEvaluationSummaryDTO;
import com.dalab.policyengine.kafka.producer.PolicyActionProducer;
import com.dalab.policyengine.mapper.PolicyMapper;
import com.dalab.policyengine.model.Policy;
import com.dalab.policyengine.model.PolicyEvaluation;
import com.dalab.policyengine.model.PolicyEvaluationStatus;
import com.dalab.policyengine.model.PolicyRule;
import com.dalab.policyengine.model.PolicyStatus;
import com.dalab.policyengine.repository.PolicyEvaluationRepository;
import com.dalab.policyengine.repository.PolicyRepository;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.protobuf.Timestamp;

import jakarta.persistence.criteria.Predicate;

@Service
@Transactional
public class PolicyEvaluationService implements IPolicyEvaluationService {

    private static final Logger log = LoggerFactory.getLogger(PolicyEvaluationService.class);

    private final PolicyRepository policyRepository;
    private final PolicyEvaluationRepository policyEvaluationRepository;
    private final PolicyMapper policyMapper;
    private final RulesEngine rulesEngine;
    private final ObjectMapper objectMapper; // For converting asset data and actions
    private final PolicyActionProducer policyActionProducer;

    @Autowired
    public PolicyEvaluationService(PolicyRepository policyRepository,
                                 PolicyEvaluationRepository policyEvaluationRepository,
                                 PolicyMapper policyMapper,
                                 ObjectMapper objectMapper,
                                 PolicyActionProducer policyActionProducer) {
        this.policyRepository = policyRepository;
        this.policyEvaluationRepository = policyEvaluationRepository;
        this.policyMapper = policyMapper;
        this.objectMapper = objectMapper;
        this.rulesEngine = new DefaultRulesEngine();
        this.policyActionProducer = policyActionProducer;
    }

    @Override
    public PolicyEvaluationOutputDTO evaluatePolicyForAsset(UUID policyId, PolicyEvaluationRequestDTO evaluationRequest, UUID triggeredByUserId) {
        Policy policy = policyRepository.findById(policyId)
                .orElseThrow(() -> new ResourceNotFoundException("Policy", "id", policyId.toString()));

        if (policy.getStatus() == PolicyStatus.DISABLED) {
            log.warn("Policy {} is disabled. Evaluation skipped for asset {}.", policy.getName(), evaluationRequest.getTargetAssetId());
            PolicyEvaluation evaluation = createEvaluationRecord(policy, evaluationRequest.getTargetAssetId(), PolicyEvaluationStatus.NOT_APPLICABLE, triggeredByUserId, Collections.singletonMap("reason", "Policy disabled"), null);
            return policyMapper.toPolicyEvaluationOutputDTO(evaluation, policy.getName());
        }

        Map<String, Object> assetData = fetchAssetData(evaluationRequest.getTargetAssetId());
        if (assetData == null || assetData.isEmpty()) {
            log.error("Asset data not found for assetId: {}. Evaluation cannot proceed.", evaluationRequest.getTargetAssetId());
            PolicyEvaluation evaluation = createEvaluationRecord(policy, evaluationRequest.getTargetAssetId(), PolicyEvaluationStatus.ERROR, triggeredByUserId, Collections.singletonMap("error", "Asset data not found"), null);
            return policyMapper.toPolicyEvaluationOutputDTO(evaluation, policy.getName());
        }

        Facts facts = new Facts();
        facts.put("asset", assetData); 
        if (evaluationRequest.getEvaluationContext() != null) {
            evaluationRequest.getEvaluationContext().forEach(facts::put);
        }

        Rules easyRules = new Rules();
        Map<String, Boolean> ruleResults = new HashMap<>();
        for (PolicyRule ruleEntity : policy.getRules()) {
            Rule easyRule = new MVELRule()
                    .name(ruleEntity.getName())
                    .description(ruleEntity.getDescription())
                    .priority(ruleEntity.getPriority())
                    .when(ruleEntity.getCondition())
                    .then("ruleResults.put(\"" + ruleEntity.getName() + "\", true); " +
                          "log.debug(\"Rule '{}' evaluated to true for asset {}\", \"" + ruleEntity.getName() + "\", evaluationRequest.getTargetAssetId());");
            easyRules.register(easyRule);
        }

        rulesEngine.fire(easyRules, facts);

        boolean overallPolicyConditionMet = true;
        if (StringUtils.hasText(policy.getConditionLogic())) {
            overallPolicyConditionMet = ruleResults.values().stream().allMatch(Boolean::booleanValue) && !ruleResults.isEmpty();
            log.debug("Policy condition logic '{}' evaluated to {} based on rule results: {}", policy.getConditionLogic(), overallPolicyConditionMet, ruleResults);
        } else {
            overallPolicyConditionMet = ruleResults.values().stream().allMatch(Boolean::booleanValue) && !ruleResults.isEmpty();
            log.debug("No policy condition logic. Overall result based on all rules: {}. Rule results: {}", overallPolicyConditionMet, ruleResults);
        }
        PolicyEvaluationStatus finalStatus = overallPolicyConditionMet ? PolicyEvaluationStatus.PASS : PolicyEvaluationStatus.FAIL;
        Map<String, Object> triggeredPolicyActionsSummary = null;

        PolicyEvaluation evaluation = createEvaluationRecord(policy, evaluationRequest.getTargetAssetId(), finalStatus, triggeredByUserId, 
            Map.of("rulesEvaluated", ruleResults, "factsSnapshot", objectMapper.convertValue(facts.asMap(), new TypeReference<Map<String, Object>>() {})), 
            null);

        if (finalStatus == PolicyEvaluationStatus.PASS) {
            log.info("Policy '{}' PASSED for asset '{}'", policy.getName(), evaluationRequest.getTargetAssetId());
            triggeredPolicyActionsSummary = executePolicyActions(policy, evaluationRequest.getTargetAssetId(), assetData, facts, evaluation.getId());
            evaluation.setTriggeredActions(triggeredPolicyActionsSummary); 
            policyEvaluationRepository.save(evaluation); 
        } else {
            log.info("Policy '{}' FAILED for asset '{}'", policy.getName(), evaluationRequest.getTargetAssetId());
        }
        
        return policyMapper.toPolicyEvaluationOutputDTO(evaluation, policy.getName());
    }

    @Override
    public PolicyEvaluationOutputDTO triggerPolicyEvaluation(UUID policyId, PolicyEvaluationRequestDTO evaluationRequest, UUID triggeredByUserId) {
        // Delegate to the main evaluation method
        return evaluatePolicyForAsset(policyId, evaluationRequest, triggeredByUserId);
    }

    private PolicyEvaluation createEvaluationRecord(Policy policy, String targetAssetId, PolicyEvaluationStatus status, UUID triggeredByUserId, Map<String, Object> details, Map<String, Object> triggeredActions) {
        PolicyEvaluation evaluation = new PolicyEvaluation();
        evaluation.setPolicyId(policy.getId());
        evaluation.setTargetAssetId(targetAssetId);
        evaluation.setStatus(status);
        evaluation.setEvaluationDetails(details);
        evaluation.setTriggeredActions(triggeredActions);
        evaluation.setEvaluatedAt(Instant.now());
        evaluation.setEvaluationTriggeredByUserId(triggeredByUserId);
        return policyEvaluationRepository.save(evaluation);
    }

    private Map<String, Object> fetchAssetData(String assetId) {
        log.debug("Fetching asset data for assetId: {}", assetId);
        Map<String, Object> data = new HashMap<>();
        if ("asset123".equals(assetId)) {
            data.put("name", "Sensitive S3 Bucket");
            data.put("assetType", "S3_BUCKET");
            data.put("region", "us-east-1");
            data.put("tags", Arrays.asList("PII", "FinancialData"));
            data.put("publicAccess", "true");
            data.put("sizeGB", 1024);
        } else if ("asset456".equals(assetId)) {
            data.put("name", "Dev EC2 Instance");
            data.put("assetType", "EC2_INSTANCE");
            data.put("region", "eu-west-1");
            data.put("tags", Arrays.asList("Development"));
            data.put("instanceType", "t2.micro");
        }
        return data.isEmpty() ? null : data;
    }

    private Map<String, Object> executePolicyActions(Policy policy, String assetId, Map<String, Object> assetData, Facts facts, UUID evaluationId) {
        if (policy.getActions() == null || policy.getActions().isEmpty()) {
            return Collections.emptyMap();
        }
        log.info("Executing actions for policy '{}' (id={}) on asset '{}': {}", 
            policy.getName(), policy.getId(), assetId, policy.getActions());

        Map<String, Object> executedActionsSummary = new HashMap<>();
        Instant now = Instant.now();
        Timestamp eventTimestamp = Timestamp.newBuilder().setSeconds(now.getEpochSecond()).setNanos(now.getNano()).build();

        final String effectiveAssetId = assetData.getOrDefault("id", assetId).toString(); // Prefer ID from assetData if available

        policy.getActions().forEach((actionType, actionConfig) -> {
            // For now, create a simple map-based event instead of protobuf until we add the protobuf definitions
            Map<String, Object> policyActionEvent = new HashMap<>();
            policyActionEvent.put("policyId", policy.getId().toString());
            policyActionEvent.put("assetId", effectiveAssetId);
            policyActionEvent.put("actionType", actionType);
            policyActionEvent.put("timestamp", now.toString());
            
            if (evaluationId != null) {
                policyActionEvent.put("evaluationId", evaluationId.toString());
            }
            
            if (actionConfig instanceof Map) {
                policyActionEvent.put("actionParameters", actionConfig);
            } else if (actionConfig != null) {
                policyActionEvent.put("actionParameters", Map.of("value", actionConfig));
            }

            // TODO: Replace with actual protobuf-based PolicyActionEvent when protobuf definitions are added
            log.info("Policy action event (simplified): {}", policyActionEvent);
            policyActionProducer.sendPolicyActionEvent(policyActionEvent);
            
            executedActionsSummary.put(actionType, actionConfig);
            log.debug("Created PolicyActionEvent for actionType: {} with assetId: {}", actionType, effectiveAssetId);
        });

        return executedActionsSummary;
    }

    @Override
    @Transactional(readOnly = true)
    public Page<PolicyEvaluationSummaryDTO> getPolicyEvaluations(Pageable pageable, UUID policyIdFilter, String targetAssetIdFilter, String statusFilter) {
        Specification<PolicyEvaluation> spec = (root, query, criteriaBuilder) -> {
            List<Predicate> predicates = new ArrayList<>();
            if (policyIdFilter != null) {
                predicates.add(criteriaBuilder.equal(root.get("policyId"), policyIdFilter));
            }
            if (StringUtils.hasText(targetAssetIdFilter)) {
                predicates.add(criteriaBuilder.equal(root.get("targetAssetId"), targetAssetIdFilter));
            }
            if (StringUtils.hasText(statusFilter)) {
                try {
                    predicates.add(criteriaBuilder.equal(root.get("status"), PolicyEvaluationStatus.valueOf(statusFilter.toUpperCase())));
                } catch (IllegalArgumentException e) {
                    log.warn("Invalid policy evaluation status provided for filtering: {}", statusFilter);
                    predicates.add(criteriaBuilder.disjunction());
                }
            }
            return criteriaBuilder.and(predicates.toArray(new Predicate[0]));
        };

        Page<PolicyEvaluation> evaluationPage = policyEvaluationRepository.findAll(spec, pageable);
        
        List<UUID> policyIds = evaluationPage.getContent().stream()
                                .map(PolicyEvaluation::getPolicyId)
                                .distinct()
                                .collect(Collectors.toList());
        Map<UUID, String> policyNamesMap = Collections.emptyMap();
        if (!policyIds.isEmpty()) {
            policyNamesMap = policyRepository.findAllById(policyIds).stream()
                .collect(Collectors.toMap(Policy::getId, Policy::getName));
        }
        final Map<UUID, String> finalPolicyNamesMap = policyNamesMap;

        return evaluationPage.map(eval -> policyMapper.toPolicyEvaluationSummaryDTO(eval, finalPolicyNamesMap.getOrDefault(eval.getPolicyId(), "Unknown Policy")));
    }

    @Override
    @Transactional(readOnly = true)
    public PolicyEvaluationOutputDTO getPolicyEvaluationById(UUID evaluationId) {
        PolicyEvaluation evaluation = policyEvaluationRepository.findById(evaluationId)
                .orElseThrow(() -> new ResourceNotFoundException("PolicyEvaluation", "id", evaluationId.toString()));
        Policy policy = policyRepository.findById(evaluation.getPolicyId())
            .orElse(null); // Policy might be deleted, handle gracefully
        return policyMapper.toPolicyEvaluationOutputDTO(evaluation, policy != null ? policy.getName() : "Unknown/Deleted Policy");
    }

    @Override
    public void evaluatePolicyForAssetInternal(AssetChangeEvent assetChangeEvent, UUID eventInitiatorId) {
        String assetIdStr = assetChangeEvent.getAssetId();
        log.info("Processing AssetChangeEvent for assetId: {}, eventType: {}", assetIdStr, assetChangeEvent.getEventType());

        List<Policy> activePolicies = policyRepository.findByStatus(PolicyStatus.ENABLED);
        if (activePolicies.isEmpty()) {
            log.info("No active policies found. No evaluations will be triggered for asset {}", assetIdStr);
            return;
        }

        log.info("Found {} active policies. Triggering evaluations for asset {}...", activePolicies.size(), assetIdStr);

        for (Policy policy : activePolicies) {
            try {
                PolicyEvaluationRequestDTO requestDTO = new PolicyEvaluationRequestDTO();
                requestDTO.setTargetAssetId(assetIdStr);
                evaluatePolicyForAsset(policy.getId(), requestDTO, eventInitiatorId);

            } catch (Exception e) {
                log.error("Error during evaluation of policy {} for asset {}: {}", policy.getName(), assetIdStr, e.getMessage(), e);
            }
        }
        log.info("Finished processing asset change event for assetId: {}", assetIdStr);
    }
}