package ai.decompile.ai.service;

import ai.decompile.ai.model.entity.EmbeddingChunk;
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.StaticString;
import ai.decompile.analysis.model.entity.StaticStruct;
import ai.decompile.analysis.model.entity.StaticXref;
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.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.workspace.model.entity.Binary;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.util.ArrayList;
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 org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.stereotype.Component;

@Component
@RequiredArgsConstructor
@ConditionalOnProperty(name = "app.ai.enabled", havingValue = "true", matchIfMissing = true)
public class EmbeddingChunkBuilder {
  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;
  private final ObjectMapper objectMapper;

  @Value("${app.ai.embedding.chunk-max-chars:3000}")
  private int chunkMaxChars;

  private static final int METADATA_TOP_STRINGS = 20;
  private static final int METADATA_TOP_IMPORTS = 30;
  private static final int METADATA_TOP_GLOBAL_VARS = 30;
  private static final int METADATA_TOP_STRUCTS = 20;
  private static final int METADATA_TOP_ENUMS = 20;

  public List<EmbeddingChunk> buildFunctionChunks(List<StaticFunction> functions, Binary binary) {
    var graphContext = loadGraphContext(functions);

    var chunks = new ArrayList<EmbeddingChunk>();
    for (var func : functions) {
      chunks.add(
          EmbeddingChunk.builder()
              .id(UUID.randomUUID())
              .binaryId(binary.getId())
              .chunkType("FUNCTION")
              .sourceType("STATIC_FUNCTION")
              .sourceId(func.getId())
              .content(buildFunctionChunkContent(func, binary, graphContext))
              .metadata(AiUtils.toJson(objectMapper, buildFunctionMetadata(func, graphContext)))
              .tokenCount(AiUtils.estimateTokens(func.getAssembly()))
              .createdAt(java.time.Instant.now())
              .build());
    }

    return chunks;
  }

  public EmbeddingChunk buildMetadataChunk(
      Binary binary, StaticAnalysis analysis, int functionCount) {
    var sb = new StringBuilder();
    sb.append(
        String.format(
            "[METADATA] Binary Overview\nFilename: %s\nFormat: %s\nArchitecture: %s\nCompiler: %s\nFunction count: %d\n",
            AiUtils.orUnknown(binary.getFilename()),
            AiUtils.orUnknown(binary.getFormat()),
            AiUtils.orUnknown(binary.getArchitecture()),
            AiUtils.orUnknown(binary.getCompiler()),
            functionCount));

    if (Objects.nonNull(analysis.getEntryPoint())) {
      sb.append("Entry point: ").append(analysis.getEntryPoint()).append("\n");
    }

    var segments = segmentRepository.findByAnalysisIdOrderByNameAsc(analysis.getId());
    if (!segments.isEmpty()) {
      sb.append("Segments:\n");
      for (var seg : segments) {
        sb.append(
            String.format(
                "  %s: %s-%s\n", seg.getName(), seg.getStartAddress(), seg.getEndAddress()));
      }
    }

    var strings = stringRepository.findByAnalysisIdOrderByAddressAsc(analysis.getId());
    if (!strings.isEmpty()) {
      sb.append("Top strings:\n");
      strings.stream()
          .limit(METADATA_TOP_STRINGS)
          .forEach(
              s ->
                  sb.append(
                      String.format(
                          "  %s (%s): \"%s\"\n",
                          s.getAddress(), s.getEncoding(), AiUtils.truncate(s.getValue(), 200))));
    }

    var imports = importRepository.findByAnalysisIdOrderByDllAscNameAsc(analysis.getId());
    var libraryDiagnostics = new LibraryDiagnostics();
    if (!imports.isEmpty()) {
      sb.append("Imports:\n");
      imports.stream()
          .limit(METADATA_TOP_IMPORTS)
          .forEach(
              i -> {
                sb.append(String.format("  %s.%s @ %s", i.getDll(), i.getName(), i.getAddress()));
                if (Objects.nonNull(i.getResolvedName())) {
                  sb.append(" → ").append(i.getResolvedName());
                  collectLibraryInfo(libraryDiagnostics, i);
                }
                sb.append("\n");
              });
    }

    appendLibraryDiagnostics(sb, libraryDiagnostics, imports);

    var globalVars = dataLabelRepository.findByAnalysisIdOrderByNameAsc(analysis.getId());
    if (!globalVars.isEmpty()) {
      sb.append("Global variables:\n");
      globalVars.stream()
          .limit(METADATA_TOP_GLOBAL_VARS)
          .forEach(
              g ->
                  sb.append(
                      String.format(
                          "  %s @ %s (%s, %d bytes)\n",
                          g.getName(), g.getAddress(), g.getSegment(), g.getSize())));
    }

    var structs = structRepository.findByAnalysisIdOrderByNameAsc(analysis.getId());
    if (!structs.isEmpty()) {
      sb.append("Structures:\n");
      structs.stream()
          .limit(METADATA_TOP_STRUCTS)
          .forEach(
              s ->
                  sb.append(
                      String.format(
                          "  %s (%d bytes, %s)\n",
                          s.getName(), s.getSize(), s.isUnion() ? "union" : "struct")));
    }

    var enums = enumRepository.findByAnalysisIdOrderByNameAsc(analysis.getId());
    if (!enums.isEmpty()) {
      sb.append("Enums:\n");
      enums.stream()
          .limit(METADATA_TOP_ENUMS)
          .forEach(e -> sb.append(String.format("  %s (%d-bit)\n", e.getName(), e.getWidth())));
    }

    return EmbeddingChunk.builder()
        .id(UUID.randomUUID())
        .binaryId(binary.getId())
        .chunkType("METADATA")
        .content(sb.toString())
        .metadata(
            AiUtils.toJson(
                objectMapper, Map.of("chunk_type", "METADATA", "function_count", functionCount)))
        .tokenCount(AiUtils.estimateTokens(sb.toString()))
        .createdAt(java.time.Instant.now())
        .build();
  }

