package ai.decompile.ai.service;

import ai.decompile.ai.model.entity.EmbeddingChunk;
import ai.decompile.ai.model.repository.EmbeddingChunkRepository;
import ai.decompile.workspace.model.entity.Binary;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.UUID;
import java.util.stream.Stream;
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.Component;

@Log4j2
@Component
@RequiredArgsConstructor
@ConditionalOnProperty(name = "app.ai.enabled", havingValue = "true", matchIfMissing = true)
public class ChatContextBuilder {
  private static final int MIN_TOKEN_LENGTH = 3;

  private final EmbeddingService embeddingService;
  private final EmbeddingChunkRepository embeddingChunkRepository;

  @Value("${app.ai.rag.top-k:5}")
  private int topK;

  @Value("${app.ai.rag.similarity-threshold:0.6}")
  private double similarityThreshold;

  @Value("${app.ai.rag.fallback-min-ratio:0.4}")
  private double fallbackMinRatio;

  public String buildContext(UUID binaryId, String query) {
    var chunks = new LinkedHashSet<EmbeddingChunk>();
    enrichWithSemanticAndTextualSearch(chunks, binaryId, query);

    return formatPreview(chunks);
  }

  public String buildSystemPrompt(String context, Binary binary) {
    var binaryIdStr = binary.getId().toString();

    return String.format(
        """
        You are an expert reverse engineer and malware analyst.
        You are analyzing a binary file with the following characteristics.
        When referencing code, mention function names, addresses, and the binary context.
        Use the analysis tools available to you.
        For a complete function analysis: find the function, fetch its full assembly with
        call graph, and check both callers and callees.
        The context below is only a semantic hint — never rely on it for details.

        [BINARY OVERVIEW]
        Filename: %s
        Format: %s
        Architecture: %s
        Compiler: %s
        Binary ID: %s

        %s
        """,
        AiUtils.orUnknown(binary.getFilename()),
        AiUtils.orUnknown(binary.getFormat()),
        AiUtils.orUnknown(binary.getArchitecture()),
        AiUtils.orUnknown(binary.getCompiler()),
        binaryIdStr,
        context);
  }

  private void enrichWithSemanticAndTextualSearch(
      LinkedHashSet<EmbeddingChunk> merged, UUID binaryId, String query) {
    try {
      var semanticChunks = semanticSearch(binaryId, query);
      log.debug(
          "Semantic search returned {} chunks for binary {}", semanticChunks.size(), binaryId);
      merged.addAll(semanticChunks);

      if (semanticChunks.size() < topK * fallbackMinRatio) {
        var tokens = extractMeaningfulTokens(query);
        enrichWithTextualFallback(merged, binaryId, tokens);
      }
    } catch (Exception e) {
      log.warn("Semantic/textual search failed: {}", e.getMessage());
    }
  }

  private void enrichWithTextualFallback(
      LinkedHashSet<EmbeddingChunk> merged, UUID binaryId, List<String> tokens) {
    int remaining = topK - merged.size();

    for (var token : tokens) {
      if (remaining <= 0 || merged.size() >= topK) {
        break;
      }

      remaining = addTextChunksToMerged(merged, binaryId, token, remaining);
    }
  }

  private int addTextChunksToMerged(
      LinkedHashSet<EmbeddingChunk> merged, UUID binaryId, String token, int limit) {
    var textChunks = embeddingChunkRepository.searchByContent(binaryId, token, limit);

    for (var chunk : textChunks) {
      if (limit <= 0) {
        break;
      }

      if (merged.add(chunk)) {
        limit--;
      }
    }

    return limit;
  }

  private String formatPreview(LinkedHashSet<EmbeddingChunk> chunks) {
    if (chunks.isEmpty()) {
      return "";
    }

    var sb = new StringBuilder();

    for (var chunk : chunks) {
      sb.append("[").append(chunk.getChunkType()).append("] sourceId=").append(chunk.getSourceId());

      if (java.util.Objects.nonNull(chunk.getSourceType())) {
        sb.append(" sourceType=").append(chunk.getSourceType());
      }

      sb.append("\n");
    }

    log.debug("RAG preview: {} chunks formatted", chunks.size());

    return sb.toString();
  }

  private List<EmbeddingChunk> semanticSearch(UUID binaryId, String query) {
    var queryEmbedding = embeddingService.embed(query);
    var vectorStr = AiUtils.toVectorString(queryEmbedding);
    return embeddingChunkRepository.findSimilar(binaryId, vectorStr, similarityThreshold, topK);
  }

  private List<String> extractMeaningfulTokens(String query) {
    return Stream.of(query.split("[\\s,.;:!?()\\[\\]{}\"']+"))
        .filter(t -> t.length() >= MIN_TOKEN_LENGTH)
        .distinct()
        .limit(10)
        .toList();
  }
}
