#include "cst_agent.h"

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

typedef struct {
  const char *p;
  size_t n;
} line_t;

typedef struct {
  line_t *v;
  size_t n;
  size_t cap;
} lines_t;

static int lines_push(lines_t *L, const char *p, size_t n) {
  if (L->n == L->cap) {
    size_t nc = L->cap ? L->cap * 2 : 64;
    line_t *nv = realloc(L->v, nc * sizeof(*nv));
    if (!nv)
      return -1;
    L->v = nv;
    L->cap = nc;
  }
  L->v[L->n].p = p;
  L->v[L->n].n = n;
  L->n++;
  return 0;
}

/* Split into lines without '\n'. A trailing newline does not add an empty
 * final line. */
static int split_lines(const char *s, size_t n, lines_t *out) {
  size_t start = 0;
  for (size_t i = 0; i < n; i++) {
    if (s[i] == '\n') {
      if (lines_push(out, s + start, i - start) != 0)
        return -1;
      start = i + 1;
    }
  }
  if (start < n && lines_push(out, s + start, n - start) != 0)
    return -1;
  return 0;
}

static size_t rtrim_len(const line_t *l) {
  size_t n = l->n;
  while (n > 0 && (l->p[n - 1] == ' ' || l->p[n - 1] == '\t' ||
                   l->p[n - 1] == '\r'))
    n--;
  return n;
}

static int line_eq(const line_t *a, const line_t *b, int loose) {
  if (!loose)
    return a->n == b->n && memcmp(a->p, b->p, a->n) == 0;
  size_t na = rtrim_len(a), nb = rtrim_len(b);
  return na == nb && memcmp(a->p, b->p, na) == 0;
}

static int seq_at(const lines_t *file, size_t pos, const lines_t *seq,
                  int loose) {
  if (pos + seq->n > file->n)
    return 0;
  for (size_t i = 0; i < seq->n; i++) {
    if (!line_eq(&file->v[pos + i], &seq->v[i], loose))
      return 0;
  }
  return 1;
}

/* Find seq in file at or after min_pos. hint<0 means unknown. Returns
 * position, -1 when absent, -2 when ambiguous without a hint. */
static long locate(const lines_t *file, size_t min_pos, const lines_t *seq,
                   long hint, int loose) {
  if (seq->n > file->n)
    return -1;
  size_t last = file->n - seq->n;
  if (hint >= 0) {
    long best = -1;
    long best_d = 0;
    for (size_t p = min_pos; p <= last; p++) {
      if (!seq_at(file, p, seq, loose))
        continue;
      long d = (long)p - hint;
      if (d < 0)
        d = -d;
      if (best < 0 || d < best_d) {
        best = (long)p;
        best_d = d;
      }
    }
    return best;
  }
  long found = -1;
  for (size_t p = min_pos; p <= last; p++) {
    if (seq_at(file, p, seq, loose)) {
      if (found >= 0)
        return -2;
      found = (long)p;
    }
  }
  return found;
}

static int parse_hunk_header(const char *p, size_t n, long *a, long *b,
                             long *c, long *d) {
  /* "@@ -a[,b] +c[,d] @@..." */
  char tmp[128];
  if (n >= sizeof(tmp))
    n = sizeof(tmp) - 1;
  memcpy(tmp, p, n);
  tmp[n] = 0;
  *b = *d = 1;
  if (sscanf(tmp, "@@ -%ld,%ld +%ld,%ld", a, b, c, d) == 4)
    return 0;
  *b = *d = 1;
  if (sscanf(tmp, "@@ -%ld +%ld,%ld", a, c, d) == 3)
    return 0;
  *b = *d = 1;
  if (sscanf(tmp, "@@ -%ld,%ld +%ld", a, b, c) == 3)
    return 0;
  *b = *d = 1;
  if (sscanf(tmp, "@@ -%ld +%ld", a, c) == 2)
    return 0;
  return -1;
}