  public List<EmbeddingChunk> buildStringChunks(List<StaticString> strings, Binary binary) {
    var chunks = new ArrayList<EmbeddingChunk>();
    for (var s : strings) {
      var sb = new StringBuilder();
      sb.append(String.format("[STRING] at %s (%s)\n", s.getAddress(), s.getEncoding()));
      if (Objects.nonNull(s.getReferencedBy()) && !s.getReferencedBy().isBlank()) {
        sb.append("Referenced by: ").append(s.getReferencedBy()).append("\n");
      }
      sb.append("\"").append(s.getValue()).append("\"");
      var content = sb.toString();
      chunks.add(
          EmbeddingChunk.builder()
              .id(UUID.randomUUID())
              .binaryId(binary.getId())
              .chunkType("STRING")
              .sourceType("STATIC_STRING")
              .sourceId(s.getId())
              .content(content)
              .metadata(
                  AiUtils.toJson(
                      objectMapper,
                      Map.of(
                          "address",
                          s.getAddress(),
                          "encoding",
                          s.getEncoding(),
                          "length",
                          s.getValue().length(),
                          "referenced_by",
                          Objects.nonNull(s.getReferencedBy()) ? s.getReferencedBy() : "")))
              .tokenCount(AiUtils.estimateTokens(content))
              .createdAt(java.time.Instant.now())
              .build());
    }

    return chunks;
  }

  public List<EmbeddingChunk> buildImportChunks(List<StaticImport> imports, Binary binary) {
    var chunks = new ArrayList<EmbeddingChunk>();
    for (var imp : imports) {
      var content = buildImportContent(imp);
      var metadata = buildImportMetadata(imp);
      chunks.add(
          EmbeddingChunk.builder()
              .id(UUID.randomUUID())
              .binaryId(binary.getId())
              .chunkType("IMPORT")
              .sourceType("STATIC_IMPORT")
              .sourceId(imp.getId())
              .content(content)
              .metadata(AiUtils.toJson(objectMapper, metadata))
              .tokenCount(AiUtils.estimateTokens(content))
              .createdAt(java.time.Instant.now())
              .build());
    }

    return chunks;
  }

