src/main/java/com/dalab/autocompliance/service/impl/ComplianceServiceImpl.java

package com.dalab.autocompliance.service.impl;

import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.stream.Collectors;

import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import com.dalab.autocompliance.dto.AssetComplianceStatusDTO;
import com.dalab.autocompliance.dto.ComplianceControlDTO;
import com.dalab.autocompliance.dto.ComplianceReportDTO;
import com.dalab.autocompliance.dto.ComplianceReportDefinitionDTO;
import com.dalab.autocompliance.dto.ControlEvaluationRequestDTO;
import com.dalab.autocompliance.dto.ControlEvaluationResponseDTO;
import com.dalab.autocompliance.dto.ReportGenerationRequestDTO;
import com.dalab.autocompliance.dto.ReportGenerationResponseDTO;
import com.dalab.autocompliance.dto.ReportJobStatusDTO;
import com.dalab.autocompliance.model.entity.ComplianceControlEntity;
import com.dalab.autocompliance.model.entity.ComplianceGeneratedReportEntity;
import com.dalab.autocompliance.model.entity.ComplianceReportDefinitionEntity;
import com.dalab.autocompliance.model.entity.ComplianceReportJobEntity;
import com.dalab.autocompliance.model.entity.ControlEvaluationJobEntity;
import com.dalab.autocompliance.model.mapper.ComplianceControlMapper;
import com.dalab.autocompliance.model.mapper.ComplianceGeneratedReportMapper;
import com.dalab.autocompliance.model.mapper.ComplianceReportDefinitionMapper;
import com.dalab.autocompliance.model.mapper.ComplianceReportJobMapper;
import com.dalab.autocompliance.model.mapper.ControlEvaluationJobMapper;
import com.dalab.autocompliance.model.repository.ComplianceControlRepository;
import com.dalab.autocompliance.model.repository.ComplianceGeneratedReportRepository;
import com.dalab.autocompliance.model.repository.ComplianceReportDefinitionRepository;
import com.dalab.autocompliance.model.repository.ComplianceReportJobRepository;
import com.dalab.autocompliance.model.repository.ControlEvaluationJobRepository;
import com.dalab.autocompliance.service.IComplianceService;
import com.dalab.autocompliance.service.exception.ControlNotFoundException;
import com.dalab.autocompliance.service.exception.GeneratedReportNotFoundException;
import com.dalab.autocompliance.service.exception.ReportDefinitionNotFoundException;
import com.dalab.autocompliance.service.exception.ReportJobNotFoundException;

import jakarta.annotation.PostConstruct;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;

@Service("complianceService")
@RequiredArgsConstructor
@Slf4j
public class ComplianceServiceImpl implements IComplianceService {

    private final ComplianceReportDefinitionRepository reportDefinitionRepository;
    private final ComplianceReportJobRepository reportJobRepository;
    private final ComplianceGeneratedReportRepository generatedReportRepository;
    private final ComplianceControlRepository controlRepository;
    private final ControlEvaluationJobRepository controlEvaluationJobRepository;

    private final ComplianceReportDefinitionMapper reportDefinitionMapper;
    private final ComplianceReportJobMapper reportJobMapper;
    private final ComplianceGeneratedReportMapper generatedReportMapper;
    private final ComplianceControlMapper controlMapper;
    private final ControlEvaluationJobMapper controlEvaluationJobMapper;

    @PostConstruct
    @Transactional
    public void initDefaultData() {
        initDefaultReportDefinitions();
        initDefaultControls();
    }

