package ai.decompile.ai.service.tools;

import ai.decompile.ai.service.AiUtils;
import ai.decompile.analysis.model.entity.StaticFunction;
import ai.decompile.analysis.model.repository.StaticAnalysisRepository;
import ai.decompile.analysis.model.repository.StaticFunctionRepository;
import ai.decompile.analysis.model.repository.StaticXrefRepository;
import java.util.ArrayDeque;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Objects;
import java.util.UUID;
import lombok.RequiredArgsConstructor;
import lombok.extern.log4j.Log4j2;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.data.domain.PageRequest;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;

@Log4j2
@Component
@RequiredArgsConstructor
@Transactional(readOnly = true)
@ConditionalOnProperty(name = "app.ai.enabled", havingValue = "true", matchIfMissing = true)
@SuppressWarnings("unused")
public class FunctionTools {
  private final StaticFunctionRepository functionRepository;
  private final StaticAnalysisRepository analysisRepository;
  private final StaticXrefRepository xrefRepository;

  private static final int MAX_CALLEES_PER_LEVEL = 10;
  private static final int MAX_TOTAL_FUNCTIONS = 20;
  private static final int MAX_ASSEMBLY_PER_FUNCTION = 5000;
  private static final int MAX_LIST_FUNCTIONS = 50;

  @Tool(
      description =
          """
      Find functions by name or address. Use this when the user mentions a specific
      function name (e.g., "sub_401000", "WinMain") or addresses (e.g., "0x401000").
      Returns up to 10 matching functions with their id, name, address, size and signature.
      """)
  public List<FunctionSummary> searchFunctions(
      @ToolParam(
              description =
                  "The function name or address to search for, extracted from the user's message "
                      + "(e.g. for 'explain WinMain' pass 'WinMain', for 'what is sub_401000' pass 'sub_401000')")
          String query,
      @ToolParam(
              description =
                  "The binary ID being analyzed, as shown in the [BINARY OVERVIEW] section. "
                      + "Always pass the exact ID shown as Binary ID.")
          String binaryId) {

    log.debug("searchFunctions called: query='{}', binaryId={}", query, binaryId);

    if (Objects.isNull(query) || query.isBlank()) {
      log.warn("searchFunctions: model called with null/blank query, returning empty list");
      return List.of();
    }

    if (Objects.isNull(binaryId) || binaryId.isBlank()) {
      log.warn("searchFunctions: model called with null/blank binaryId, returning empty list");
      return List.of();
    }

    UUID binaryUuid;

    try {
      binaryUuid = UUID.fromString(binaryId);
    } catch (IllegalArgumentException e) {
      log.warn("searchFunctions: invalid binaryId format '{}'", binaryId);
      return List.of();
    }

    var analysis =
        analysisRepository.findFirstByBinaryIdOrderByCreatedAtDesc(binaryUuid).orElse(null);

    if (Objects.isNull(analysis)) {
      log.debug("searchFunctions: no analysis found for binary {}", binaryUuid);
      return List.of();
    }

    var byName =
        functionRepository.findByNameIgnoreCaseContainingAndAnalysisId(query, analysis.getId());

    log.debug("searchFunctions: name search for '{}' returned {} results", query, byName.size());

    var results = new LinkedHashSet<>(byName);

    var addressQuery = query;

    if (addressQuery.startsWith("0x") || addressQuery.startsWith("0X")) {
      addressQuery = addressQuery.substring(2);
    }

    var byAddress =
        functionRepository.findByAddressIgnoreCaseContainingAndAnalysisId(
            addressQuery, analysis.getId());

    log.debug(
        "searchFunctions: address search for '{}' returned {} results",
        addressQuery,
        byAddress.size());
    results.addAll(byAddress);

    var limited = results.stream().limit(10).toList();

    log.debug(
        "searchFunctions: returning {} functions total (from {} combined)",
        limited.size(),
        results.size());

    return limited.stream().map(this::toSummary).toList();
  }

  @Tool(
      description =
          """
      Get the full assembly code of a single function. Use this when you need to
      analyze a function's implementation in detail — instruction by instruction.
      The assembly is truncated to 8000 characters.
      """)
  public FunctionAssembly getFunctionAssembly(
      @ToolParam(description = "The function ID (UUID from searchFunctions result)")
          String functionId) {

    log.debug("getFunctionAssembly called: functionId={}", functionId);

    var func = functionRepository.findById(UUID.fromString(functionId)).orElse(null);

    if (Objects.isNull(func)) {
      return new FunctionAssembly("unknown", "unknown", "Function not found: " + functionId);
    }

    return new FunctionAssembly(func.getName(), func.getAddress(), func.getAssembly());
  }

  @Tool(
      description =
          """
      Get the decompiled pseudocode of a function (C-like syntax from IDA Hex-Rays
      or Ghidra decompiler). Use this for high-level understanding of the function's
      logic without reading assembly. Returns empty if no decompiled code is available.
      """)
  public DecompiledResult getFunctionDecompiled(
      @ToolParam(description = "The function ID (UUID from searchFunctions result)")
          String functionId) {

    log.debug("getFunctionDecompiled called: functionId={}", functionId);

    var func = functionRepository.findById(UUID.fromString(functionId)).orElse(null);

    if (Objects.isNull(func)) {
      return new DecompiledResult("unknown", "unknown", "");
    }

    var code = func.getDecompiledCode();

    return new DecompiledResult(
        func.getName(), func.getAddress(), Objects.nonNull(code) ? code : "");
  }