  private String buildImportContent(StaticImport imp) {
    var sb = new StringBuilder();
    sb.append(String.format("[IMPORT] %s.%s @ %s", imp.getDll(), imp.getName(), imp.getAddress()));

    if (Objects.nonNull(imp.getResolvedName())) {
      sb.append("\n  Resolved: ").append(imp.getResolvedName());

      var sym = imp.getLibrarySymbol();
      if (Objects.nonNull(sym)) {
        sb.append(" [").append(sym.getLibrary().getName()).append("]");

        if (Objects.nonNull(sym.getClassName())) {
          sb.append(" Class: ").append(sym.getClassName());
        }

        if (Objects.nonNull(sym.getCategory())) {
          sb.append(" Category: ").append(sym.getCategory());
        }
      }
    }

    return sb.toString();
  }

  private Map<String, Object> buildImportMetadata(StaticImport imp) {
    var metadata = new HashMap<String, Object>();
    metadata.put("dll", imp.getDll());
    metadata.put("name", imp.getName());
    metadata.put("address", imp.getAddress());
    metadata.put("ordinal", imp.getOrdinal());

    if (Objects.nonNull(imp.getResolvedName())) {
      metadata.put("resolved_name", imp.getResolvedName());

      var sym = imp.getLibrarySymbol();
      if (Objects.nonNull(sym)) {
        metadata.put("library_name", sym.getLibrary().getName());
        metadata.put("class_name", Objects.nonNull(sym.getClassName()) ? sym.getClassName() : "");
        metadata.put("category", Objects.nonNull(sym.getCategory()) ? sym.getCategory() : "");
      }
    }

    return metadata;
  }

  public List<EmbeddingChunk> buildGlobalVarChunks(
      List<StaticDataLabel> globalVars, Binary binary) {
    var chunks = new ArrayList<EmbeddingChunk>();
    for (var gv : globalVars) {
      var sb = new StringBuilder();
      sb.append(
          String.format(
              "[GLOBAL VAR] %s at %s (%s, %d bytes)\n",
              gv.getName(), gv.getAddress(), gv.getSegment(), gv.getSize()));
      if (Objects.nonNull(gv.getDataXrefs())
          && !gv.getDataXrefs().isBlank()
          && !"[]".equals(gv.getDataXrefs().trim())) {
        sb.append("Data xrefs: ")
            .append(AiUtils.truncate(gv.getDataXrefs(), chunkMaxChars / 4))
            .append("\n");
      }
      var content = sb.toString();
      chunks.add(
          EmbeddingChunk.builder()
              .id(UUID.randomUUID())
              .binaryId(binary.getId())
              .chunkType("GLOBAL_VAR")
              .sourceType("STATIC_DATA_LABEL")
              .sourceId(gv.getId())
              .content(content)
              .metadata(
                  AiUtils.toJson(
                      objectMapper,
                      Map.of(
                          "name",
                          gv.getName(),
                          "address",
                          gv.getAddress(),
                          "segment",
                          gv.getSegment(),
                          "size",
                          gv.getSize())))
              .tokenCount(AiUtils.estimateTokens(content))
              .createdAt(java.time.Instant.now())
              .build());
    }

    return chunks;
  }

  public List<EmbeddingChunk> buildStructChunks(List<StaticStruct> structs, Binary binary) {
    var chunks = new ArrayList<EmbeddingChunk>();
    for (var s : structs) {
      var content =
          String.format(
              "[STRUCT] %s (%d bytes, %s)\n%s",
              s.getName(),
              s.getSize(),
              s.isUnion() ? "union" : "struct",
              AiUtils.truncate(s.getMembers(), chunkMaxChars / 2));
      chunks.add(
          EmbeddingChunk.builder()
              .id(UUID.randomUUID())
              .binaryId(binary.getId())
              .chunkType("STRUCT")
              .sourceType("STATIC_STRUCT")
              .sourceId(s.getId())
              .content(content)
              .metadata(
                  AiUtils.toJson(objectMapper, Map.of("name", s.getName(), "size", s.getSize())))
              .tokenCount(AiUtils.estimateTokens(content))
              .createdAt(java.time.Instant.now())
              .build());
    }

    return chunks;
  }

