#include "agent.h"

#include <curl/curl.h>
#include <cmath>
#include <limits>
#include <stdlib.h>
#include <string.h>

static char *trim_slashes(char *s) {
    size_t n = strlen(s);
    while (n && s[n - 1] == '/')
        s[--n] = '\0';
    return s;
}

void provider_probe_free(provider_probe *p) {
    free(p->detail);
    p->detail = NULL;
}

const char *agent_default_base_url(provider_kind p) {
    switch (p) {
    case PROV_OPENAI:
        return "https://api.openai.com/v1";
    case PROV_ANTHROPIC:
        return "https://api.anthropic.com";
    case PROV_OLLAMA:
        return "http://localhost:11434";
    }
    return "https://api.openai.com/v1";
}

char *agent_models_url(const config *cfg) {
    char *base = xstrdup(cfg->base_url ? cfg->base_url
                                       : agent_default_base_url(cfg->provider));
    trim_slashes(base);
    xbuf u;
    xbuf_init(&u);
    xbuf_put(&u, base);
    switch (cfg->provider) {
    case PROV_OPENAI:
        xbuf_push(&u, '/');
        xbuf_put(&u, "models");
        break;
    case PROV_ANTHROPIC: {
        /* The canonical base is api.anthropic.com; the /v1 prefix belongs to
         * the path. If the caller already included /v1 in base_url, don't
         * duplicate it (avoids .../v1/v1/models). */
        size_t blen = strlen(base);
        if (blen >= 3 && memcmp(base + blen - 3, "/v1", 3) == 0)
            xbuf_put(&u, "/models?limit=100");
        else
            xbuf_put(&u, "/v1/models?limit=100");
        break;
    }
    case PROV_OLLAMA: {
        xbuf_put(&u, "/api/show?name=");
        char *enc = cfg->model ? curl_easy_escape(NULL, cfg->model, 0) : NULL;
        xbuf_put(&u, enc ? enc : "");
        free(enc);
        break;
    }
    }
    char *out = xbuf_steal(&u);
    free(base);
    return out;
}

void agent_auth(const config *cfg, char **header, char **value) {
    if (header)
        *header = NULL;
    if (value)
        *value = NULL;
    switch (cfg->provider) {
    case PROV_OPENAI:
        if (cfg->api_key) {
            if (header)
                *header = xstrdup("Authorization");
            xbuf b;
            xbuf_init(&b);
            xbuf_fmt(&b, "Bearer %s", cfg->api_key);
            if (value)
                *value = xbuf_steal(&b);
            else
                xbuf_free(&b);
        }
        break;
    case PROV_ANTHROPIC:
        if (cfg->api_key) {
            if (header)
                *header = xstrdup("x-api-key");
            if (value)
                *value = xstrdup(cfg->api_key);
        }
        break;
    case PROV_OLLAMA:
        break;
    }
}

int agent_http_get(const char *url, const char *auth_header,
                   const char *auth_value, long *status, char **body,
                   char **err) {
    if (body)
        *body = NULL;
    if (err)
        *err = NULL;
    CURL *c = curl_easy_init();
    if (!c) {
        if (err)
            *err = xstrdup("could not initialize curl");
        return -1;
    }
    xbuf b;
    xbuf_init(&b);
    struct curl_slist *h = NULL;
    if (auth_header && auth_value) {
        xbuf hb;
        xbuf_init(&hb);
        xbuf_fmt(&hb, "%s: %s", auth_header, auth_value);
        h = curl_slist_append(h, hb.data);
        xbuf_free(&hb);
        curl_easy_setopt(c, CURLOPT_HTTPHEADER, h);
    }
    curl_easy_setopt(c, CURLOPT_URL, url);
    curl_easy_setopt(c, CURLOPT_TIMEOUT, 5L);
    curl_easy_setopt(c, CURLOPT_NOSIGNAL, 1L);
    const char *ca = llsh_ca_bundle();
    if (ca)
        curl_easy_setopt(c, CURLOPT_CAINFO, ca);
    curl_easy_setopt(c, CURLOPT_WRITEFUNCTION, xbuf_curl_write);
    curl_easy_setopt(c, CURLOPT_WRITEDATA, &b);
    CURLcode rc = curl_easy_perform(c);
    curl_slist_free_all(h);
    if (rc != CURLE_OK) {
        char *e = xstrdup(curl_easy_strerror(rc));
        free(b.data);
        curl_easy_cleanup(c);
        if (err)
            *err = e;
        return -1;
    }
    long code = 0;
    curl_easy_getinfo(c, CURLINFO_RESPONSE_CODE, &code);
    if (status)
        *status = code;
    if (body)
        *body = xbuf_steal(&b);
    else
        free(b.data);
    curl_easy_cleanup(c);
    return 0;
}

