package ai.decompile.ai.service;

import ai.decompile.ai.model.entity.EmbeddingChunk;
import ai.decompile.ai.model.repository.EmbeddingChunkRepository;
import ai.decompile.analysis.model.entity.StaticAnalysis;
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.StaticImportRepository;
import ai.decompile.analysis.model.repository.StaticStringRepository;
import ai.decompile.analysis.model.repository.StaticStructRepository;
import ai.decompile.workspace.service.BinaryService;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.UUID;
import lombok.RequiredArgsConstructor;
import lombok.extern.log4j.Log4j2;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

@Log4j2
@Service
@RequiredArgsConstructor
@ConditionalOnProperty(name = "app.ai.enabled", havingValue = "true", matchIfMissing = true)
public class EmbeddingService {
  private final EmbeddingClient embeddingClient;
  private final EmbeddingChunkBuilder chunkBuilder;
  private final EmbeddingChunkRepository embeddingChunkRepository;
  private final StaticAnalysisRepository analysisRepository;
  private final StaticStringRepository stringRepository;
  private final StaticImportRepository importRepository;
  private final StaticDataLabelRepository dataLabelRepository;
  private final StaticStructRepository structRepository;
  private final StaticEnumRepository enumRepository;
  private final BinaryService binaryService;

  @Value("${app.ai.embedding.batch-size:20}")
  private int batchSize;

  public float[] embed(String text) {
    return embeddingClient.embed(text);
  }

  @Transactional
  public int indexBinary(UUID binaryId) {
    log.info("Starting embedding generation for binary {}", binaryId);
    embeddingChunkRepository.deleteByBinaryId(binaryId);
    embeddingChunkRepository.flush();

    var binary = binaryService.getBinary(binaryId);
    var analysis = getLatestAnalysis(binaryId);
    if (Objects.isNull(analysis)) {
      throw new IllegalStateException(
          "No static analysis found for binary " + binaryId + ". Run analysis first.");
    }

    var functions = analysis.getFunctions();
    log.info("Indexing {} functions for binary {}", functions.size(), binaryId);

    var functionChunks = chunkBuilder.buildFunctionChunks(functions, binary);
    var allChunks = new ArrayList<>(functionChunks);
    allChunks.add(chunkBuilder.buildMetadataChunk(binary, analysis, functions.size()));

    var strings = stringRepository.findByAnalysisIdOrderByAddressAsc(analysis.getId());
    if (!strings.isEmpty()) {
      allChunks.addAll(chunkBuilder.buildStringChunks(strings, binary));
    }

    var imports = importRepository.findByAnalysisIdOrderByDllAscNameAsc(analysis.getId());
    if (!imports.isEmpty()) {
      allChunks.addAll(chunkBuilder.buildImportChunks(imports, binary));
    }

    var globalVars = dataLabelRepository.findByAnalysisIdOrderByNameAsc(analysis.getId());
    if (!globalVars.isEmpty()) {
      allChunks.addAll(chunkBuilder.buildGlobalVarChunks(globalVars, binary));
    }

    var structs = structRepository.findByAnalysisIdOrderByNameAsc(analysis.getId());
    if (!structs.isEmpty()) {
      allChunks.addAll(chunkBuilder.buildStructChunks(structs, binary));
    }

    var enums = enumRepository.findByAnalysisIdOrderByNameAsc(analysis.getId());
    if (!enums.isEmpty()) {
      allChunks.addAll(chunkBuilder.buildEnumChunks(enums, binary));
    }

    embedAndPersistChunks(allChunks);

    log.info("Embedding complete for binary {}: {} chunks", binaryId, allChunks.size());
    return allChunks.size();
  }

  private void embedAndPersistChunks(List<EmbeddingChunk> chunks) {
    for (var i = 0; i < chunks.size(); i += batchSize) {
      var batch = chunks.subList(i, Math.min(i + batchSize, chunks.size()));
      var texts = batch.stream().map(EmbeddingChunk::getContent).toList();
      var embeddings = embeddingClient.embed(texts);

      for (int j = 0; j < batch.size(); j++) {
        var c = batch.get(j);
        embeddingChunkRepository.insertChunk(
            c.getId(),
            c.getBinaryId(),
            c.getChunkType(),
            c.getSourceType(),
            c.getSourceId(),
            c.getContent(),
            AiUtils.toVectorString(embeddings.get(j)),
            c.getMetadata(),
            c.getTokenCount(),
            c.getCreatedAt());
      }

      log.debug(
          "Saved batch {}-{}/{} chunks",
          i + 1,
          Math.min(i + batchSize, chunks.size()),
          chunks.size());
    }
  }

  private StaticAnalysis getLatestAnalysis(UUID binaryId) {
    return analysisRepository.findFirstByBinaryIdOrderByCreatedAtDesc(binaryId).orElse(null);
  }
}
