#define _POSIX_C_SOURCE 200809L

#include "prompt_template.h"

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

static char* dup_range(const char* s, size_t n) {
    char* out = (char*)malloc(n + 1U);
    if (!out) return NULL;
    if (n > 0) {
        memcpy(out, s, n);
    }
    out[n] = '\0';
    return out;
}

static char* ltrim_inplace(char* s) {
    if (!s) return s;
    while (*s && isspace((unsigned char)*s)) s++;
    return s;
}

static void rtrim_inplace(char* s) {
    if (!s) return;
    size_t n = strlen(s);
    while (n > 0 && isspace((unsigned char)s[n - 1])) {
        s[n - 1] = '\0';
        n--;
    }
}

static int starts_with(const char* s, const char* prefix) {
    if (!s || !prefix) return 0;
    size_t n = strlen(prefix);
    return strncmp(s, prefix, n) == 0;
}

static int add_line(char*** lines, int* count, int* cap, char* line) {
    if (!lines || !count || !cap) return -1;
    if (*count >= *cap) {
        int next = (*cap == 0) ? 64 : (*cap * 2);
        char** grown = (char**)realloc(*lines, (size_t)next * sizeof(char*));
        if (!grown) return -1;
        *lines = grown;
        *cap = next;
    }
    (*lines)[*count] = line;
    (*count)++;
    return 0;
}

static int split_lines_inplace(char* s, char*** out_lines, int* out_count) {
    if (!s || !out_lines || !out_count) return -1;

    char** lines = NULL;
    int count = 0;
    int cap = 0;

    char* p = s;
    while (*p) {
        char* start = p;
        while (*p && *p != '\n') p++;
        if (*p == '\n') {
            *p = '\0';
            p++;
        }
        if (add_line(&lines, &count, &cap, start) != 0) {
            free(lines);
            return -1;
        }
    }

    *out_lines = lines;
    *out_count = count;
    return 0;
}