  @Tool(
      description =
          """
      Get the assembly code of a function AND all functions it calls,
      recursively up to the specified depth. Use this to fully understand
      a function's implementation including its callees.
      Each level is indented, showing caller relationships clearly.
      """)
  public String getFunctionCallGraph(
      @ToolParam(description = "The function ID (UUID from searchFunctions result)")
          String functionId,
      @ToolParam(
              description =
                  "Maximum depth of callee recursion (1-3, default 3). "
                      + "Level 1 = direct callees only, Level 2 = callees of callees, etc.")
          int depth) {

    log.debug("getFunctionCallGraph called: functionId={}, depth={}", functionId, depth);

    var rootFunc = functionRepository.findById(UUID.fromString(functionId)).orElse(null);

    if (Objects.isNull(rootFunc)) {
      return "Function not found for id: " + functionId;
    }

    var safeDepth = Math.clamp(depth, 1, 3);
    var sb = new StringBuilder();
    var visited = new LinkedHashSet<UUID>();
    var totalFunctions = 0;

    var queue = new ArrayDeque<CallGraphNode>();
    queue.add(new CallGraphNode(rootFunc, 0));

    while (!queue.isEmpty() && totalFunctions < MAX_TOTAL_FUNCTIONS) {
      var node = queue.poll();
      var func = node.function();

      if (!visited.add(func.getId())) {
        continue;
      }

      appendFunctionBlock(sb, func, node.level(), visited.size() == 1);
      totalFunctions++;

      if (node.level() >= safeDepth) {
        continue;
      }

      var callees = xrefRepository.findByCallerId(func.getId());
      var added = 0;

      for (var xref : callees) {
        if (added >= MAX_CALLEES_PER_LEVEL) {
          break;
        }

        var callee = xref.getCallee();

        if (visited.contains(callee.getId())) {
          continue;
        }

        queue.add(new CallGraphNode(callee, node.level() + 1));
        added++;
      }
    }

    if (totalFunctions >= MAX_TOTAL_FUNCTIONS && !queue.isEmpty()) {
      sb.append("\n[Truncated: reached maximum of ")
          .append(MAX_TOTAL_FUNCTIONS)
          .append(" functions]");
    }

    log.debug("getFunctionCallGraph: returned {} functions (depth={})", totalFunctions, safeDepth);

    return sb.toString();
  }

  @Tool(
      description =
          """
      List functions in the binary, optionally filtered by name pattern.
      Use this when the user asks to "list all functions", wants to browse
      the function catalog, or search for functions matching a pattern.
      Returns up to 50 results with id, name, address, size and signature.
      """)
  public List<FunctionSummary> listFunctions(
      @ToolParam(
              description =
                  "Optional: name pattern to filter (substring match). Leave empty to list all functions.")
          String nameFilter,
      @ToolParam(
              description =
                  "The binary ID being analyzed, as shown in the [BINARY OVERVIEW] section.")
          String binaryId) {

    log.debug("listFunctions called: nameFilter='{}', binaryId={}", nameFilter, binaryId);

    UUID binaryUuid;

    try {
      binaryUuid = UUID.fromString(binaryId);
    } catch (IllegalArgumentException e) {
      return List.of();
    }

    var analysis =
        analysisRepository.findFirstByBinaryIdOrderByCreatedAtDesc(binaryUuid).orElse(null);

    if (Objects.isNull(analysis)) {
      return List.of();
    }

    if (Objects.nonNull(nameFilter) && !nameFilter.isBlank()) {
      return functionRepository
          .findByNameIgnoreCaseContainingAndAnalysisId(nameFilter, analysis.getId())
          .stream()
          .limit(MAX_LIST_FUNCTIONS)
          .map(this::toSummary)
          .toList();
    }

    var page =
        functionRepository.findByAnalysisId(
            analysis.getId(), PageRequest.of(0, MAX_LIST_FUNCTIONS));

    return page.getContent().stream().map(this::toSummary).toList();
  }

  private FunctionSummary toSummary(StaticFunction f) {
    return new FunctionSummary(
        f.getId(),
        f.getName(),
        f.getAddress(),
        Objects.nonNull(f.getFunctionSize()) ? f.getFunctionSize() : 0,
        Objects.nonNull(f.getSignature()) ? f.getSignature() : "");
  }

  private void appendFunctionBlock(
      StringBuilder sb, StaticFunction func, int level, boolean isRoot) {
    var indent = "  ".repeat(level);
    var labelPrefix = isRoot ? "[LEVEL " + level + "]" : "[LEVEL " + level + " → CALLED BY]";

    sb.append(indent)
        .append(labelPrefix)
        .append(" ")
        .append(func.getName())
        .append(" at ")
        .append(func.getAddress());

    if (Objects.nonNull(func.getFunctionSize()) && func.getFunctionSize() > 0) {
      sb.append(" (").append(func.getFunctionSize()).append(" bytes)");
    }

    sb.append("\n");

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

    if (Objects.nonNull(func.getAssembly()) && !func.getAssembly().isBlank()) {
      var truncated = AiUtils.truncate(func.getAssembly(), MAX_ASSEMBLY_PER_FUNCTION);
      sb.append(indent).append(truncated).append("\n");
    } else {
      sb.append(indent).append("[no assembly available]\n");
    }
  }

  private record CallGraphNode(StaticFunction function, int level) {}

  public record FunctionSummary(
      UUID id, String name, String address, int functionSize, String signature) {}

  public record DecompiledResult(String name, String address, String decompiledCode) {}

  public record FunctionAssembly(String name, String address, String assembly) {}
}
