backend/src/main/java/ai/decompile/analysis/service/AnalysisService.java

package ai.decompile.analysis.service;

import ai.decompile.analysis.model.dto.AnalysisResponse;
import ai.decompile.analysis.model.dto.FunctionDetailResponse;
import ai.decompile.analysis.model.dto.FunctionSummaryResponse;
import ai.decompile.analysis.model.dto.LabelResponse;
import ai.decompile.analysis.model.dto.XrefResponse;
import ai.decompile.analysis.model.entity.StaticAnalysis;
import ai.decompile.analysis.model.entity.StaticFunction;
import ai.decompile.analysis.model.entity.StaticLabel;
import ai.decompile.analysis.model.entity.StaticXref;
import ai.decompile.analysis.model.repository.StaticAnalysisRepository;
import ai.decompile.analysis.model.repository.StaticFunctionRepository;
import ai.decompile.analysis.model.repository.StaticLabelRepository;
import ai.decompile.analysis.model.repository.StaticXrefRepository;
import ai.decompile.common.exception.NotFoundException;
import ai.decompile.engine.model.AnalysisResult;
import ai.decompile.workspace.model.entity.Binary;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.function.Function;
import lombok.RequiredArgsConstructor;
import lombok.extern.log4j.Log4j2;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.Pageable;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

@Log4j2
@Service
@RequiredArgsConstructor
public class AnalysisService {
  private final StaticAnalysisRepository analysisRepository;
  private final StaticFunctionRepository functionRepository;
  private final StaticXrefRepository xrefRepository;
  private final StaticLabelRepository labelRepository;

  @Transactional
  public AnalysisResponse saveResult(
      Binary binary, UUID jobId, String engine, AnalysisResult result) {
    replaceExistingAnalysis(binary, engine);

    var analysis = createAnalysis(binary, jobId, engine, result);
    var nameToId = saveFunctions(analysis, result);
    saveXrefs(result, nameToId);

    return AnalysisResponse.from(analysis);
  }

  public List<AnalysisResponse> listAnalyses(UUID binaryId) {
    return analysisRepository.findByBinaryIdOrderByCreatedAtDesc(binaryId).stream()
        .map(AnalysisResponse::from)
        .toList();
  }

  public AnalysisResponse getAnalysis(UUID analysisId) {
    return AnalysisResponse.from(findAnalysis(analysisId));
  }

  public Page<FunctionSummaryResponse> listFunctions(UUID analysisId, Pageable pageable) {
    assertAnalysisExists(analysisId);
    return functionRepository
        .findByAnalysisId(analysisId, pageable)
        .map(FunctionSummaryResponse::from);
  }

  public FunctionDetailResponse getFunction(UUID functionId) {
    var func = findFunction(functionId);
    int calleeCount = xrefRepository.countByCallerId(functionId);
    int callerCount = xrefRepository.countByCalleeId(functionId);
    var labels =
        labelRepository.findByFunctionId(functionId).stream().map(LabelResponse::from).toList();
    return FunctionDetailResponse.from(func, callerCount, calleeCount, labels);
  }

  public List<XrefResponse> getCallers(UUID functionId) {
    return listXrefs(functionId, xrefRepository::findByCalleeId);
  }

  public List<XrefResponse> getCallees(UUID functionId) {
    return listXrefs(functionId, xrefRepository::findByCallerId);
  }

  private void replaceExistingAnalysis(Binary binary, String engine) {
    analysisRepository
        .findByBinaryIdAndEngine(binary.getId(), engine)
        .ifPresent(
            existing -> {
              log.info(
                  "Deleting existing {} analysis for binary {} (id={})",
                  engine,
                  binary.getId(),
                  existing.getId());
              analysisRepository.delete(existing);
              analysisRepository.flush();
            });
  }

  private StaticAnalysis createAnalysis(
      Binary binary, UUID jobId, String engine, AnalysisResult result) {
    return analysisRepository.save(
        StaticAnalysis.builder()
            .binary(binary)
            .jobId(jobId)
            .engine(engine)
            .status("COMPLETED")
            .functionCount(result.functions().size())
            .build());
  }

  private Map<String, UUID> saveFunctions(StaticAnalysis analysis, AnalysisResult result) {
    var nameToId = new HashMap<String, UUID>();
    for (var funcResult : result.functions()) {
      var func =
          functionRepository.save(
              StaticFunction.builder()
                  .analysis(analysis)
                  .name(funcResult.name())
                  .address(funcResult.address())
                  .assembly(funcResult.assembly())
                  .decompiledCode(funcResult.decompiledCode())
                  .signature(funcResult.signature())
                  .build());
      nameToId.put(funcResult.name(), func.getId());

      for (var label : funcResult.labels()) {
        labelRepository.save(
            StaticLabel.builder()
                .function(func)
                .name(label.name())
                .address(label.address())
                .build());
      }
    }
    return nameToId;
  }

  private void saveXrefs(AnalysisResult result, Map<String, UUID> nameToId) {
    for (var funcResult : result.functions()) {
      var callerId = nameToId.get(funcResult.name());
      if (callerId == null) {
        continue;
      }

      for (var calleeName : funcResult.xrefNames()) {
        var calleeId = nameToId.get(calleeName);
        if (calleeId == null) {
          log.debug(
              "Skipping xref from '{}' to unknown function '{}' (likely external/lib)",
              funcResult.name(),
              calleeName);
          continue;
        }

        xrefRepository.save(
            StaticXref.builder()
                .caller(functionRepository.getReferenceById(callerId))
                .callee(functionRepository.getReferenceById(calleeId))
                .type("CALL")
                .build());
      }
    }
  }

  private List<XrefResponse> listXrefs(UUID functionId, Function<UUID, List<StaticXref>> finder) {
    findFunction(functionId);
    return finder.apply(functionId).stream().map(XrefResponse::from).toList();
  }

  private void assertAnalysisExists(UUID analysisId) {
    if (!analysisRepository.existsById(analysisId)) {
      throw new NotFoundException("Static analysis with id '" + analysisId + "' not found");
    }
  }

  private StaticAnalysis findAnalysis(UUID analysisId) {
    return analysisRepository
        .findById(analysisId)
        .orElseThrow(
            () -> new NotFoundException("Static analysis with id '" + analysisId + "' not found"));
  }

  private StaticFunction findFunction(UUID functionId) {
    return functionRepository
        .findById(functionId)
        .orElseThrow(
            () -> new NotFoundException("Function with id '" + functionId + "' not found"));
  }
}