static int append_text(char** buf, size_t* cap, size_t* used, const char* s) {
    if (!buf || !cap || !used || !s) return -1;
    size_t n = strlen(s);
    if (*used + n + 1U > *cap) {
        size_t next = *cap;
        while (*used + n + 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, n);
    *used += n;
    (*buf)[*used] = '\0';
    return 0;
}

static const char* map_variable_tool_name(const char* var_name) {
    if (!var_name || var_name[0] == '\0') return NULL;

    if (strcmp(var_name, "admin_profile") == 0) return "nostr_admin_profile";
    if (strcmp(var_name, "admin_notes") == 0) return "nostr_admin_notes";
    if (strcmp(var_name, "admin_relays") == 0) return "nostr_admin_relays";
    if (strcmp(var_name, "adopted_skills") == 0) return "adopted_skills";
    if (strcmp(var_name, "triggering_event") == 0) return "trigger_event";
    return NULL;
}

static char* resolve_inline_variables_local(const char* tpl, tools_context_t* tools_ctx) {
    if (!tpl) return strdup("");

    size_t cap = strlen(tpl) + 64U;
    size_t used = 0U;
    char* out = (char*)malloc(cap);
    if (!out) return NULL;
    out[0] = '\0';

    const char* p = tpl;
    while (*p) {
        const char* open = strstr(p, "{{");
        if (!open) {
            if (append_text(&out, &cap, &used, p) != 0) {
                free(out);
                return NULL;
            }
            break;
        }

        if (open > p) {
            char* prefix = dup_range(p, (size_t)(open - p));
            if (!prefix) {
                free(out);
                return NULL;
            }
            int rc = append_text(&out, &cap, &used, prefix);
            free(prefix);
            if (rc != 0) {
                free(out);
                return NULL;
            }
        }

        const char* close = strstr(open + 2, "}}");
        if (!close) {
            if (append_text(&out, &cap, &used, open) != 0) {
                free(out);
                return NULL;
            }
            break;
        }

        char* var = dup_range(open + 2, (size_t)(close - (open + 2)));
        if (!var) {
            free(out);
            return NULL;
        }
        char* var_trim = ltrim_inplace(var);
        rtrim_inplace(var_trim);

        const char* replacement = "";
        char* replacement_owned = NULL;

        if (strcmp(var_trim, "message") == 0) {
            replacement = (tools_ctx && tools_ctx->template_current_user_message)
                              ? tools_ctx->template_current_user_message
                              : "";
        } else if (strcmp(var_trim, "dm_history") == 0) {
            replacement = "";
        } else if (tools_ctx) {
            const char* tool_name = map_variable_tool_name(var_trim);
            if (tool_name && tool_name[0] != '\0') {
                char* tool_result = tools_execute(tools_ctx, tool_name, "{}");
                if (tool_result) {
                    cJSON* root = cJSON_Parse(tool_result);
                    if (root && cJSON_IsObject(root)) {
                        cJSON* success = cJSON_GetObjectItemCaseSensitive(root, "success");
                        if (!success || !cJSON_IsBool(success) || cJSON_IsTrue(success)) {
                            cJSON* content = cJSON_GetObjectItemCaseSensitive(root, "content");
                            if (content && cJSON_IsString(content) && content->valuestring) {
                                replacement_owned = strdup(content->valuestring);
                            } else if (content) {
                                replacement_owned = cJSON_PrintUnformatted(content);
                            }
                        }
                    }
                    cJSON_Delete(root);
                    free(tool_result);
                }
            }

            if (!replacement_owned && tools_ctx->template_skill_lookup && var_trim[0] != '\0') {
                const char* skill_content = tools_ctx->template_skill_lookup(
                    tools_ctx->template_skill_lookup_user_data,
                    var_trim);
                if (skill_content && skill_content[0] != '\0') {
                    replacement_owned = strdup(skill_content);
                }
            }

            replacement = replacement_owned ? replacement_owned : "";
        }

        if (append_text(&out, &cap, &used, replacement) != 0) {
            free(replacement_owned);
            free(var);
            free(out);
            return NULL;
        }

        free(replacement_owned);
        free(var);
        p = close + 2;
    }

    return out;
}

static void init_section_defaults(prompt_template_section_t* sec) {
    if (!sec) return;
    memset(sec->name, 0, sizeof(sec->name));
    memset(sec->role, 0, sizeof(sec->role));
    snprintf(sec->role, sizeof(sec->role), "system");
    sec->content_template = NULL;
    sec->tool_name = NULL;
    sec->tool_args = NULL;
    sec->result_field = NULL;
    sec->limit = 0;
    sec->skip_if_empty = 0;
    sec->provider_name = NULL;
    sec->provider_content_template = NULL;
}

static int parse_int_or_zero(const char* s) {
    if (!s) return 0;
    while (*s && isspace((unsigned char)*s)) s++;
    return atoi(s);
}

int prompt_template_parse(const char* skill_content, prompt_template_t* out_template) {
    if (!skill_content || !out_template) {
        return -1;
    }

    memset(out_template, 0, sizeof(*out_template));

    const char* marker = strstr(skill_content, PROMPT_TEMPLATE_MARKER);
    if (!marker) {
        return -1;
    }

    size_t personality_len = (size_t)(marker - skill_content);
    out_template->personality = dup_range(skill_content, personality_len);
    if (!out_template->personality) {
        return -1;
    }
    rtrim_inplace(out_template->personality);

    const char* after = marker + strlen(PROMPT_TEMPLATE_MARKER);
    while (*after == '\r' || *after == '\n') after++;

    char* tpl = strdup(after);
    if (!tpl) {
        prompt_template_free(out_template);
        return -1;
    }

    char** lines = NULL;
    int line_count = 0;
    if (split_lines_inplace(tpl, &lines, &line_count) != 0) {
        free(tpl);
        prompt_template_free(out_template);
        return -1;
    }

    int current = -1;
    int i = 0;
    while (i < line_count) {
        char* raw = lines[i];
        rtrim_inplace(raw);
        char* line = ltrim_inplace(raw);

        if (*line == '\0') {
            i++;
            continue;
        }

        if (starts_with(line, "- section:")) {
            if (out_template->section_count >= PROMPT_TEMPLATE_MAX_SECTIONS) {
                break;
            }

            current = out_template->section_count;
            init_section_defaults(&out_template->sections[current]);
            out_template->section_count++;

            char* name = line + strlen("- section:");
            name = ltrim_inplace(name);
            rtrim_inplace(name);
            snprintf(out_template->sections[current].name,
                     sizeof(out_template->sections[current].name),
                     "%s",
                     name);
            i++;
            continue;
        }

        if (current < 0) {
            i++;
            continue;
        }

        if (starts_with(line, "role:")) {
            char* role = line + strlen("role:");
            role = ltrim_inplace(role);
            rtrim_inplace(role);
            snprintf(out_template->sections[current].role,
                     sizeof(out_template->sections[current].role),
                     "%s",
                     (*role) ? role : "system");
            i++;
            continue;
        }

        if (starts_with(line, "limit:")) {
            char* lim = line + strlen("limit:");
            out_template->sections[current].limit = parse_int_or_zero(lim);
            i++;
            continue;
        }

        if (starts_with(line, "skip_if_empty:")) {
            char* flag = line + strlen("skip_if_empty:");
            flag = ltrim_inplace(flag);
            rtrim_inplace(flag);
            out_template->sections[current].skip_if_empty =
                (strcmp(flag, "true") == 0 || strcmp(flag, "1") == 0) ? 1 : 0;
            i++;
            continue;
        }

        if (starts_with(line, "tool:")) {
            char* val = line + strlen("tool:");
            val = ltrim_inplace(val);
            rtrim_inplace(val);
            free(out_template->sections[current].tool_name);
            out_template->sections[current].tool_name = strdup(val);
            i++;
            continue;
        }

        if (starts_with(line, "args:")) {
            char* val = line + strlen("args:");
            val = ltrim_inplace(val);
            rtrim_inplace(val);
            free(out_template->sections[current].tool_args);
            out_template->sections[current].tool_args = strdup((*val) ? val : "{}");
            i++;
            continue;
        }

        if (starts_with(line, "result_field:")) {
            char* val = line + strlen("result_field:");
            val = ltrim_inplace(val);
            rtrim_inplace(val);
            free(out_template->sections[current].result_field);
            out_template->sections[current].result_field = strdup(val);
            i++;
            continue;
        }

        if (starts_with(line, "content:")) {
            char* val = line + strlen("content:");
            val = ltrim_inplace(val);

            free(out_template->sections[current].content_template);
            out_template->sections[current].content_template = NULL;

            if (strcmp(val, "|") == 0) {
                char* acc = strdup("");
                size_t cap = acc ? 1U : 0U;
                size_t used = 0U;
                if (!acc) {
                    free(lines);
                    free(tpl);
                    prompt_template_free(out_template);
                    return -1;
                }

                i++;
                while (i < line_count) {
                    char* next_raw = lines[i];
                    char* next_ltrim = ltrim_inplace(next_raw);

                    if (starts_with(next_ltrim, "- section:") ||
                        starts_with(next_ltrim, "role:") ||
                        starts_with(next_ltrim, "limit:") ||
                        starts_with(next_ltrim, "skip_if_empty:") ||
                        starts_with(next_ltrim, "tool:") ||
                        starts_with(next_ltrim, "args:") ||
                        starts_with(next_ltrim, "result_field:") ||
                        starts_with(next_ltrim, "content:") ||
                        starts_with(next_ltrim, "provider:")) {
                        break;
                    }

                    char* piece = next_raw;
                    if (strncmp(piece, "    ", 4) == 0) piece += 4;
                    else if (strncmp(piece, "  ", 2) == 0) piece += 2;
                    rtrim_inplace(piece);

                    if (append_text(&acc, &cap, &used, piece) != 0 ||
                        append_text(&acc, &cap, &used, "\n") != 0) {
                        free(acc);
                        free(lines);
                        free(tpl);
                        prompt_template_free(out_template);
                        return -1;
                    }
                    i++;
                }

                rtrim_inplace(acc);
                out_template->sections[current].content_template = acc;
                continue;
            }

            rtrim_inplace(val);
            out_template->sections[current].content_template = strdup(val);
            i++;
            continue;
        }

        if (starts_with(line, "provider:")) {
            i++;
            if (i >= line_count) continue;

            char* provider_line = ltrim_inplace(lines[i]);
            rtrim_inplace(provider_line);
            char* colon = strchr(provider_line, ':');
            if (!colon) {
                continue;
            }

            *colon = '\0';
            char* provider_name = ltrim_inplace(provider_line);
            rtrim_inplace(provider_name);
            char* provider_val = ltrim_inplace(colon + 1);
            rtrim_inplace(provider_val);

            free(out_template->sections[current].provider_name);
            out_template->sections[current].provider_name = strdup(provider_name);
            free(out_template->sections[current].provider_content_template);
            out_template->sections[current].provider_content_template = NULL;

            if (strcmp(provider_val, "|") == 0) {
                char* acc = strdup("");
                size_t cap = acc ? 1U : 0U;
                size_t used = 0U;
                if (!acc) {
                    free(lines);
                    free(tpl);
                    prompt_template_free(out_template);
                    return -1;
                }

                i++;
                while (i < line_count) {
                    char* next = ltrim_inplace(lines[i]);
                    if (starts_with(next, "- section:") ||
                        starts_with(next, "role:") ||
                        starts_with(next, "limit:") ||
                        starts_with(next, "skip_if_empty:") ||
                        starts_with(next, "tool:") ||
                        starts_with(next, "args:") ||
                        starts_with(next, "result_field:") ||
                        starts_with(next, "content:") ||
                        starts_with(next, "provider:")) {
                        break;
                    }
                    char* piece = lines[i];
                    if (strncmp(piece, "    ", 4) == 0) piece += 4;
                    else if (strncmp(piece, "  ", 2) == 0) piece += 2;
                    rtrim_inplace(piece);
                    if (append_text(&acc, &cap, &used, piece) != 0 ||
                        append_text(&acc, &cap, &used, "\n") != 0) {
                        free(acc);
                        free(lines);
                        free(tpl);
                        prompt_template_free(out_template);
                        return -1;
                    }
                    i++;
                }
                rtrim_inplace(acc);
                out_template->sections[current].provider_content_template = acc;
                continue;
            }

            out_template->sections[current].provider_content_template = strdup(provider_val);
            i++;
            continue;
        }

        i++;
    }

    free(lines);
    free(tpl);
    return 0;
}

static int append_message_object(cJSON* messages, const char* role, const char* content) {
    if (!messages || !role) return -1;
    cJSON* msg = cJSON_CreateObject();
    if (!msg) return -1;
    cJSON_AddStringToObject(msg, "role", role);
    cJSON_AddStringToObject(msg, "content", content ? content : "");
    cJSON_AddItemToArray(messages, msg);
    return 0;
}

static char* extract_tool_result_content(const prompt_template_section_t* sec, const char* tool_result_json) {
    if (!tool_result_json) return strdup("");

    cJSON* root = cJSON_Parse(tool_result_json);
    if (!root || !cJSON_IsObject(root)) {
        cJSON_Delete(root);
        return strdup(tool_result_json);
    }

    cJSON* success = cJSON_GetObjectItemCaseSensitive(root, "success");
    if (success && cJSON_IsBool(success) && !cJSON_IsTrue(success)) {
        cJSON_Delete(root);
        return strdup("");
    }

    const char* field_name = (sec && sec->result_field && sec->result_field[0] != '\0')
                                 ? sec->result_field
                                 : "content";

    cJSON* field = cJSON_GetObjectItemCaseSensitive(root, field_name);
    if (!field) {
        field = cJSON_GetObjectItemCaseSensitive(root, "content");
    }

    char* out = NULL;
    if (field && cJSON_IsString(field) && field->valuestring) {
        out = strdup(field->valuestring);
    } else if (field) {
        out = cJSON_PrintUnformatted(field);
    } else {
        out = cJSON_PrintUnformatted(root);
    }

    cJSON_Delete(root);
    return out ? out : strdup("");
}

cJSON* prompt_template_build_messages(const prompt_template_t* tmpl,
                                      const char* provider_name,
                                      tools_context_t* tools_ctx,
                                      cJSON* dm_history_messages,
                                      int dm_history_default_limit,
                                      prompt_template_emit_hook_fn emit_hook,
                                      void* emit_hook_user_data) {
    if (!tmpl) return NULL;

    cJSON* out = cJSON_CreateArray();
    if (!out) return NULL;

    int out_idx = 0;
    if (tmpl->personality && tmpl->personality[0] != '\0') {
        if (append_message_object(out, "system", tmpl->personality) != 0) {
            cJSON_Delete(out);
            return NULL;
        }
        if (emit_hook) emit_hook("system_prompt", out_idx, emit_hook_user_data);
        out_idx++;
    }

    for (int i = 0; i < tmpl->section_count; i++) {
        const prompt_template_section_t* sec = &tmpl->sections[i];
        const char* role = sec->role[0] ? sec->role : "system";

        if (strcmp(role, "expand") == 0) {
            if (!dm_history_messages || !cJSON_IsArray(dm_history_messages)) {
                continue;
            }
            int total = cJSON_GetArraySize(dm_history_messages);
            int lim = sec->limit > 0 ? sec->limit : dm_history_default_limit;
            int start = (lim > 0 && total > lim) ? (total - lim) : 0;

            for (int j = start; j < total; j++) {
                cJSON* item = cJSON_GetArrayItem(dm_history_messages, j);
                if (!item || !cJSON_IsObject(item)) continue;
                cJSON* dup = cJSON_Duplicate(item, 1);
                if (!dup) {
                    cJSON_Delete(out);
                    return NULL;
                }
                cJSON_AddItemToArray(out, dup);
                if (emit_hook) emit_hook(sec->name[0] ? sec->name : "context_part", out_idx, emit_hook_user_data);
                out_idx++;
            }
            continue;
        }

        char* resolved = NULL;

        if (sec->tool_name && sec->tool_name[0] != '\0' && tools_ctx) {
            const char* args_json = (sec->tool_args && sec->tool_args[0] != '\0') ? sec->tool_args : "{}";
            char* tool_result = tools_execute(tools_ctx, sec->tool_name, args_json);
            resolved = extract_tool_result_content(sec, tool_result);
            free(tool_result);
        } else {
            const char* tpl = sec->content_template;
            if (provider_name && sec->provider_name && sec->provider_content_template &&
                strcmp(provider_name, sec->provider_name) == 0) {
                tpl = sec->provider_content_template;
            }
            resolved = resolve_inline_variables_local(tpl ? tpl : "", tools_ctx);
        }
        if (!resolved) {
            cJSON_Delete(out);
            return NULL;
        }

        if (sec->skip_if_empty) {
            char* chk = ltrim_inplace(resolved);
            if (chk && *chk == '\0') {
                free(resolved);
                continue;
            }
        }

        if (append_message_object(out, role, resolved) != 0) {
            free(resolved);
            cJSON_Delete(out);
            return NULL;
        }
        if (emit_hook) emit_hook(sec->name[0] ? sec->name : "context_part", out_idx, emit_hook_user_data);
        out_idx++;
        free(resolved);
    }

    return out;
}

void prompt_template_free(prompt_template_t* tmpl) {
    if (!tmpl) return;

    free(tmpl->personality);
    tmpl->personality = NULL;

    for (int i = 0; i < tmpl->section_count; i++) {
        free(tmpl->sections[i].content_template);
        tmpl->sections[i].content_template = NULL;
        free(tmpl->sections[i].tool_name);
        tmpl->sections[i].tool_name = NULL;
        free(tmpl->sections[i].tool_args);
        tmpl->sections[i].tool_args = NULL;
        free(tmpl->sections[i].result_field);
        tmpl->sections[i].result_field = NULL;
        free(tmpl->sections[i].provider_name);
        tmpl->sections[i].provider_name = NULL;
        free(tmpl->sections[i].provider_content_template);
        tmpl->sections[i].provider_content_template = NULL;
        tmpl->sections[i].name[0] = '\0';
        tmpl->sections[i].role[0] = '\0';
        tmpl->sections[i].limit = 0;
        tmpl->sections[i].skip_if_empty = 0;
    }

    tmpl->section_count = 0;
}