  public List<EmbeddingChunk> buildEnumChunks(List<StaticEnum> enums, Binary binary) {
    var chunks = new ArrayList<EmbeddingChunk>();
    for (var e : enums) {
      var content =
          String.format(
              "[ENUM] %s (%d-bit)\n%s",
              e.getName(), e.getWidth(), AiUtils.truncate(e.getConstants(), chunkMaxChars / 2));
      chunks.add(
          EmbeddingChunk.builder()
              .id(UUID.randomUUID())
              .binaryId(binary.getId())
              .chunkType("ENUM")
              .sourceType("STATIC_ENUM")
              .sourceId(e.getId())
              .content(content)
              .metadata(
                  AiUtils.toJson(objectMapper, Map.of("name", e.getName(), "width", e.getWidth())))
              .tokenCount(AiUtils.estimateTokens(content))
              .createdAt(java.time.Instant.now())
              .build());
    }

    return chunks;
  }

  private FunctionGraphContext loadGraphContext(List<StaticFunction> functions) {
    var nameToCallers =
        loadXrefNames(functions, xrefRepository::findByCalleeId, xr -> xr.getCaller().getName());

    var nameToCallees =
        loadXrefNames(functions, xrefRepository::findByCallerId, xr -> xr.getCallee().getName());

    var funcToLabels = loadLabels(functions);

    return new FunctionGraphContext(nameToCallers, nameToCallees, funcToLabels);
  }

  private Map<String, List<String>> loadXrefNames(
      List<StaticFunction> functions,
      Function<UUID, List<StaticXref>> queryFn,
      Function<StaticXref, String> nameFn) {
    var m = new HashMap<String, List<String>>();
    for (var f : functions) {
      var xrefs = queryFn.apply(f.getId());
      if (!xrefs.isEmpty()) {
        m.put(f.getName(), xrefs.stream().map(nameFn).distinct().toList());
      }
    }

    return m;
  }

  private Map<UUID, List<String>> loadLabels(List<StaticFunction> functions) {
    var m = new HashMap<UUID, List<String>>();
    for (var f : functions) {
      var labels = labelRepository.findByFunctionId(f.getId());
      if (!labels.isEmpty()) {
        m.put(
            f.getId(),
            labels.stream().map(l -> l.getAddress() + " (" + l.getName() + ")").toList());
      }
    }

    return m;
  }

  private String buildFunctionChunkContent(
      StaticFunction func, Binary binary, FunctionGraphContext ctx) {
    var sb = new StringBuilder();
    sb.append(
        String.format(
            "[FUNCTION] %s at %s (%d bytes)\n",
            func.getName(),
            func.getAddress(),
            Objects.nonNull(func.getFunctionSize()) ? func.getFunctionSize() : 0));
    if (Objects.nonNull(func.getSignature()) && !func.getSignature().isBlank()) {
      sb.append("Signature: ").append(func.getSignature()).append("\n");
    }
    sb.append(
        String.format(
            "Binary: %s | %s | %s | %s\n",
            AiUtils.orUnknown(binary.getFilename()),
            AiUtils.orUnknown(binary.getArchitecture()),
            AiUtils.orUnknown(binary.getCompiler()),
            AiUtils.orUnknown(binary.getFormat())));

    var callers = ctx.nameToCallers.getOrDefault(func.getName(), List.of());
    if (!callers.isEmpty()) {
      sb.append("Callers: ").append(String.join(", ", callers)).append("\n");
    }

    var callees = ctx.nameToCallees.getOrDefault(func.getName(), List.of());
    if (!callees.isEmpty()) {
      sb.append("Callees: ").append(String.join(", ", callees)).append("\n");
    }

    var labels = ctx.funcToLabels.getOrDefault(func.getId(), List.of());
    if (!labels.isEmpty()) {
      sb.append("Labels:\n");
      for (var l : labels) {
        sb.append("  ").append(l).append("\n");
      }
    }

    sb.append("\n");
    if (Objects.nonNull(func.getFlags()) && !func.getFlags().isBlank()) {
      sb.append("Flags: ").append(func.getFlags()).append("\n");
    }

    if (Objects.nonNull(func.getDecompiledCode()) && !func.getDecompiledCode().isBlank()) {
      sb.append("[DECOMPILED]\n")
          .append(AiUtils.truncate(func.getDecompiledCode(), chunkMaxChars / 2))
          .append("\n");
    }

    if (Objects.nonNull(func.getStackFrame()) && !func.getStackFrame().isBlank()) {
      sb.append("[STACK FRAME]\n")
          .append(AiUtils.truncate(func.getStackFrame(), chunkMaxChars / 4))
          .append("\n");
    }

    if (Objects.nonNull(func.getComments()) && !func.getComments().isBlank()) {
      sb.append("[COMMENTS]\n")
          .append(AiUtils.truncate(func.getComments(), chunkMaxChars / 4))
          .append("\n");
    }

    if (Objects.nonNull(func.getCfgBlocks())
        && !func.getCfgBlocks().isBlank()
        && !"[]".equals(func.getCfgBlocks().trim())) {
      sb.append("[CFG] ")
          .append(AiUtils.truncate(func.getCfgBlocks(), chunkMaxChars / 4))
          .append("\n");
    }

    sb.append("[ASSEMBLY]\n").append(AiUtils.truncate(func.getAssembly(), chunkMaxChars));

    return sb.toString();
  }