    public void initDefaultReportDefinitions() {
        if (reportDefinitionRepository.count() == 0) {
            log.info("Initializing default compliance report definitions...");
            List<ComplianceReportDefinitionEntity> defaultDefinitions = new ArrayList<>();
            defaultDefinitions.add(ComplianceReportDefinitionEntity.builder()
                .reportType("cis-gcp-foundations-v1.3")
                .displayName("CIS Google Cloud Platform Foundations Benchmark v1.3")
                .description("Checks for compliance against the CIS GCP Foundations Benchmark v1.3.")
                .version("1.3.0")
                .applicableComplianceFrameworks(List.of("CIS Benchmark"))
                .targetAssetTypes(List.of("GCP_PROJECT", "GCP_IAM_POLICY", "GCP_VPC_NETWORK"))
                .generationParametersDefinition(Map.of("gcpProjectId", "string", "targetRegions", "list_string"))
                .detailsLink("https://www.cisecurity.org/benchmark/google_cloud_computing_platform/")
                .build());

            defaultDefinitions.add(ComplianceReportDefinitionEntity.builder()
                .reportType("pci-dss-aws-v3.2.1")
                .displayName("PCI DSS v3.2.1 for AWS Infrastructure")
                .description("Assesses AWS resources against applicable PCI DSS v3.2.1 requirements.")
                .version("3.2.1")
                .applicableComplianceFrameworks(List.of("PCI DSS v3.2.1"))
                .targetAssetTypes(List.of("AWS_EC2_INSTANCE", "AWS_S3_BUCKET", "AWS_VPC"))
                .generationParametersDefinition(Map.of("awsAccountId", "string", "awsRegion", "string"))
                .detailsLink("https://aws.amazon.com/compliance/pci-dss/")
                .build());
            reportDefinitionRepository.saveAll(defaultDefinitions);
            log.info("{} default compliance report definitions initialized.", defaultDefinitions.size());
        }
    }

    public void initDefaultControls() {
        if (controlRepository.count() == 0) {
            log.info("Initializing default compliance controls...");
            List<ComplianceControlEntity> defaultControls = new ArrayList<>();
            defaultControls.add(ComplianceControlEntity.builder()
                .controlId("CIS-GCP-1.1")
                .name("Ensure that corporate login credentials are used instead of gmail.com accounts")
                .description("GCP supports using Google Workspace or Cloud Identity to manage user identities. It is recommended to use managed corporate accounts instead of personal gmail.com accounts for managing cloud resources.")
                .severity("MEDIUM")
                .applicableFrameworks(List.of("CIS GCP Foundations Benchmark v1.3"))
                .targetAssetTypes(List.of("GCP_IAM_POLICY"))
                .remediationSteps("Migrate any gmail.com accounts with IAM roles to managed corporate accounts.")
                .evaluationParametersDefinition(Map.of("gcpProjectId", "string"))
                .detailsLink("https://www.cisecurity.org/benchmark/google_cloud_computing_platform/") // Point to specific control if possible
                .enabled(true)
                .build());
            defaultControls.add(ComplianceControlEntity.builder()
                .controlId("PCI-REQ-8.2.3")
                .name("Use strong passwords (or passphrases)")
                .description("Ensure passwords/passphrases meet minimum complexity requirements as defined by PCI DSS requirement 8.2.3.")
                .severity("HIGH")
                .applicableFrameworks(List.of("PCI DSS v3.2.1"))
                .targetAssetTypes(List.of("AWS_IAM_USER", "Linux_Server")) // Example asset types
                .remediationSteps("Configure password policies to enforce complexity. For existing users, enforce password change at next login.")
                .evaluationParametersDefinition(Map.of("minPasswordLength", "integer", "requiresUppercase", "boolean")) // Example params
                .detailsLink("https://www.pcisecuritystandards.org/") // Point to specific control
                .enabled(true)
                .build());
            controlRepository.saveAll(defaultControls);
            log.info("{} default compliance controls initialized.", defaultControls.size());
        }
    }

    @Override
    @Transactional(readOnly = true)
    public List<ComplianceReportDefinitionDTO> listAvailableReportDefinitions() {
        log.debug("Fetching all available report definitions.");
        List<ComplianceReportDefinitionEntity> entities = reportDefinitionRepository.findAll();
        return reportDefinitionMapper.toDtoList(entities);
    }

