#include "mrjunejune/inference_bridge.h"

#include <assert.h>
#include <pthread.h>
#include <stdio.h>
#include <string.h>
#include <time.h>

typedef struct {
  pthread_mutex_t mutex;
  pthread_cond_t condition;
  int accepted;
  int deltas;
  int completed;
  int usage;
  int custom;
  int done;
  int closed;
  char delta[128];
  char custom_json[256];
  int64 input_tokens;
  int64 output_tokens;
} Test_State;

static void Handle_Event(
    const Inference_Event *p_event,
    void *p_user_data)
{
  Test_State *p_state = p_user_data;
  pthread_mutex_lock(&p_state->mutex);
  if (strcmp(p_event->type, "turn.accepted") == 0)
    p_state->accepted++;
  else if (strcmp(p_event->type, "assistant.delta") == 0)
  {
    p_state->deltas++;
    snprintf(p_state->delta, sizeof(p_state->delta), "%s", p_event->delta);
  }
  else if (strcmp(p_event->type, "assistant.completed") == 0)
    p_state->completed++;
  else if (strcmp(p_event->type, "assistant.usage") == 0)
  {
    p_state->usage++;
    p_state->input_tokens = p_event->input_tokens;
    p_state->output_tokens = p_event->output_tokens;
  }
  else if (strcmp(p_event->type, "tool.started") == 0)
  {
    p_state->custom++;
    snprintf(
        p_state->custom_json,
        sizeof(p_state->custom_json),
        "%s",
        p_event->raw_json);
  }
  else if (strcmp(p_event->type, "turn.done") == 0 &&
           strncmp(p_event->request_id, "request-", 8) == 0)
  {
    p_state->done++;
  }
  else if (strcmp(p_event->type, "bridge.closed") == 0)
    p_state->closed++;
  pthread_cond_broadcast(&p_state->condition);
  pthread_mutex_unlock(&p_state->mutex);
}

int main(int argc, char **argv)
{
  assert(argc == 2);
  Test_State state = {0};
  assert(pthread_mutex_init(&state.mutex, NULL) == 0);
  assert(pthread_cond_init(&state.condition, NULL) == 0);

  Inference_Bridge *p_bridge = Inference_Bridge_Create(
      argv[1], argv[1], Handle_Event, &state);
  assert(p_bridge);
  assert(Inference_Bridge_Start(p_bridge));
  assert(Inference_Bridge_Is_Ready(p_bridge));
  assert(Inference_Bridge_Start_Turn(
      p_bridge, "request-1", "conversation-1", "hello"));

  struct timespec deadline;
  clock_gettime(CLOCK_REALTIME, &deadline);
  deadline.tv_sec += 5;
  pthread_mutex_lock(&state.mutex);
  while (state.done == 0)
  {
    int result = pthread_cond_timedwait(
        &state.condition, &state.mutex, &deadline);
    assert(result == 0);
  }
  pthread_mutex_unlock(&state.mutex);

  assert(state.accepted == 1);
  assert(state.deltas == 1);
  assert(state.completed == 1);
  assert(state.usage == 1);
  assert(state.custom == 1);
  assert(strcmp(state.delta, "hello \"traveler\"\\path") == 0);
  assert(strstr(state.custom_json, "\"name\":\"mock-search\""));
  assert(state.input_tokens == 7);
  assert(state.output_tokens == 3);

  char large_prompt[32 * 1024 + 1];
  memset(large_prompt, 0x01, sizeof(large_prompt) - 1);
  large_prompt[sizeof(large_prompt) - 1] = '\0';
  assert(Inference_Bridge_Start_Turn(
      p_bridge,
      "request-large",
      "conversation-1",
      large_prompt));
  clock_gettime(CLOCK_REALTIME, &deadline);
  deadline.tv_sec += 5;
  pthread_mutex_lock(&state.mutex);
  while (state.done < 2)
  {
    int result = pthread_cond_timedwait(
        &state.condition, &state.mutex, &deadline);
    assert(result == 0);
  }
  pthread_mutex_unlock(&state.mutex);

  Inference_Bridge_Destroy(p_bridge);
  assert(state.closed == 1);
  pthread_cond_destroy(&state.condition);
  pthread_mutex_destroy(&state.mutex);
  printf("Inference bridge tests passed\n");
  return 0;
}