  private Map<String, Object> buildFunctionMetadata(StaticFunction func, FunctionGraphContext ctx) {
    Map<String, Object> m = new HashMap<>();
    m.put("function_name", func.getName());
    m.put("address", func.getAddress());
    m.put("callers", ctx.nameToCallers.getOrDefault(func.getName(), List.of()));
    m.put("callees", ctx.nameToCallees.getOrDefault(func.getName(), List.of()));
    m.put("label_count", ctx.funcToLabels.getOrDefault(func.getId(), List.of()).size());
    m.put("assembly_length", Objects.nonNull(func.getAssembly()) ? func.getAssembly().length() : 0);
    m.put("flags", Objects.nonNull(func.getFlags()) ? func.getFlags() : "");
    m.put(
        "has_decompiled_code",
        Objects.nonNull(func.getDecompiledCode()) && !func.getDecompiledCode().isBlank());
    return m;
  }

  private void collectLibraryInfo(LibraryDiagnostics diagnostics, StaticImport imp) {
    var sym = imp.getLibrarySymbol();
    if (Objects.isNull(sym)) {
      return;
    }

    var lib = sym.getLibrary();
    if (Objects.nonNull(lib)) {
      diagnostics.libraryNames.add(lib.getName());
    }

    if (Objects.nonNull(sym.getClassName())) {
      diagnostics.classNames.add(sym.getClassName());
      var category = sym.getCategory();
      if (Objects.nonNull(category)) {
        diagnostics.categoryCounts.merge(category, 1, Integer::sum);
      }
    }
  }

  private void appendLibraryDiagnostics(
      StringBuilder sb, LibraryDiagnostics diagnostics, List<StaticImport> imports) {
    if (diagnostics.libraryNames.isEmpty()) {
      return;
    }

    var resolvedCount = imports.stream().filter(i -> Objects.nonNull(i.getResolvedName())).count();

    sb.append("\n[LIBRARY DIAGNOSTICS]\n");
    sb.append("Resolved imports: ")
        .append(resolvedCount)
        .append("/")
        .append(imports.size())
        .append("\n");

    sb.append("Detected libraries: ")
        .append(String.join(", ", diagnostics.libraryNames))
        .append("\n");

    if (!diagnostics.classNames.isEmpty()) {
      var classList = diagnostics.classNames.stream().sorted().toList();
      sb.append("Detected classes (")
          .append(classList.size())
          .append("): ")
          .append(String.join(", ", classList))
          .append("\n");
    }

    if (!diagnostics.categoryCounts.isEmpty()) {
      sb.append("Category distribution: ");
      var cats =
          diagnostics.categoryCounts.entrySet().stream()
              .sorted((a, b) -> b.getValue().compareTo(a.getValue()))
              .map(e -> e.getKey() + " (" + e.getValue() + ")")
              .toList();
      sb.append(String.join(", ", cats)).append("\n");
    }
  }

  private static class LibraryDiagnostics {
    final java.util.Set<String> libraryNames = new java.util.LinkedHashSet<>();
    final java.util.Set<String> classNames = new java.util.LinkedHashSet<>();
    final java.util.Map<String, Integer> categoryCounts = new java.util.HashMap<>();
  }

  record FunctionGraphContext(
      Map<String, List<String>> nameToCallers,
      Map<String, List<String>> nameToCallees,
      Map<UUID, List<String>> funcToLabels) {}
}
