#include "mm_tool.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define MM_OPEN "<minimax:tool_call>"
#define MM_CLOSE "</minimax:tool_call>"
#define MM_OPEN_LEN 19
#define MM_CLOSE_LEN 20

struct mm_scan {
  mm_on_event_fn on_event;
  void *userdata;
  char *buf;
  size_t len;
  size_t cap;
  size_t pending_peak;
  int in_tool; /* 1 if last open is unmatched */
  char idbuf[CTL_MAX_ID];
  unsigned id_seq;
};

static void emit(mm_scan_t *s, ctl_event_t *ev) {
  if (s->on_event)
    s->on_event(ev, s->userdata);
}

static void emit_text(mm_scan_t *s, const char *t, size_t n) {
  if (!t || !n)
    return;
  ctl_event_t ev = {0};
  ev.type = CTL_EV_TEXT;
  ev.text = t;
  ev.text_len = n;
  emit(s, &ev);
}

static int ensure(mm_scan_t *s, size_t need) {
  if (s->len + need + 1 <= s->cap)
    return 0;
  size_t nc = s->cap ? s->cap * 2 : 4096;
  while (nc < s->len + need + 1)
    nc *= 2;
  char *nb = realloc(s->buf, nc);
  if (!nb)
    return -1;
  s->buf = nb;
  s->cap = nc;
  return 0;
}

static void bump_pending(mm_scan_t *s) {
  if (s->len > s->pending_peak)
    s->pending_peak = s->len;
}

mm_scan_t *mm_scan_create(mm_on_event_fn on_event, void *userdata) {
  mm_scan_t *s = calloc(1, sizeof(*s));
  if (!s)
    return NULL;
  s->on_event = on_event;
  s->userdata = userdata;
  s->id_seq = 1;
  return s;
}

void mm_scan_destroy(mm_scan_t *s) {
  if (!s)
    return;
  free(s->buf);
  free(s);
}

void mm_scan_reset(mm_scan_t *s) {
  if (!s)
    return;
  s->len = 0;
  s->in_tool = 0;
  if (s->buf)
    s->buf[0] = 0;
}

size_t mm_scan_pending(const mm_scan_t *s) {
  return s ? (s->in_tool ? s->len : 0) : 0;
}

static void clear_call(ctl_tool_call_t *c) {
  for (int i = 0; i < c->nargs; i++) {
    free(c->args[i].value);
    c->args[i].value = NULL;
  }
  memset(c, 0, sizeof(*c));
}

static int add_arg(ctl_tool_call_t *c, const char *key, const char *val,
                   size_t vlen) {
  if (c->nargs >= CTL_MAX_ARGS)
    return -1;
  if (strlen(key) >= CTL_MAX_KEY)
    return -1;
  snprintf(c->args[c->nargs].key, CTL_MAX_KEY, "%s", key);
  c->args[c->nargs].value = malloc(vlen + 1);
  if (!c->args[c->nargs].value)
    return -1;
  memcpy(c->args[c->nargs].value, val, vlen);
  c->args[c->nargs].value[vlen] = 0;
  c->args[c->nargs].value_len = vlen;
  c->nargs++;
  return 0;
}

/* Extract name="..." or name='...' from a tag start. */
static int extract_name_attr(const char *tag, char *out, size_t out_sz) {
  const char *p = strstr(tag, "name=");
  if (!p)
    return -1;
  p += 5;
  char quote = 0;
  if (*p == '"' || *p == '\'') {
    quote = *p++;
  }
  size_t i = 0;
  while (*p && i + 1 < out_sz) {
    if (quote) {
      if (*p == quote)
        break;
    } else {
      if (*p == ' ' || *p == '>' || *p == '/')
        break;
    }
    out[i++] = *p++;
  }
  out[i] = 0;
  return i > 0 ? 0 : -1;
}

static void emit_call(mm_scan_t *s, ctl_tool_call_t *call) {
  /* Only emit allowlisted tools — prevents "cmd" etc. from becoming tools. */
  if (!call->name[0])
    return;
  if (strcmp(call->name, "read_file") != 0 &&
      strcmp(call->name, "run_shell") != 0 &&
      strcmp(call->name, "write_file") != 0 &&
      strcmp(call->name, "apply_patch") != 0)
    return;
  ctl_event_t ev = {0};
  ev.type = CTL_EV_CALL_OPEN;
  ev.call = call;
  emit(s, &ev);
  ev.type = CTL_EV_CALL_READY;
  ev.call = call;
  emit(s, &ev);
}

