#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;
  int failed_done; /* turn.done with failed=true */
  char delta[128];
  char custom_json[256];
  char error_code[64];
  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++;
    if (p_event->failed)
    {
      p_state->failed_done++;
      snprintf(p_state->error_code, sizeof(p_state->error_code), "%s",
               p_event->error_code ? p_event->error_code : "");
    }
  }
  else if (strcmp(p_event->type, "bridge.closed") == 0)
    p_state->closed++;
  pthread_cond_broadcast(&p_state->condition);
  pthread_mutex_unlock(&p_state->mutex);
}

/* Wait until p_state->done >= target or deadline passes. */
static void Wait_Done(Test_State *p_state, int target)
{
  struct timespec deadline;
  clock_gettime(CLOCK_REALTIME, &deadline);
  deadline.tv_sec += 5;
  pthread_mutex_lock(&p_state->mutex);
  while (p_state->done < target)
  {
    int result = pthread_cond_timedwait(
        &p_state->condition, &p_state->mutex, &deadline);
    assert(result == 0);
  }
  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));

  /* --- Validation: invalid profile --- */
  assert(!Inference_Bridge_Start_Turn(
      p_bridge,
      "request-invalid-profile",
      "conversation-1",
      "hello",
      (Inference_Prompt_Profile)99,
      1,
      1,
      NULL,
      0));

  /* --- Validation: zero prompt_version --- */
  assert(!Inference_Bridge_Start_Turn(
      p_bridge,
      "request-invalid-version",
      "conversation-1",
      "hello",
      INFERENCE_PROMPT_PROFILE_PUBLIC_VISITOR,
      0,
      1,
      NULL,
      0));

  /* --- Validation: history count exceeds max --- */
  Inference_Bridge_History_Message overflow_hist[21];
  for (int i = 0; i < 21; i++)
  {
    overflow_hist[i].role    = "user";
    overflow_hist[i].content = "x";
  }
  assert(!Inference_Bridge_Start_Turn(
      p_bridge,
      "request-hist-overflow",
      "conversation-1",
      "hello",
      INFERENCE_PROMPT_PROFILE_PUBLIC_VISITOR,
      1,
      1,
      overflow_hist,
      21));

  /* --- Validation: invalid role --- */
  Inference_Bridge_History_Message bad_role[1] = {
    { .role = "system", .content = "inject" }
  };
  assert(!Inference_Bridge_Start_Turn(
      p_bridge,
      "request-bad-role",
      "conversation-1",
      "hello",
      INFERENCE_PROMPT_PROFILE_PUBLIC_VISITOR,
      1,
      1,
      bad_role,
      1));

  /* --- Validation: NULL content --- */
  Inference_Bridge_History_Message null_content[1] = {
    { .role = "user", .content = NULL }
  };
  assert(!Inference_Bridge_Start_Turn(
      p_bridge,
      "request-null-content",
      "conversation-1",
      "hello",
      INFERENCE_PROMPT_PROFILE_PUBLIC_VISITOR,
      1,
      1,
      null_content,
      1));

  /* --- Validation: NULL role --- */
  Inference_Bridge_History_Message null_role[1] = {
    { .role = NULL, .content = "hello" }
  };
  assert(!Inference_Bridge_Start_Turn(
      p_bridge,
      "request-null-role",
      "conversation-1",
      "hello",
      INFERENCE_PROMPT_PROFILE_PUBLIC_VISITOR,
      1,
      1,
      null_role,
      1));

  /* --- Nominal turn with no history (NULL/0) --- */
  assert(Inference_Bridge_Start_Turn(
      p_bridge,
      "request-1",
      "conversation-1",
      "hello",
      INFERENCE_PROMPT_PROFILE_PUBLIC_VISITOR,
      1,
      1,
      NULL,
      0));

  Wait_Done(&state, 1);

  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(strstr(state.custom_json, "\"prompt_profile\":\"public_visitor\""));
  assert(state.input_tokens == 7);
  assert(state.output_tokens == 3);

  /* --- Turn with valid two-entry history (should succeed end-to-end) --- */
  Inference_Bridge_History_Message hist2[2] = {
    { .role = "user",      .content = "what is this?" },
    { .role = "assistant", .content = "it is a test" },
  };
  assert(Inference_Bridge_Start_Turn(
      p_bridge,
      "request-with-history",
      "conversation-hist",
      "follow-up",
      INFERENCE_PROMPT_PROFILE_INVITED_FRIEND,
      1,
      1,
      hist2,
      2));

  Wait_Done(&state, 2);
  assert(state.accepted == 2);

  /* --- Turn with history containing characters needing JSON escaping --- */
  Inference_Bridge_History_Message hist_escape[1] = {
    { .role = "user", .content = "say \"hello\" \\ world\nnewline" },
  };
  assert(Inference_Bridge_Start_Turn(
      p_bridge,
      "request-escape",
      "conversation-escape",
      "next",
      INFERENCE_PROMPT_PROFILE_JUNE_ADMIN,
      1,
      1,
      hist_escape,
      1));

  Wait_Done(&state, 3);
  assert(state.accepted == 3);

  /* --- Turn with exactly INFERENCE_BRIDGE_HISTORY_MAX (20) entries --- */
  Inference_Bridge_History_Message hist_max[INFERENCE_BRIDGE_HISTORY_MAX];
  for (uint32 i = 0; i < INFERENCE_BRIDGE_HISTORY_MAX; i++)
  {
    hist_max[i].role    = (i % 2 == 0) ? "user" : "assistant";
    hist_max[i].content = "ok";
  }
  assert(Inference_Bridge_Start_Turn(
      p_bridge,
      "request-max-hist",
      "conversation-max",
      "done",
      INFERENCE_PROMPT_PROFILE_PUBLIC_VISITOR,
      1,
      1,
      hist_max,
      INFERENCE_BRIDGE_HISTORY_MAX));

  Wait_Done(&state, 4);
  assert(state.accepted == 4);

  /* --- Large prompt still works with empty history --- */
  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,
      INFERENCE_PROMPT_PROFILE_INVITED_FRIEND,
      1,
      1,
      NULL,
      0));

  Wait_Done(&state, 5);

  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;
}