static bool str_eq(const cJSON *v, const char *s) {
    return cJSON_IsString(v) && v->valuestring &&
           strcmp(v->valuestring, s) == 0;
}

static unsigned long num_u64(const cJSON *v) {
    if (cJSON_IsNumber(v) && std::isfinite(v->valuedouble) &&
        v->valuedouble >= 0 && std::floor(v->valuedouble) == v->valuedouble &&
        v->valuedouble < (double)std::numeric_limits<unsigned long>::max())
        return (unsigned long)v->valuedouble;
    return 0;
}

static bool openai_model_matches_entry(const cJSON *entry, const char *model) {
    if (str_eq(cJSON_GetObjectItem(entry, "id"), model))
        return true;
    if (str_eq(cJSON_GetObjectItem(entry, "name"), model))
        return true;
    const cJSON *al = cJSON_GetObjectItem(entry, "aliases");
    if (cJSON_IsArray(al)) {
        const cJSON *x = NULL;
        cJSON_ArrayForEach(x, al)
            if (str_eq(x, model))
                return true;
    }
    return false;
}

bool openai_model_present(const char *body, const char *model) {
    cJSON *v = cJSON_Parse(body);
    if (!v)
        return false;
    bool found = false;
    cJSON *arr = cJSON_GetObjectItem(v, "data");
    if (!arr)
        arr = cJSON_GetObjectItem(v, "models");
    if (cJSON_IsArray(arr)) {
        const cJSON *m = NULL;
        cJSON_ArrayForEach(m, arr)
            if (openai_model_matches_entry(m, model)) {
                found = true;
                break;
            }
    }
    cJSON_Delete(v);
    return found;
}

bool anthropic_model_present(const char *body, const char *model) {
    cJSON *v = cJSON_Parse(body);
    if (!v)
        return false;
    bool found = false;
    const cJSON *arr = cJSON_GetObjectItem(v, "data");
    if (cJSON_IsArray(arr)) {
        const cJSON *m = NULL;
        cJSON_ArrayForEach(m, arr)
            if (str_eq(cJSON_GetObjectItem(m, "id"), model)) {
                found = true;
                break;
            }
    }
    cJSON_Delete(v);
    return found;
}

unsigned long extract_context_limit(const cJSON *json, const config *cfg) {
    if (!json)
        return 0;
    switch (cfg->provider) {
    case PROV_OPENAI: {
        unsigned long v = num_u64(cJSON_GetObjectItem(json, "context_length"));
        if (v)
            return v;
        const cJSON *data = cJSON_GetObjectItem(json, "data");
        if (!cJSON_IsArray(data))
            return 0;
        const cJSON *want = NULL;
        const cJSON *m = NULL;
        cJSON_ArrayForEach(m, data)
            if (openai_model_matches_entry(m, cfg->model)) {
                want = m;
                break;
            }
        if (!want && data->child)
            want = data->child;
        if (!want)
            return 0;
        v = num_u64(cJSON_GetObjectItem(want, "context_length"));
        if (v)
            return v;
        v = num_u64(cJSON_GetObjectItem(cJSON_GetObjectItem(want, "meta"),
                                        "n_ctx"));
        if (v)
            return v;
        return num_u64(cJSON_GetObjectItem(want, "context_window"));
    }
    case PROV_OLLAMA: {
        unsigned long v = num_u64(cJSON_GetObjectItem(
            cJSON_GetObjectItem(json, "model_details"), "context_length"));
        if (v)
            return v;
        const cJSON *mf = cJSON_GetObjectItem(json, "modelfiles");
        if (cJSON_IsArray(mf) && mf->child)
            v = num_u64(cJSON_GetObjectItem(
                cJSON_GetObjectItem(mf->child, "parameters"),
                "context_length"));
        return v;
    }
    case PROV_ANTHROPIC: {
        const cJSON *data = cJSON_GetObjectItem(json, "data");
        if (!cJSON_IsArray(data) || !data->child)
            return 0;
        const cJSON *want = NULL;
        const cJSON *m = NULL;
        cJSON_ArrayForEach(m, data)
            if (str_eq(cJSON_GetObjectItem(m, "id"), cfg->model)) {
                want = m;
                break;
            }
        if (!want)
            want = data->child;
        return num_u64(cJSON_GetObjectItem(want, "max_tokens"));
    }
    }
    return 0;
}