/*
 * Parse one complete tool_call inner body (between open/close tags).
 * Emits one CALL_READY per invoke.
 */
static void parse_tool_block(mm_scan_t *s, const char *inner, size_t n) {
  char *block = malloc(n + 1);
  if (!block)
    return;
  memcpy(block, inner, n);
  block[n] = 0;

  /* Normalize: support both <invoke name="x"> and bare "invoke x" lines. */
  char *p = block;
  while (*p) {
    /* Skip whitespace */
    while (*p == ' ' || *p == '\n' || *p == '\r' || *p == '\t')
      p++;
    if (!*p)
      break;

    ctl_tool_call_t call;
    memset(&call, 0, sizeof(call));
    snprintf(call.id, sizeof(call.id), "mm%u", s->id_seq++);

    if (strncmp(p, "<invoke", 7) == 0) {
      char *gt = strchr(p, '>');
      if (!gt) {
        break;
      }
      char tag[256];
      size_t tlen = (size_t)(gt - p) + 1;
      if (tlen >= sizeof(tag))
        tlen = sizeof(tag) - 1;
      memcpy(tag, p, tlen);
      tag[tlen] = 0;
      if (extract_name_attr(tag, call.name, sizeof(call.name)) != 0) {
        p = gt + 1;
        continue;
      }
      p = gt + 1;
      char *end_inv = strstr(p, "</invoke>");
      if (!end_inv) {
        /* unclosed invoke — treat rest as parameters until end */
        end_inv = block + n;
      }
      /* parameters / args inside */
      char *q = p;
      while (q < end_inv) {
        char *arg = strstr(q, "<parameter");
        char *arg2 = strstr(q, "<arg");
        char *use = NULL;
        int is_param = 0;
        if (arg && (!arg2 || arg <= arg2) && arg < end_inv) {
          use = arg;
          is_param = 1;
        } else if (arg2 && arg2 < end_inv) {
          use = arg2;
          is_param = 0;
        } else {
          break;
        }
        char *agt = strchr(use, '>');
        if (!agt || agt >= end_inv)
          break;
        char atag[256];
        size_t al = (size_t)(agt - use) + 1;
        if (al >= sizeof(atag))
          al = sizeof(atag) - 1;
        memcpy(atag, use, al);
        atag[al] = 0;
        char key[CTL_MAX_KEY];
        if (extract_name_attr(atag, key, sizeof(key)) != 0) {
          q = agt + 1;
          continue;
        }
        const char *close = is_param ? "</parameter>" : "</arg>";
        char *aend = strstr(agt + 1, close);
        if (!aend || aend > end_inv) {
          q = agt + 1;
          continue;
        }
        size_t vlen = (size_t)(aend - (agt + 1));
        add_arg(&call, key, agt + 1, vlen);
        q = aend + strlen(close);
      }
      emit_call(s, &call);
      clear_call(&call);
      if (end_inv < block + n && strncmp(end_inv, "</invoke>", 9) == 0)
        p = end_inv + 9;
      else
        p = end_inv;
      continue;
    }

    /* Bare: invoke NAME\n */
    if (strncmp(p, "invoke ", 7) == 0) {
      p += 7;
      while (*p == ' ' || *p == '\t')
        p++;
      size_t ni = 0;
      while (p[ni] && p[ni] != '\n' && p[ni] != '\r' && p[ni] != ' ' &&
             p[ni] != '<' && ni + 1 < sizeof(call.name)) {
        call.name[ni] = p[ni];
        ni++;
      }
      call.name[ni] = 0;
      p += ni;
      while (*p == '\n' || *p == '\r' || *p == ' ' || *p == '\t')
        p++;

      /* Collect <arg>...</arg> until next invoke or end */
      while (*p) {
        while (*p == ' ' || *p == '\n' || *p == '\r' || *p == '\t')
          p++;
        if (strncmp(p, "invoke ", 7) == 0 || strncmp(p, "<invoke", 7) == 0)
          break;
        if (strncmp(p, "<arg", 4) != 0 && strncmp(p, "<parameter", 10) != 0)
          break;
        int is_param = (strncmp(p, "<parameter", 10) == 0);
        char *agt = strchr(p, '>');
        if (!agt)
          break;
        char atag[256];
        size_t al = (size_t)(agt - p) + 1;
        if (al >= sizeof(atag))
          al = sizeof(atag) - 1;
        memcpy(atag, p, al);
        atag[al] = 0;
        char key[CTL_MAX_KEY];
        if (extract_name_attr(atag, key, sizeof(key)) != 0) {
          p = agt + 1;
          continue;
        }
        const char *close = is_param ? "</parameter>" : "</arg>";
        char *aend = strstr(agt + 1, close);
        if (!aend)
          break;
        size_t vlen = (size_t)(aend - (agt + 1));
        add_arg(&call, key, agt + 1, vlen);
        p = aend + strlen(close);
      }
      if (call.name[0]) {
        emit_call(s, &call);
      }
      clear_call(&call);
      continue;
    }

    /* Unknown content — skip one char */
    p++;
  }
  free(block);
}