int cst_patch_apply_buffer(const char *orig, size_t olen, const char *hunks,
                           size_t hlen, char **out, size_t *out_len,
                           int *nhunks, char *err, size_t err_sz) {
  lines_t file = {0}, hl = {0};
  lines_t old_seq = {0}, new_seq = {0};
  int rc = -1;
  int count = 0;
  long delta = 0;
  size_t min_pos = 0;
  int orig_final_nl = olen == 0 || orig[olen - 1] == '\n';
  int final_nl = orig_final_nl;
  if (err && err_sz)
    err[0] = 0;
  if (out)
    *out = NULL;

  if (split_lines(orig ? orig : "", olen, &file) != 0 ||
      split_lines(hunks ? hunks : "", hlen, &hl) != 0) {
    snprintf(err, err_sz, "oom");
    goto done;
  }

  size_t i = 0;
  while (i < hl.n) {
    const line_t *h = &hl.v[i];
    if (!(h->n >= 2 && h->p[0] == '@' && h->p[1] == '@')) {
      i++;
      continue;
    }
    long a = -1, b = -1, c = -1, d = -1;
    int have_hint = parse_hunk_header(h->p, h->n, &a, &b, &c, &d) == 0;
    i++;
    old_seq.n = new_seq.n = 0;
    int last_kind = 0; /* ' ', '-', '+' of previous body line */
    int new_no_nl = 0, old_no_nl = 0;
    for (; i < hl.n; i++) {
      const line_t *l = &hl.v[i];
      if (l->n >= 2 && l->p[0] == '@' && l->p[1] == '@')
        break;
      if (l->n == 0) {
        /* models often drop the leading space of blank context lines; a
         * blank line at the very end of the patch is just a terminator */
        if (i + 1 == hl.n)
          break;
        lines_push(&old_seq, l->p, 0);
        lines_push(&new_seq, l->p, 0);
        last_kind = ' ';
        continue;
      }
      char tag = l->p[0];
      if (tag == ' ') {
        lines_push(&old_seq, l->p + 1, l->n - 1);
        lines_push(&new_seq, l->p + 1, l->n - 1);
      } else if (tag == '-') {
        if (l->n >= 4 && memcmp(l->p, "--- ", 4) == 0 && i + 1 < hl.n &&
            hl.v[i + 1].n >= 4 && memcmp(hl.v[i + 1].p, "+++ ", 4) == 0)
          break; /* next file header */
        lines_push(&old_seq, l->p + 1, l->n - 1);
      } else if (tag == '+') {
        lines_push(&new_seq, l->p + 1, l->n - 1);
      } else if (tag == '\\') {
        if (last_kind == '+')
          new_no_nl = 1;
        else if (last_kind == '-')
          old_no_nl = 1;
        else if (last_kind == ' ')
          new_no_nl = old_no_nl = 1;
        continue;
      } else {
        break; /* diff --git, index ..., prose */
      }
      last_kind = tag;
    }

    long hint = -1;
    if (have_hint) {
      /* "-a,0" inserts after line a; otherwise a is 1-based start */
      hint = (b == 0 ? a : a - 1) + delta;
      if (hint < 0)
        hint = 0;
    }

    long pos;
    if (old_seq.n == 0) {
      if (hint >= 0)
        pos = hint > (long)file.n ? (long)file.n : hint;
      else if (file.n == 0)
        pos = 0;
      else {
        snprintf(err, err_sz,
                 "hunk %d: pure insertion needs @@ -a,0 +c,d @@ line numbers",
                 count + 1);
        goto done;
      }
    } else {
      pos = locate(&file, min_pos, &old_seq, hint, 0);
      if (pos == -1)
        pos = locate(&file, min_pos, &old_seq, hint, 1);
      if (pos == -1 && min_pos > 0) {
        pos = locate(&file, 0, &old_seq, hint, 0);
        if (pos == -1)
          pos = locate(&file, 0, &old_seq, hint, 1);
      }
      if (pos == -2) {
        snprintf(err, err_sz,
                 "hunk %d: ambiguous, context matches several places; add "
                 "line numbers or more context",
                 count + 1);
        goto done;
      }
      if (pos < 0) {
        snprintf(err, err_sz,
                 "hunk %d: context/removed lines not found in file (first: "
                 "\"%.*s\")",
                 count + 1, (int)(old_seq.v[0].n > 80 ? 80 : old_seq.v[0].n),
                 old_seq.v[0].p);
        goto done;
      }
    }

    /* splice: file[pos .. pos+old) -> new_seq */
    size_t old_n = old_seq.n, new_n = new_seq.n;
    size_t need = file.n - old_n + new_n;
    if (need > file.cap) {
      line_t *nv = realloc(file.v, (need + 16) * sizeof(line_t));
      if (!nv) {
        snprintf(err, err_sz, "oom");
        goto done;
      }
      file.v = nv;
      file.cap = need + 16;
    }
    memmove(&file.v[pos + new_n], &file.v[pos + old_n],
            (file.n - (size_t)pos - old_n) * sizeof(line_t));
    for (size_t k = 0; k < new_n; k++)
      file.v[pos + k] = new_seq.v[k];
    int touches_end = (size_t)pos + new_n == need;
    file.n = need;
    delta += (long)new_n - (long)old_n;
    min_pos = (size_t)pos + new_n;
    if (touches_end) {
      if (new_no_nl)
        final_nl = 0;
      else if (old_no_nl)
        final_nl = 1;
    }
    count++;
  }

  if (count == 0) {
    snprintf(err, err_sz, "no @@ hunk found");
    goto done;
  }

  {
    size_t total = 1;
    for (size_t k = 0; k < file.n; k++)
      total += file.v[k].n + 1;
    char *buf = malloc(total);
    if (!buf) {
      snprintf(err, err_sz, "oom");
      goto done;
    }
    size_t o = 0;
    for (size_t k = 0; k < file.n; k++) {
      memcpy(buf + o, file.v[k].p, file.v[k].n);
      o += file.v[k].n;
      if (k + 1 < file.n || final_nl)
        buf[o++] = '\n';
    }
    buf[o] = 0;
    *out = buf;
    if (out_len)
      *out_len = o;
  }
  if (nhunks)
    *nhunks = count;
  rc = 0;

done:
  free(file.v);
  free(hl.v);
  free(old_seq.v);
  free(new_seq.v);
  return rc;
}