    @Override
    @Transactional
    public ReportGenerationResponseDTO generateComplianceReport(String reportType, ReportGenerationRequestDTO request) {
        log.info("Request to generate report type: {} with parameters: {}", reportType, request.getParameters());

        ComplianceReportDefinitionEntity definitionEntity = reportDefinitionRepository.findByReportTypeIgnoreCase(reportType)
                .orElseThrow(() -> {
                    log.warn("Report type '{}' not found.", reportType);
                    return new ReportDefinitionNotFoundException("Report type " + reportType + " not found.");
                });

        // Validate parameters (basic check for key presence)
        if (request.getParameters() == null || 
            !request.getParameters().keySet().containsAll(definitionEntity.getGenerationParametersDefinition().keySet())) {
            String errorMsg = String.format("Missing required parameters for report type '%s'. Required: %s, Provided: %s",
                                   reportType, definitionEntity.getGenerationParametersDefinition().keySet(), 
                                   request.getParameters() != null ? request.getParameters().keySet() : "null");
            log.warn(errorMsg);
            // For immediate feedback, we can return a response, though throwing an exception is also an option.
            return ReportGenerationResponseDTO.builder()
                    .jobId(null).reportType(reportType).status("FAILED_VALIDATION")
                    .submittedAt(LocalDateTime.now()).message(errorMsg)
                    .build();
        }

        ComplianceReportJobEntity jobEntity = reportJobMapper.fromGenerationRequest(request, reportType);
        jobEntity.setJobId(UUID.randomUUID().toString());
        jobEntity.setStatus("QUEUED");
        jobEntity.setSubmittedAt(LocalDateTime.now());
        jobEntity.setReportType(definitionEntity.getReportType()); // Ensure canonical report type
        // triggeredBy and notificationEmail are mapped by mapper if present in request

        reportJobRepository.save(jobEntity);
        log.info("Report generation job {} for report type '{}' accepted and queued.", jobEntity.getJobId(), reportType);

        // TODO: Trigger actual asynchronous report generation (e.g., via @Async method or Kafka message)
        // This async process would later update the jobEntity status and link to a ComplianceGeneratedReportEntity.

        return ReportGenerationResponseDTO.builder()
                .jobId(jobEntity.getJobId()).reportType(reportType).status("ACCEPTED")
                .submittedAt(jobEntity.getSubmittedAt()).message("Report generation accepted and queued.")
                .build();
    }

    @Override
    @Transactional(readOnly = true)
    public ReportJobStatusDTO getReportGenerationJobStatus(String jobId) {
        log.debug("Fetching status for job ID: {}", jobId);
        ComplianceReportJobEntity jobEntity = reportJobRepository.findById(jobId)
                .orElseThrow(() -> {
                    log.warn("Job ID '{}' not found.", jobId);
                    return new ReportJobNotFoundException("Job ID " + jobId + " not found.");
                });
        return reportJobMapper.toDto(jobEntity);
    }

    @Override
    @Transactional(readOnly = true)
    public ComplianceReportDTO getGeneratedReport(String reportId) {
        log.debug("Fetching generated report for report ID: {}", reportId);
        ComplianceGeneratedReportEntity reportEntity = generatedReportRepository.findById(reportId)
                .orElseThrow(() -> {
                    log.warn("Report ID '{}' not found.", reportId);
                    return new GeneratedReportNotFoundException("Generated report with ID " + reportId + " not found.");
                });
        return generatedReportMapper.toDto(reportEntity);
    }