/* Drop stray close tags that MiniMax sometimes emits outside a tool block. */
static void strip_stray_closes(mm_scan_t *s) {
  for (;;) {
    char *c = strstr(s->buf, MM_CLOSE);
    if (!c)
      return;
    /* Only strip if it appears before any open (orphan). */
    char *o = strstr(s->buf, MM_OPEN);
    if (o && o < c)
      return;
    size_t at = (size_t)(c - s->buf);
    size_t rest = s->len - at - MM_CLOSE_LEN;
    if (at > 0)
      emit_text(s, s->buf, at);
    memmove(s->buf, c + MM_CLOSE_LEN, rest);
    s->len = rest;
    s->buf[s->len] = 0;
  }
}

/* Process buffer: emit text before tools, parse complete tool blocks. */
static void drain(mm_scan_t *s) {
  if (!s->buf)
    return;
  for (;;) {
    if (!s->in_tool) {
      strip_stray_closes(s);
      char *open = strstr(s->buf, MM_OPEN);
      if (!open) {
        /* Keep a short tail that could be a partial open tag. */
        size_t keep = MM_OPEN_LEN - 1;
        if (s->len <= keep)
          return;
        size_t emit_n = s->len - keep;
        /* Don't split a UTF-8 codepoint across emit/keep. */
        while (emit_n > 0 &&
               ((unsigned char)s->buf[emit_n] & 0xC0) == 0x80)
          emit_n--;
        if (emit_n == 0)
          return;
        emit_text(s, s->buf, emit_n);
        memmove(s->buf, s->buf + emit_n, s->len - emit_n);
        s->len -= emit_n;
        s->buf[s->len] = 0;
        return;
      }
      size_t before = (size_t)(open - s->buf);
      if (before)
        emit_text(s, s->buf, before);
      /* drop open tag from buffer, enter tool mode */
      size_t rest = s->len - before - MM_OPEN_LEN;
      memmove(s->buf, open + MM_OPEN_LEN, rest);
      s->len = rest;
      s->buf[s->len] = 0;
      s->in_tool = 1;
      bump_pending(s);
      continue;
    }

    /* in tool: wait for close */
    char *close = strstr(s->buf, MM_CLOSE);
    if (!close)
      return; /* pending */
    size_t inner_len = (size_t)(close - s->buf);
    parse_tool_block(s, s->buf, inner_len);
    size_t rest = s->len - inner_len - MM_CLOSE_LEN;
    memmove(s->buf, close + MM_CLOSE_LEN, rest);
    s->len = rest;
    s->buf[s->len] = 0;
    s->in_tool = 0;
  }
}

int mm_scan_feed(mm_scan_t *s, const char *data, size_t len) {
  if (!s || !data)
    return -1;
  if (ensure(s, len) != 0)
    return -1;
  memcpy(s->buf + s->len, data, len);
  s->len += len;
  s->buf[s->len] = 0;
  bump_pending(s);
  drain(s);
  return 0;
}

int mm_scan_finish(mm_scan_t *s) {
  if (!s)
    return -1;
  drain(s);
  if (s->in_tool) {
    /* Incomplete tool — do NOT execute; emit as visible text so UI isn't stuck */
    emit_text(s, MM_OPEN, MM_OPEN_LEN);
    if (s->len)
      emit_text(s, s->buf, s->len);
    s->len = 0;
    s->in_tool = 0;
  } else if (s->len) {
    emit_text(s, s->buf, s->len);
    s->len = 0;
  }
  return 0;
}