package ai.decompile.analysis.service;

import ai.decompile.analysis.model.dto.AnalysisResponse;
import ai.decompile.analysis.model.dto.DataLabelResponse;
import ai.decompile.analysis.model.dto.EnumResponse;
import ai.decompile.analysis.model.dto.FunctionDetailResponse;
import ai.decompile.analysis.model.dto.FunctionSummaryResponse;
import ai.decompile.analysis.model.dto.ImportResponse;
import ai.decompile.analysis.model.dto.LabelResponse;
import ai.decompile.analysis.model.dto.SegmentResponse;
import ai.decompile.analysis.model.dto.StringResponse;
import ai.decompile.analysis.model.dto.StructResponse;
import ai.decompile.analysis.model.dto.XrefResponse;
import ai.decompile.analysis.model.entity.StaticAnalysis;
import ai.decompile.analysis.model.entity.StaticDataLabel;
import ai.decompile.analysis.model.entity.StaticEnum;
import ai.decompile.analysis.model.entity.StaticFunction;
import ai.decompile.analysis.model.entity.StaticImport;
import ai.decompile.analysis.model.entity.StaticLabel;
import ai.decompile.analysis.model.entity.StaticSegment;
import ai.decompile.analysis.model.entity.StaticString;
import ai.decompile.analysis.model.entity.StaticStruct;
import ai.decompile.analysis.model.entity.StaticXref;
import ai.decompile.analysis.model.repository.StaticAnalysisRepository;
import ai.decompile.analysis.model.repository.StaticDataLabelRepository;
import ai.decompile.analysis.model.repository.StaticEnumRepository;
import ai.decompile.analysis.model.repository.StaticFunctionRepository;
import ai.decompile.analysis.model.repository.StaticImportRepository;
import ai.decompile.analysis.model.repository.StaticLabelRepository;
import ai.decompile.analysis.model.repository.StaticSegmentRepository;
import ai.decompile.analysis.model.repository.StaticStringRepository;
import ai.decompile.analysis.model.repository.StaticStructRepository;
import ai.decompile.analysis.model.repository.StaticXrefRepository;
import ai.decompile.common.exception.NotFoundException;
import ai.decompile.engine.model.AnalysisResult;
import ai.decompile.engine.model.SegmentEntry;
import ai.decompile.workspace.model.entity.Binary;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
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;
  private final StaticSegmentRepository segmentRepository;
  private final StaticStringRepository stringRepository;
  private final StaticImportRepository importRepository;
  private final StaticDataLabelRepository dataLabelRepository;
  private final StaticStructRepository structRepository;
  private final StaticEnumRepository enumRepository;

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

    var analysis = createAnalysis(binary, jobId, engine, result);
    saveSegments(analysis, result.segments());
    saveStrings(analysis, result.strings());
    saveImports(analysis, result.imports());
    saveGlobalVars(analysis, result.globalVars());
    saveStructures(analysis, result.structures());
    saveEnums(analysis, result.enums());
    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);
    var resolution = buildNameResolutionMap(analysisId);
    return functionRepository
        .findByAnalysisId(analysisId, pageable)
        .map(f -> FunctionSummaryResponse.from(f, resolution));
  }

  public FunctionDetailResponse getFunction(UUID functionId) {
    var func = findFunction(functionId);
    var resolution = buildNameResolutionMap(func.getAnalysis().getId());
    var resolvedName = resolution.get(func.getName());
    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, resolvedName);
  }

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

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

  public List<SegmentResponse> getSegments(UUID analysisId) {
    assertAnalysisExists(analysisId);
    return segmentRepository.findByAnalysisIdOrderByNameAsc(analysisId).stream()
        .map(SegmentResponse::from)
        .toList();
  }

  public List<StringResponse> getStrings(UUID analysisId) {
    assertAnalysisExists(analysisId);
    return stringRepository.findByAnalysisIdOrderByAddressAsc(analysisId).stream()
        .map(StringResponse::from)
        .toList();
  }

  public List<ImportResponse> getImports(UUID analysisId) {
    assertAnalysisExists(analysisId);
    return importRepository.findByAnalysisIdOrderByDllAscNameAsc(analysisId).stream()
        .map(ImportResponse::from)
        .toList();
  }

  public List<DataLabelResponse> getGlobalVars(UUID analysisId) {
    assertAnalysisExists(analysisId);
    return dataLabelRepository.findByAnalysisIdOrderByNameAsc(analysisId).stream()
        .map(DataLabelResponse::from)
        .toList();
  }

  public List<StructResponse> getStructures(UUID analysisId) {
    assertAnalysisExists(analysisId);
    return structRepository.findByAnalysisIdOrderByNameAsc(analysisId).stream()
        .map(StructResponse::from)
        .toList();
  }

  public List<EnumResponse> getEnums(UUID analysisId) {
    assertAnalysisExists(analysisId);
    return enumRepository.findByAnalysisIdOrderByNameAsc(analysisId).stream()
        .map(EnumResponse::from)
        .toList();
  }

  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 void saveSegments(StaticAnalysis analysis, List<SegmentEntry> entries) {
    for (var entry : entries) {
      segmentRepository.save(
          StaticSegment.builder()
              .analysis(analysis)
              .name(entry.name())
              .startAddress(entry.startAddress())
              .endAddress(entry.endAddress())
              .build());
    }
  }

  private void saveStrings(
      StaticAnalysis analysis, List<ai.decompile.engine.model.StringEntry> entries) {
    for (var entry : entries) {
      stringRepository.save(
          StaticString.builder()
              .analysis(analysis)
              .address(entry.address())
              .value(entry.value())
              .encoding(entry.encoding())
              .referencedBy(entry.referencedBy())
              .build());
    }
  }

  private void saveImports(
      StaticAnalysis analysis, List<ai.decompile.engine.model.ImportEntry> entries) {
    for (var entry : entries) {
      importRepository.save(
          StaticImport.builder()
              .analysis(analysis)
              .dll(entry.dll())
              .name(entry.name())
              .address(entry.address())
              .ordinal(entry.ordinal())
              .build());
    }
  }

  private void saveGlobalVars(
      StaticAnalysis analysis, List<ai.decompile.engine.model.GlobalVarEntry> entries) {
    for (var entry : entries) {
      dataLabelRepository.save(
          StaticDataLabel.builder()
              .analysis(analysis)
              .name(entry.name())
              .address(entry.address())
              .segment(entry.segment())
              .size(entry.size())
              .dataXrefs(entry.dataXrefs())
              .build());
    }
  }

  private void saveStructures(
      StaticAnalysis analysis, List<ai.decompile.engine.model.StructEntry> entries) {
    for (var entry : entries) {
      structRepository.save(
          StaticStruct.builder()
              .analysis(analysis)
              .name(entry.name())
              .size(entry.size())
              .isUnion(entry.isUnion())
              .members(entry.members())
              .build());
    }
  }

  private void saveEnums(
      StaticAnalysis analysis, List<ai.decompile.engine.model.EnumEntry> entries) {
    for (var entry : entries) {
      enumRepository.save(
          StaticEnum.builder()
              .analysis(analysis)
              .name(entry.name())
              .width(entry.width())
              .constants(entry.constants())
              .build());
    }
  }

  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())
            .entryPoint(result.entryPoint())
            .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())
                  .functionSize(funcResult.size())
                  .flags(funcResult.flags())
                  .assembly(funcResult.assembly())
                  .decompiledCode(funcResult.decompiledCode())
                  .signature(funcResult.signature())
                  .stackFrame(funcResult.stackFrame())
                  .comments(funcResult.comments())
                  .cfgBlocks(funcResult.cfgBlocks())
                  .build());
      nameToId.put(funcResult.name(), func.getId());

      for (var label : funcResult.labels()) {
        labelRepository.save(
            StaticLabel.builder()
                .function(func)
                .name(label.name())
                .address(label.address())
                .type(label.type())
                .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 xrefEntry : funcResult.xrefs()) {
        var calleeId = nameToId.get(xrefEntry.name());
        if (calleeId == null) {
          log.debug(
              "Skipping xref from '{}' to unknown function '{}' (likely external/lib)",
              funcResult.name(),
              xrefEntry.name());
          continue;
        }

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

  private List<XrefResponse> listXrefs(UUID functionId, Function<UUID, List<StaticXref>> finder) {
    var func = findFunction(functionId);
    var resolution = buildNameResolutionMap(func.getAnalysis().getId());
    return finder.apply(functionId).stream().map(xr -> XrefResponse.from(xr, resolution)).toList();
  }

  private Map<String, String> buildNameResolutionMap(UUID analysisId) {
    var imports = importRepository.findByAnalysisIdOrderByDllAscNameAsc(analysisId);
    if (imports.isEmpty()) {
      return Map.of();
    }

    var map = new HashMap<String, String>();
    for (var imp : imports) {
      if (Objects.isNull(imp.getResolvedName())) {
        continue;
      }

      map.putIfAbsent(imp.getName(), imp.getResolvedName());
      if (imp.getOrdinal() > 0) {
        map.putIfAbsent("ord_" + imp.getOrdinal(), imp.getResolvedName());
        map.putIfAbsent("j_ord_" + imp.getOrdinal(), imp.getResolvedName());
        map.putIfAbsent(imp.getDll() + "_" + imp.getOrdinal(), imp.getResolvedName());
      }
    }

    return map;
  }

  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"));
  }
}