    @Override
    @Transactional(readOnly = true)
    public AssetComplianceStatusDTO getAssetComplianceStatus(String assetId) {
        log.debug("Fetching compliance status for asset ID: {}", assetId);
        
        // This is a simplified implementation. A more robust solution would involve:
        // 1. Querying `ComplianceGeneratedReportEntity` for reports containing this assetId (potentially using the native query or service-layer filtering).
        // 2. Aggregating findings from those reports.
        // 3. Determining overall status based on those findings.
        // For now, we mock it by looking for the dummy asset in any dummy report.

        List<ComplianceGeneratedReportEntity> reportsWithAsset = generatedReportRepository.findReportsContainingAssetId(assetId);

        if (reportsWithAsset.isEmpty()) {
            log.warn("Asset ID '{}' not found in any generated reports or no compliance data available.", assetId);
            return AssetComplianceStatusDTO.builder()
                    .assetId(assetId)
                    .overallComplianceStatus("UNKNOWN") // Or NOT_EVALUATED
                    .lastEvaluatedAt(LocalDateTime.now()) // Or null if never evaluated
                    .build();
        }

        List<AssetComplianceStatusDTO.ComplianceFindingSummaryDTO> assetFindingsSummaries = new ArrayList<>();
        List<String> relevantReportIds = new ArrayList<>();
        int compliantCount = 0;
        int nonCompliantCount = 0;
        String assetType = "UNKNOWN";
        String assetName = "UNKNOWN"; 
        LocalDateTime lastEvalTime = null;

        for (ComplianceGeneratedReportEntity reportEntity : reportsWithAsset) {
            relevantReportIds.add(reportEntity.getReportId());
            if (reportEntity.getFindingsData() != null) {
                for (com.dalab.autocompliance.model.entity.ReportFindingData findingData : reportEntity.getFindingsData()) {
                    if (assetId.equals(findingData.getAssetId())) {
                        // Update asset details from the first finding containing them (could be more sophisticated)
                        if ("UNKNOWN".equals(assetType) && findingData.getAssetType() != null) assetType = findingData.getAssetType();
                        if ("UNKNOWN".equals(assetName) && findingData.getAssetName() != null) assetName = findingData.getAssetName();
                        
                        if ("NON_COMPLIANT".equalsIgnoreCase(findingData.getStatus())) {
                            nonCompliantCount++;
                            assetFindingsSummaries.add(AssetComplianceStatusDTO.ComplianceFindingSummaryDTO.builder()
                                    .checkId(findingData.getCheckId())
                                    .description(findingData.getDescription())
                                    .severity(findingData.getSeverity())
                                    .reportId(reportEntity.getReportId())
                                    .findingTimestamp(reportEntity.getGenerationTimestamp()) 
                                    .build());
                        } else if ("COMPLIANT".equalsIgnoreCase(findingData.getStatus())) {
                            compliantCount++;
                        }
                        // Update last evaluation time for this asset based on reports
                        if (lastEvalTime == null || reportEntity.getGenerationTimestamp().isAfter(lastEvalTime)) {
                            lastEvalTime = reportEntity.getGenerationTimestamp();
                        }
                    }
                }
            }
        }

        String overallStatus = "UNKNOWN";
        if (compliantCount > 0 || nonCompliantCount > 0) {
            overallStatus = nonCompliantCount > 0 ? "NON_COMPLIANT" : "COMPLIANT";
        } else if (!reportsWithAsset.isEmpty()) {
             overallStatus = "NO_FINDINGS"; // Or EVALUATED_COMPLIANT if all checks were compliant
        }

        return AssetComplianceStatusDTO.builder()
                .assetId(assetId)
                .assetType(assetType)
                .assetName(assetName)
                .overallComplianceStatus(overallStatus)
                .lastEvaluatedAt(lastEvalTime != null ? lastEvalTime : LocalDateTime.now()) // Fallback, should ideally be from data
                .totalChecksApplied(compliantCount + nonCompliantCount) // This is simplified, total checks could be different
                .compliantChecks(compliantCount)
                .nonCompliantChecks(nonCompliantCount)
                .recentNonCompliantFindings(assetFindingsSummaries.stream().limit(5).collect(Collectors.toList()))
                .relevantReportIds(relevantReportIds)
                .build();
    }

    @Override
    @Transactional(readOnly = true)
    public List<ComplianceControlDTO> listAvailableControls() {
        log.debug("Fetching all available and enabled compliance controls.");
        List<ComplianceControlEntity> entities = controlRepository.findByEnabledTrue(); 
        return controlMapper.toDtoList(entities);
    }

    @Override
    @Transactional
    public ControlEvaluationResponseDTO evaluateControl(String controlId, ControlEvaluationRequestDTO request) {
        log.info("Request to evaluate control ID: {} with request: {}", controlId, request);

        ComplianceControlEntity controlEntity = controlRepository.findByControlIdIgnoreCase(controlId)
                .orElseThrow(() -> {
                    log.warn("Control ID '{}' not found.", controlId);
                    return new ControlNotFoundException("Control ID " + controlId + " not found.");
                });
        
        // Create and save the evaluation job
        ControlEvaluationJobEntity jobEntity = controlEvaluationJobMapper.fromRequestDto(request, controlId);
        jobEntity = controlEvaluationJobRepository.save(jobEntity);
        
        log.info("Created control evaluation job {} for control {}", jobEntity.getJobId(), controlId);
        
        // TODO: Implement actual control evaluation logic
        // This would typically involve:
        // 1. Fetching relevant assets based on request parameters
        // 2. Applying the control's evaluation logic
        // 3. Updating the job status and results
        
        return controlEvaluationJobMapper.toResponseDto(jobEntity);
    }
}