package ai.decompile.analysis.service;

import ai.decompile.analysis.model.repository.LibraryAliasRepository;
import ai.decompile.analysis.model.repository.LibraryRepository;
import ai.decompile.analysis.model.repository.LibrarySymbolRepository;
import ai.decompile.analysis.model.repository.StaticImportRepository;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import java.util.UUID;
import lombok.RequiredArgsConstructor;
import lombok.extern.log4j.Log4j2;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

@Log4j2
@Service
@RequiredArgsConstructor
public class SymbolResolutionService {
  private final LibraryAliasRepository aliasRepository;
  private final LibrarySymbolRepository symbolRepository;
  private final LibraryRepository libraryRepository;
  private final StaticImportRepository importRepository;

  @Transactional
  public int resolveImports(UUID analysisId) {
    var imports = importRepository.findByAnalysisIdOrderByDllAscNameAsc(analysisId);
    if (imports.isEmpty()) {
      return 0;
    }

    var lookup = buildLibraryLookup();
    if (lookup.isEmpty()) {
      log.debug("No libraries with symbols found, skipping resolution for analysis {}", analysisId);
      return 0;
    }

    var resolvedCount = 0;
    for (var imp : imports) {
      if (imp.getOrdinal() == 0) {
        continue;
      }

      if (Objects.nonNull(imp.getResolvedName())) {
        continue;
      }

      var libraryId = lookup.get(imp.getDll().toLowerCase());
      if (Objects.isNull(libraryId)) {
        libraryId = lookup.get(imp.getDll().toLowerCase().replace(".dll", ""));
      }

      if (Objects.isNull(libraryId)) {
        continue;
      }

      var symbol = symbolRepository.findByLibraryIdAndOrdinal(libraryId, imp.getOrdinal());
      if (symbol.isEmpty()) {
        continue;
      }

      imp.setLibrarySymbol(symbol.get());
      imp.setResolvedName(symbol.get().getSymbolName());
      importRepository.save(imp);
      resolvedCount++;
    }

    var totalWithOrdinal = imports.stream().filter(i -> i.getOrdinal() != 0).count();
    log.info(
        "Symbol resolution: resolved {}/{} ordinal imports ({}%) for analysis {}",
        resolvedCount,
        totalWithOrdinal,
        totalWithOrdinal > 0 ? (resolvedCount * 100L / totalWithOrdinal) : 0,
        analysisId);

    return resolvedCount;
  }

  @Transactional
  public int resolveAllUnresolved() {
    var allUnresolved =
        importRepository.findAll().stream()
            .filter(
                i ->
                    i.getOrdinal() != 0
                        && Objects.isNull(i.getResolvedName())
                        && Objects.isNull(i.getLibrarySymbolId()))
            .toList();

    if (allUnresolved.isEmpty()) {
      return 0;
    }

    var lookup = buildLibraryLookup();
    if (lookup.isEmpty()) {
      return 0;
    }

    var resolvedCount = 0;
    for (var imp : allUnresolved) {
      var libraryId = lookup.get(imp.getDll().toLowerCase());
      if (Objects.isNull(libraryId)) {
        libraryId = lookup.get(imp.getDll().toLowerCase().replace(".dll", ""));
      }

      if (Objects.isNull(libraryId)) {
        continue;
      }

      var symbol = symbolRepository.findByLibraryIdAndOrdinal(libraryId, imp.getOrdinal());
      if (symbol.isEmpty()) {
        continue;
      }

      imp.setLibrarySymbol(symbol.get());
      imp.setResolvedName(symbol.get().getSymbolName());
      importRepository.save(imp);
      resolvedCount++;
    }

    log.info("Global symbol resolution: resolved {} previously unresolved imports", resolvedCount);
    return resolvedCount;
  }

  private Map<String, UUID> buildLibraryLookup() {
    var lookup = new HashMap<String, UUID>();

    for (var alias : aliasRepository.findAll()) {
      lookup.putIfAbsent(alias.getAlias().toLowerCase(), alias.getLibrary().getId());
    }

    for (var lib : libraryRepository.findAll()) {
      for (var sym : lib.getSymbols()) {
        if (Objects.nonNull(sym.getModuleName()) && !sym.getModuleName().isBlank()) {
          lookup.putIfAbsent(sym.getModuleName().toLowerCase(), lib.getId());
        }
      }
    }

    return lookup;
  }
}
