package ai.decompile.ai.service;

import ai.decompile.ai.model.dto.StreamEvent;
import ai.decompile.ai.model.entity.ChatSession;
import ai.decompile.ai.service.tools.BinaryTools;
import ai.decompile.ai.service.tools.FunctionTools;
import ai.decompile.ai.service.tools.ImportTools;
import ai.decompile.ai.service.tools.StringTools;
import ai.decompile.ai.service.tools.XrefTools;
import ai.decompile.workspace.model.entity.Binary;
import ai.decompile.workspace.service.BinaryService;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.util.List;
import java.util.Objects;
import java.util.UUID;
import lombok.RequiredArgsConstructor;
import lombok.extern.log4j.Log4j2;
import org.jspecify.annotations.NonNull;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.support.ToolCallbacks;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.definition.ToolDefinition;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;

@Log4j2
@Service
@RequiredArgsConstructor
@ConditionalOnProperty(name = "app.ai.enabled", havingValue = "true", matchIfMissing = true)
public class ChatService {
  private final StreamingModelProvider streamingModelProvider;
  private final ChatSessionService sessionService;
  private final ChatMessageConverter messageConverter;
  private final ChatContextBuilder contextBuilder;
  private final BinaryService binaryService;
  private final FunctionTools functionTools;
  private final XrefTools xrefTools;
  private final StringTools stringTools;
  private final ImportTools importTools;
  private final BinaryTools binaryTools;
  private final ObjectMapper objectMapper;

  private static final int SUMMARY_MAX_LENGTH = 120;

  public Flux<StreamEvent> chat(UUID sessionId, String userMessageContent) {
    var session = sessionService.getSession(sessionId);
    var binary = binaryService.getBinary(session.getBinaryId());

    sessionService.saveMessage(session, "USER", userMessageContent);

    var messagesList = buildChatMessagesList(session, binary, userMessageContent);
    var streamingModel = streamingModelProvider.resolve();
    var assistantContent = new StringBuilder();

    if (streamingModel instanceof ChatModel chatModel) {
      Sinks.Many<StreamEvent> eventSink = Sinks.many().multicast().onBackpressureBuffer(64, false);
      var observedCallbacks =
          wrapToolCallbacks(
              ToolCallbacks.from(functionTools, xrefTools, stringTools, importTools, binaryTools),
              eventSink);

      Flux<StreamEvent> contentFlux =
          ChatClient.create(chatModel)
              .prompt(new Prompt(messagesList))
              .tools((Object[]) observedCallbacks)
              .stream()
              .content()
              .map(content -> (StreamEvent) new StreamEvent.Chunk(content))
              .doOnNext(chunk -> assistantContent.append(((StreamEvent.Chunk) chunk).content()))
              .doOnComplete(
                  () -> {
                    log.debug("Chat completed for session {}", sessionId);
                    saveAssistantMessage(session, assistantContent.toString());
                    eventSink.tryEmitComplete();
                  })
              .doOnError(
                  e -> {
                    log.error("Chat stream error for session {}", sessionId, e);
                    eventSink.tryEmitComplete();
                  });

      return Flux.merge(contentFlux, eventSink.asFlux())
          .concatWith(Mono.just(new StreamEvent.Done(sessionId.toString())));
    }

    return Flux.just(
        new StreamEvent.Error(
            "Model does not support ChatModel. Tool calling requires ChatModel."));
  }

  private ToolCallback[] wrapToolCallbacks(
      ToolCallback[] callbacks, Sinks.Many<StreamEvent> eventSink) {
    var wrapped = new ToolCallback[callbacks.length];
    for (var i = 0; i < callbacks.length; i++) {
      wrapped[i] = new ObservedToolCallback(callbacks[i], eventSink);
    }
    return wrapped;
  }

  private List<Message> buildChatMessagesList(
      ChatSession session, Binary binary, String userMessageContent) {
    var historyMessages = sessionService.getMessages(session.getId());

    var context = contextBuilder.buildContext(session.getBinaryId(), userMessageContent);
    var systemPrompt = contextBuilder.buildSystemPrompt(context, binary);

    return messageConverter.buildMessageList(
        systemPrompt,
        messageConverter.getRecentHistoryExcludingLastMessage(historyMessages),
        userMessageContent);
  }

  private void saveAssistantMessage(ChatSession session, String fullContent) {
    if (Objects.nonNull(fullContent) && !fullContent.isBlank()) {
      sessionService.saveMessage(session, "ASSISTANT", fullContent);
    }
  }

  private class ObservedToolCallback implements ToolCallback {
    private final ToolCallback delegate;
    private final Sinks.Many<StreamEvent> eventSink;

    ObservedToolCallback(ToolCallback delegate, Sinks.Many<StreamEvent> eventSink) {
      this.delegate = delegate;
      this.eventSink = eventSink;
    }

    @Override
    public @NonNull ToolDefinition getToolDefinition() {
      return delegate.getToolDefinition();
    }

    @Override
    public @NonNull String call(@NonNull String input) {
      var toolName = delegate.getToolDefinition().name();
      eventSink.tryEmitNext(new StreamEvent.ToolCallStart(toolName));
      var start = System.currentTimeMillis();
      try {
        var result = delegate.call(input);
        var duration = System.currentTimeMillis() - start;
        eventSink.tryEmitNext(
            new StreamEvent.ToolCallEnd(toolName, duration, buildSummary(result)));
        return result;
      } catch (Exception e) {
        var duration = System.currentTimeMillis() - start;
        eventSink.tryEmitNext(
            new StreamEvent.ToolCallEnd(toolName, duration, "Error: " + e.getMessage()));
        throw e;
      }
    }

    private String buildSummary(String rawResult) {
      if (Objects.isNull(rawResult) || rawResult.isBlank()) {
        return "No results";
      }

      try {
        var node = objectMapper.readTree(rawResult);
        if (node.isArray()) {
          var count = node.size();
          if (count == 0) {
            return "0 results";
          }

          var sb = new StringBuilder();

          sb.append(count);
          sb.append(" result");

          if (count > 1) {
            sb.append("s");
          }

          var first = node.get(0);
          if (first.has("name")) {
            sb.append(": ");
            sb.append(first.get("name").asText());
          } else if (first.has("functionName")) {
            sb.append(": ");
            sb.append(first.get("functionName").asText());
          } else if (first.has("address")) {
            sb.append(": ");
            sb.append(first.get("address").asText());
          }

          if (count > 1) {
            sb.append(", ...");
          }

          return truncate(sb.toString());
        }

        if (node.has("name") && node.has("address")) {
          return truncate(node.get("name").asText() + " at " + node.get("address").asText());
        }

        if (node.has("overview")) {
          return truncate(node.get("overview").asText());
        }

        if (node.has("summary")) {
          return truncate(node.get("summary").asText());
        }

        return truncate(rawResult);
      } catch (Exception e) {
        return truncate(rawResult);
      }
    }

    private String truncate(String text) {
      if (text.length() <= SUMMARY_MAX_LENGTH) {
        return text;
      }
      return text.substring(0, SUMMARY_MAX_LENGTH - 3) + "...";
    }
  }
}
