#define _POSIX_C_SOURCE 200809L

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

static int append_text(char** buf, size_t* cap, size_t* used, const char* s, size_t len) {
    if (!buf || !cap || !used || !s) return -1;
    if (*used + len + 1U > *cap) {
        size_t next = *cap;
        while (*used + len + 1U > next) {
            next = (next == 0U) ? 256U : (next * 2U);
        }
        char* grown = (char*)realloc(*buf, next);
        if (!grown) return -1;
        *buf = grown;
        *cap = next;
    }
    memcpy(*buf + *used, s, len);
    *used += len;
    (*buf)[*used] = '\0';
    return 0;
}

int context_roles_split(const char* markdown, context_roles_t* out) {
    if (!out) return -1;
    memset(out, 0, sizeof(*out));

    if (!markdown || markdown[0] == '\0') {
        out->system_content = strdup("");
        out->user_content = strdup("");
        out->assistant_content = strdup("");
        return 0;
    }

    size_t sys_cap = 1024, sys_used = 0;
    size_t usr_cap = 1024, usr_used = 0;
    size_t ast_cap = 1024, ast_used = 0;

    char* sys_buf = (char*)malloc(sys_cap);
    char* usr_buf = (char*)malloc(usr_cap);
    char* ast_buf = (char*)malloc(ast_cap);

    if (!sys_buf || !usr_buf || !ast_buf) {
        free(sys_buf);
        free(usr_buf);
        free(ast_buf);
        return -1;
    }

    sys_buf[0] = '\0';
    usr_buf[0] = '\0';
    ast_buf[0] = '\0';

    char** current_buf = &sys_buf;
    size_t* current_cap = &sys_cap;
    size_t* current_used = &sys_used;

    const char* p = markdown;
    while (*p) {
        const char* eol = strchr(p, '\n');
        size_t line_len = eol ? (size_t)(eol - p) : strlen(p);

        if (strncmp(p, "system:", 7) == 0 && (line_len == 7 || (line_len == 8 && p[7] == '\r'))) {
            current_buf = &sys_buf;
            current_cap = &sys_cap;
            current_used = &sys_used;
        } else if (strncmp(p, "user:", 5) == 0 && (line_len == 5 || (line_len == 6 && p[5] == '\r'))) {
            current_buf = &usr_buf;
            current_cap = &usr_cap;
            current_used = &usr_used;
        } else if (strncmp(p, "assistant:", 10) == 0 && (line_len == 10 || (line_len == 11 && p[10] == '\r'))) {
            current_buf = &ast_buf;
            current_cap = &ast_cap;
            current_used = &ast_used;
        } else {
            if (append_text(current_buf, current_cap, current_used, p, line_len) != 0) {
                free(sys_buf);
                free(usr_buf);
                free(ast_buf);
                return -1;
            }
            if (eol) {
                if (append_text(current_buf, current_cap, current_used, "\n", 1) != 0) {
                    free(sys_buf);
                    free(usr_buf);
                    free(ast_buf);
                    return -1;
                }
            }
        }

        if (!eol) break;
        p = eol + 1;
    }

    // Trim trailing newlines
    while (sys_used > 0 && (sys_buf[sys_used - 1] == '\n' || sys_buf[sys_used - 1] == '\r')) {
        sys_buf[--sys_used] = '\0';
    }
    while (usr_used > 0 && (usr_buf[usr_used - 1] == '\n' || usr_buf[usr_used - 1] == '\r')) {
        usr_buf[--usr_used] = '\0';
    }
    while (ast_used > 0 && (ast_buf[ast_used - 1] == '\n' || ast_buf[ast_used - 1] == '\r')) {
        ast_buf[--ast_used] = '\0';
    }

    // Trim leading newlines
    char* sys_start = sys_buf;
    while (*sys_start == '\n' || *sys_start == '\r') sys_start++;
    if (sys_start > sys_buf) {
        memmove(sys_buf, sys_start, sys_used - (sys_start - sys_buf) + 1);
    }

    char* usr_start = usr_buf;
    while (*usr_start == '\n' || *usr_start == '\r') usr_start++;
    if (usr_start > usr_buf) {
        memmove(usr_buf, usr_start, usr_used - (usr_start - usr_buf) + 1);
    }

    char* ast_start = ast_buf;
    while (*ast_start == '\n' || *ast_start == '\r') ast_start++;
    if (ast_start > ast_buf) {
        memmove(ast_buf, ast_start, ast_used - (ast_start - ast_buf) + 1);
    }

    out->system_content = sys_buf;
    out->user_content = usr_buf;
    out->assistant_content = ast_buf;

    return 0;
}

void context_roles_free(context_roles_t* roles) {
    if (!roles) return;
    free(roles->system_content);
    free(roles->user_content);
    free(roles->assistant_content);
    memset(roles, 0, sizeof(*roles));
}