unsigned long agent_fetch_context_limit(const config *cfg) {
    char *url = agent_models_url(cfg);
    char *h = NULL, *v = NULL;
    agent_auth(cfg, &h, &v);
    long status = 0;
    char *body = NULL;
    char *err = NULL;
    int rc = agent_http_get(url, h, v, &status, &body, &err);
    free(h);
    free(v);
    free(url);
    free(err);
    if (rc != 0) {
        free(body);
        return 0;
    }
    unsigned long limit = 0;
    /* No status gate (parity): a non-2xx body just has to parse and carry
     * the field. */
    cJSON *j = cJSON_Parse(body);
    if (j) {
        limit = extract_context_limit(j, cfg);
        cJSON_Delete(j);
    }
    free(body);
    (void)status;
    return limit;
}

static const char *status_phrase(long code) {
    switch (code) {
    case 200: return "200 OK";
    case 201: return "201 Created";
    case 204: return "204 No Content";
    case 400: return "400 Bad Request";
    case 401: return "401 Unauthorized";
    case 403: return "403 Forbidden";
    case 404: return "404 Not Found";
    case 405: return "405 Method Not Allowed";
    case 408: return "408 Request Timeout";
    case 429: return "429 Too Many Requests";
    case 500: return "500 Internal Server Error";
    case 502: return "502 Bad Gateway";
    case 503: return "503 Service Unavailable";
    case 504: return "504 Gateway Timeout";
    default: {
        static char buf[32];
        snprintf(buf, sizeof buf, "%ld", code);
        return buf;
    }
    }
}

provider_probe agent_probe_provider(const config *cfg) {
    provider_probe p = {false, NULL};
    char *base = xstrdup(cfg->base_url ? cfg->base_url
                                       : agent_default_base_url(cfg->provider));
    trim_slashes(base);
    char *url = agent_models_url(cfg);
    char *h = NULL, *v = NULL;
    agent_auth(cfg, &h, &v);
    long status = 0;
    char *body = NULL;
    char *err = NULL;
    if (agent_http_get(url, h, v, &status, &body, &err) != 0) {
        xbuf fb;
        xbuf_init(&fb);
        xbuf_fmt(&fb, "could not reach %s (connection: %s); is the endpoint "
                      "up and reachable?",
                 base, err ? err : "unknown");
        free(err);
        free(h);
        free(v);
        free(url);
        free(base);
        free(body);
        p.detail = xbuf_steal(&fb);
        return p;
    }
    free(h);
    free(v);
    free(url);
    bool present = false;
    switch (cfg->provider) {
    case PROV_OPENAI:
        present = openai_model_present(body ? body : "", cfg->model);
        break;
    case PROV_ANTHROPIC:
        present = anthropic_model_present(body ? body : "", cfg->model);
        break;
    case PROV_OLLAMA:
        present = (status >= 200 && status < 300);
        break;
    }
    bool success = status >= 200 && status < 300;
    bool client_err = status >= 400 && status < 500;
    xbuf fb;
    xbuf_init(&fb);
    switch (cfg->provider) {
    case PROV_OPENAI:
        if (success) {
            if (present) {
                p.ok = true;
                xbuf_fmt(&fb, "model `%s` found in the provider list at %s",
                         cfg->model, base);
            } else {
                xbuf_fmt(&fb, "model `%s` not in the provider's model list "
                              "at %s",
                         cfg->model, base);
            }
        } else if (client_err) {
            xbuf_fmt(&fb, "provider returned %s for %s/models (check api_key)",
                     status_phrase(status), base);
        } else {
            xbuf_fmt(&fb, "provider returned %s for %s/models",
                     status_phrase(status), base);
        }
        break;
    case PROV_ANTHROPIC:
        if (success) {
            if (present) {
                p.ok = true;
                xbuf_fmt(&fb, "model `%s` found in the provider list",
                         cfg->model);
            } else {
                xbuf_fmt(&fb, "model `%s` not in the provider's model list "
                              "at %s",
                         cfg->model, base);
            }
        } else if (client_err) {
            xbuf_fmt(&fb, "provider returned %s for %s/v1/models (check "
                          "api_key)",
                     status_phrase(status), base);
        } else {
            xbuf_fmt(&fb, "provider returned %s for %s/v1/models",
                     status_phrase(status), base);
        }
        break;
    case PROV_OLLAMA:
        if (success) {
            p.ok = true;
            xbuf_fmt(&fb, "model `%s` found via Ollama", cfg->model);
        } else {
            xbuf_fmt(&fb, "Ollama returned %s at %s (model `%s`)",
                     status_phrase(status), base, cfg->model);
        }
        break;
    }
    free(body);
    p.detail = xbuf_steal(&fb);
    free(base);
    return p;
}