#include "commands.h"
#include "cmd_util.h"

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

#include "audit.h"
#include "auth.h"
#include "config.h"
#include "db.h"
#include "sessions.h"
#include "util.h"
#include "version.h"

/* ------------------------------------------------------------------ */
/* login rate limiting (in-memory, per key)                            */
/* ------------------------------------------------------------------ */

#define RL_MAX_KEYS 16
#define RL_MAX_FAILS 5
#define RL_WINDOW 900

struct rl_entry {
    char key[64];
    int fails;
    int64_t window_end;
};

static struct rl_entry g_rl[RL_MAX_KEYS];

static struct rl_entry *rl_get(const char *key, int create)
{
    struct rl_entry *slot = NULL;
    for (size_t i = 0; i < RL_MAX_KEYS; i++) {
        if (g_rl[i].key[0] && strcmp(g_rl[i].key, key) == 0)
            return &g_rl[i];
        if (create && !g_rl[i].key[0] && !slot)
            slot = &g_rl[i];
    }
    if (create && slot) {
        snprintf(slot->key, sizeof slot->key, "%s", key);
        slot->fails = 0;
        slot->window_end = 0;
    }
    return slot;
}

static int rl_blocked(const char *key, int64_t *retry_after)
{
    struct rl_entry *e = rl_get(key, 0);
    if (!e || e->fails < RL_MAX_FAILS)
        return 0;
    int64_t now = util_now();
    if (e->window_end <= now)
        return 0;
    if (retry_after)
        *retry_after = e->window_end - now;
    return 1;
}

static void rl_fail(const char *key)
{
    struct rl_entry *e = rl_get(key, 1);
    if (!e)
        return;
    int64_t now = util_now();
    if (e->fails == 0 || e->window_end <= now)
        e->window_end = now + RL_WINDOW;
    e->fails++;
}

static void rl_ok(const char *key)
{
    struct rl_entry *e = rl_get(key, 0);
    if (e) {
        e->fails = 0;
        e->window_end = 0;
    }
}

/* ------------------------------------------------------------------ */
/* public commands                                                     */
/* ------------------------------------------------------------------ */

static yyjson_mut_val *h_health(struct req *r)
{
    yyjson_mut_val *o = yyjson_mut_obj(r->rdoc);
    yyjson_mut_obj_add_strcpy(r->rdoc, o, "status", "ok");
    return o;
}

static yyjson_mut_val *h_meta(struct req *r)
{
    yyjson_mut_val *o = yyjson_mut_obj(r->rdoc);
    yyjson_mut_obj_add_strcpy(r->rdoc, o, "server", "bokfd");
    yyjson_mut_obj_add_strcpy(r->rdoc, o, "version", BOKF_VERSION);
    yyjson_mut_obj_add_int(r->rdoc, o, "protocol", BOKF_PROTOCOL_VERSION);
    yyjson_mut_val *features = yyjson_mut_arr(r->rdoc);
    yyjson_mut_arr_add_strcpy(r->rdoc, features, "describe");
    yyjson_mut_arr_add_strcpy(r->rdoc, features, "agent.instructions");
    yyjson_mut_arr_add_strcpy(r->rdoc, features, "orgs");
    yyjson_mut_arr_add_strcpy(r->rdoc, features, "tokens");
    yyjson_mut_arr_add_strcpy(r->rdoc, features, "backup");
    yyjson_mut_obj_add_val(r->rdoc, o, "features", features);
    yyjson_mut_val *limits = yyjson_mut_obj(r->rdoc);
    yyjson_mut_obj_add_int(r->rdoc, limits, "max_line_bytes",
                           g_cfg.max_line_bytes);
    yyjson_mut_obj_add_int(r->rdoc, limits, "max_attachment_bytes",
                           g_cfg.max_attachment_bytes);
    yyjson_mut_obj_add_int(r->rdoc, limits, "session_ttl_seconds",
                           g_cfg.session_ttl);
    yyjson_mut_obj_add_val(r->rdoc, o, "limits", limits);
    yyjson_mut_obj_add_bool(r->rdoc, o, "tcp_enabled", g_cfg.tcp_enabled != 0);
    char ts[32];
    util_iso8601(util_now(), ts, sizeof ts);
    yyjson_mut_obj_add_strcpy(r->rdoc, o, "time", ts);
    return o;
}

static yyjson_mut_val *orgs_for_user(sqlite3 *db, yyjson_mut_doc *doc,
                                     int64_t user_id, int64_t only_org,
                                     int64_t *first_org)
{
    yyjson_mut_val *arr = yyjson_mut_arr(doc);
    const char *sql = only_org
                          ? "SELECT o.id,o.name,m.role FROM memberships m"
                            " JOIN orgs o ON o.id=m.org_id"
                            " WHERE m.user_id=?1 AND o.id=?2 ORDER BY o.id"
                          : "SELECT o.id,o.name,m.role FROM memberships m"
                            " JOIN orgs o ON o.id=m.org_id"
                            " WHERE m.user_id=?1 ORDER BY o.id";
    sqlite3_stmt *st = NULL;
    if (sqlite3_prepare_v2(db, sql, -1, &st, NULL) != SQLITE_OK)
        return arr;
    sqlite3_bind_int64(st, 1, user_id);
    if (only_org)
        sqlite3_bind_int64(st, 2, only_org);
    while (sqlite3_step(st) == SQLITE_ROW) {
        int64_t id = sqlite3_column_int64(st, 0);
        yyjson_mut_val *o = yyjson_mut_arr_add_obj(doc, arr);
        yyjson_mut_obj_add_int(doc, o, "id", id);
        yyjson_mut_obj_add_strcpy(doc, o, "name", sq(sqlite3_column_text(st, 1)));
        yyjson_mut_obj_add_strcpy(doc, o, "role", sq(sqlite3_column_text(st, 2)));
        if (first_org && *first_org == 0)
            *first_org = id;
    }
    sqlite3_finalize(st);
    return arr;
}

static yyjson_mut_val *session_payload(struct req *r, struct session *s,
                                       int64_t user_id, const char *username,
                                       const char *display,
                                       yyjson_mut_val *orgs,
                                       int64_t active_org)
{
    yyjson_mut_val *o = yyjson_mut_obj(r->rdoc);
    yyjson_mut_obj_add_strcpy(r->rdoc, o, "session", s->id);
    yyjson_mut_val *u = yyjson_mut_obj(r->rdoc);
    yyjson_mut_obj_add_int(r->rdoc, u, "id", user_id);
    yyjson_mut_obj_add_strcpy(r->rdoc, u, "username", username ? username : "");
    yyjson_mut_obj_add_strcpy(r->rdoc, u, "display_name",
                           display ? display : "");
    yyjson_mut_obj_add_bool(r->rdoc, u, "is_admin", s->is_admin != 0);
    yyjson_mut_obj_add_val(r->rdoc, o, "user", u);
    yyjson_mut_obj_add_val(r->rdoc, o, "orgs", orgs);
    if (active_org)
        yyjson_mut_obj_add_int(r->rdoc, o, "active_org", active_org);
    else
        yyjson_mut_obj_add_null(r->rdoc, o, "active_org");
    return o;
}

static yyjson_mut_val *h_session_open(struct req *r)
{
    const char *method = arg_str(r->args, "method");
    if (!method)
        return fail(r, "INVALID_ARGS", "method is required");
    char *reqjson = audit_args_json(r->args);

    if (strcmp(method, "password") == 0) {
        const char *username = arg_str(r->args, "username");
        const char *password = arg_str(r->args, "password");
        if (!username || !password) {
            free(reqjson);
            return fail(r, "INVALID_ARGS", "username and password are required");
        }
        int64_t retry = 0;
        if (rl_blocked("local", &retry)) {
            free(reqjson);
            return failf(r, "RATE_LIMITED",
                         "too many failed logins, retry in %lld seconds",
                         (long long)retry);
        }
        char *pwhash = NULL, *display = NULL;
        int is_admin = 0, disabled = 0;
        int64_t uid = auth_user_lookup(r->db, username, &pwhash, &display,
                                       &is_admin, &disabled);
        int ok = uid > 0 && !disabled && pwhash &&
                 auth_verify_password(pwhash, password) == 0;
        free(pwhash);
        if (!ok) {
            rl_fail("local");
            audit_append(r->db, 0, uid > 0 ? uid : 0, 0, "auth.fail", reqjson,
                         "AUTH_FAILED", NULL);
            free(display);
            free(reqjson);
            return fail(r, "AUTH_FAILED", "invalid credentials");
        }
        rl_ok("local");
        int64_t active = 0;
        yyjson_mut_val *orgs =
            orgs_for_user(r->db, r->rdoc, uid, 0, &active);
        struct session *s =
            sessions_create(uid, is_admin, active, 0, "read,write,admin");
        audit_append(r->db, active, uid, 0, "auth.open", reqjson, "OK", NULL);
        yyjson_mut_val *out = session_payload(r, s, uid, username, display,
                                              orgs, active);
        free(display);
        free(reqjson);
        return out;
    }

    if (strcmp(method, "token") == 0) {
        const char *token = arg_str(r->args, "token");
        if (!token) {
            free(reqjson);
            return fail(r, "INVALID_ARGS", "token is required");
        }
        unsigned char th[32];
        auth_hash_token(token, th);
        sqlite3_stmt *st = NULL;
        int rc = sqlite3_prepare_v2(
            r->db,
            "SELECT t.id,t.user_id,t.org_id,t.scopes,t.expires_at,t.revoked_at,"
            " t.label,u.username,u.display_name,u.is_admin,u.disabled_at"
            " FROM api_tokens t JOIN users u ON u.id=t.user_id"
            " WHERE t.token_hash=?1",
            -1, &st, NULL);
        if (rc != SQLITE_OK) {
            free(reqjson);
            return db_error(r);
        }
        sqlite3_bind_blob(st, 1, th, 32, SQLITE_TRANSIENT);
        int found = sqlite3_step(st) == SQLITE_ROW;
        int64_t token_id = 0, uid = 0, org_id = 0;
        char *scopes = NULL, *expires = NULL, *revoked = NULL;
        char *username = NULL, *display = NULL;
        int is_admin = 0, disabled = 0;
        if (found) {
            token_id = sqlite3_column_int64(st, 0);
            uid = sqlite3_column_int64(st, 1);
            org_id = sqlite3_column_int64(st, 2);
            scopes = xstrdup(sq(sqlite3_column_text(st, 3)));
            if (sqlite3_column_type(st, 4) != SQLITE_NULL)
                expires = xstrdup(sq(sqlite3_column_text(st, 4)));
            if (sqlite3_column_type(st, 5) != SQLITE_NULL)
                revoked = xstrdup(sq(sqlite3_column_text(st, 5)));
            username = xstrdup(sq(sqlite3_column_text(st, 7)));
            display = xstrdup(sq(sqlite3_column_text(st, 8)));
            is_admin = sqlite3_column_int(st, 9);
            disabled = sqlite3_column_type(st, 10) != SQLITE_NULL;
        }
        sqlite3_finalize(st);

        int expired = 0;
        if (expires) {
            char today[16];
            time_t t = (time_t)util_now();
            struct tm tm;
            gmtime_r(&t, &tm);
            strftime(today, sizeof today, "%Y-%m-%d", &tm);
            expired = strcmp(expires, today) < 0;
        }
        if (!found || revoked || disabled || expired) {
            audit_append(r->db, 0, uid, 0, "auth.fail", reqjson, "AUTH_FAILED",
                         NULL);
            free(scopes);
            free(expires);
            free(revoked);
            free(username);
            free(display);
            free(reqjson);
            return fail(r, "AUTH_FAILED", "invalid token");
        }
        char ts[32];
        util_iso8601(util_now(), ts, sizeof ts);
        sqlite3_stmt *up = NULL;
        if (sqlite3_prepare_v2(
                r->db, "UPDATE api_tokens SET last_used_at=?1 WHERE id=?2", -1,
                &up, NULL) == SQLITE_OK) {
            sqlite3_bind_text(up, 1, ts, -1, SQLITE_TRANSIENT);
            sqlite3_bind_int64(up, 2, token_id);
            sqlite3_step(up);
            sqlite3_finalize(up);
        }
        int64_t active = org_id;
        yyjson_mut_val *orgs =
            orgs_for_user(r->db, r->rdoc, uid, org_id, &active);
        struct session *s = sessions_create(uid, is_admin, org_id, org_id,
                                            scopes ? scopes : "read");
        s->active_org = org_id;
        s->token_id = token_id;
        audit_append(r->db, org_id, uid, token_id, "auth.open", reqjson, "OK",
                     NULL);
        yyjson_mut_val *out = session_payload(r, s, uid, username, display,
                                              orgs, org_id);
        free(scopes);
        free(expires);
        free(revoked);
        free(username);
        free(display);
        free(reqjson);
        return out;
    }

    free(reqjson);
    return fail(r, "UNSUPPORTED", "unsupported auth method");
}

static yyjson_mut_val *h_session_close(struct req *r)
{
    char *reqjson = audit_args_json(r->args);
    audit_append(r->db, r->sess->active_org, r->sess->user_id, 0,
                 "session.close", reqjson, "OK", NULL);
    free(reqjson);
    sessions_destroy(r->sess->id);
    return yyjson_mut_obj(r->rdoc);
}

static yyjson_mut_val *h_session_whoami(struct req *r)
{
    char *display = NULL, *username = NULL;
    int is_admin = 0;
    sqlite3_stmt *st = NULL;
    if (sqlite3_prepare_v2(r->db,
                           "SELECT username,display_name,is_admin FROM users"
                           " WHERE id=?1",
                           -1, &st, NULL) != SQLITE_OK)
        return db_error(r);
    sqlite3_bind_int64(st, 1, r->sess->user_id);
    if (sqlite3_step(st) == SQLITE_ROW) {
        username = xstrdup(sq(sqlite3_column_text(st, 0)));
        display = xstrdup(sq(sqlite3_column_text(st, 1)));
        is_admin = sqlite3_column_int(st, 2);
    }
    sqlite3_finalize(st);
    yyjson_mut_val *o = yyjson_mut_obj(r->rdoc);
    yyjson_mut_val *u = yyjson_mut_obj(r->rdoc);
    yyjson_mut_obj_add_int(r->rdoc, u, "id", r->sess->user_id);
    yyjson_mut_obj_add_strcpy(r->rdoc, u, "username", username ? username : "");
    yyjson_mut_obj_add_strcpy(r->rdoc, u, "display_name", display ? display : "");
    yyjson_mut_obj_add_bool(r->rdoc, u, "is_admin", is_admin != 0);
    yyjson_mut_obj_add_val(r->rdoc, o, "user", u);
    yyjson_mut_obj_add_strcpy(r->rdoc, o, "scopes", r->sess->scopes);
    if (r->sess->active_org) {
        yyjson_mut_obj_add_int(r->rdoc, o, "active_org", r->sess->active_org);
        char *role = db_membership_role(r->db, r->sess->active_org,
                                        r->sess->user_id);
        if (role)
            yyjson_mut_obj_add_strcpy(r->rdoc, o, "role", role);
        else
            yyjson_mut_obj_add_null(r->rdoc, o, "role");
        free(role);
    } else {
        yyjson_mut_obj_add_null(r->rdoc, o, "active_org");
        yyjson_mut_obj_add_null(r->rdoc, o, "role");
    }
    free(username);
    free(display);
    return o;
}

static yyjson_mut_val *h_session_list_orgs(struct req *r)
{
    int64_t first = 0;
    yyjson_mut_val *orgs = orgs_for_user(r->db, r->rdoc, r->sess->user_id, 0,
                                         &first);
    yyjson_mut_val *o = yyjson_mut_obj(r->rdoc);
    yyjson_mut_obj_add_val(r->rdoc, o, "items", orgs);
    return o;
}

static yyjson_mut_val *h_session_use_org(struct req *r)
{
    int64_t org = 0;
    if (!arg_int(r->args, "org", &org) || org <= 0)
        return fail(r, "INVALID_ARGS", "org is required");
    if (r->sess->bound_org && org != r->sess->bound_org)
        return fail(r, "ORG_FORBIDDEN", "token is bound to another org");
    char *role = db_membership_role(r->db, org, r->sess->user_id);
    if (!role)
        return fail(r, "ORG_FORBIDDEN", "not a member of this org");
    r->sess->active_org = org;
    yyjson_mut_val *o = yyjson_mut_obj(r->rdoc);
    yyjson_mut_obj_add_int(r->rdoc, o, "active_org", org);
    yyjson_mut_obj_add_strcpy(r->rdoc, o, "role", role);
    free(role);
    return o;
}


static const struct cmd_arg args_session_open[] = {
    { "method", ARG_ENUM, 1, NULL, "password,token",
      "Open with a password or an API token" },
    { "username", ARG_STR, 0, NULL, NULL, "Username for method=password" },
    { "password", ARG_STR, 0, NULL, NULL, "Password for method=password" },
    { "token", ARG_STR, 0, NULL, NULL, "API token for method=token" },
};

static const struct cmd_arg args_session_use_org[] = {
    { "org", ARG_INT, 1, NULL, NULL, "Org id" },
};

const struct command g_cmd_auth[] = {
    { "health", "Liveness probe", PERM_PUBLIC, 0, 0, 0, h_health, NULL, 0 },
    { "meta", "Server metadata and limits", PERM_PUBLIC, 0, 0, 0, h_meta, NULL,
      0 },
    { "session.open", "Open a session (password or token)", PERM_PUBLIC, 0, 0,
      0, h_session_open, CMD_ARGS(args_session_open) },
    { "session.close", "Close the current session", PERM_READ, 0, 0, 0,
      h_session_close, NULL, 0 },
    { "session.whoami", "Current user, org, role and scopes", PERM_READ, 0, 0,
      0, h_session_whoami, NULL, 0 },
    { "session.list_orgs", "Orgs the current user is a member of", PERM_READ,
      0, 0, 0, h_session_list_orgs, NULL, 0 },
    { "session.use_org", "Switch active org", PERM_READ, 0, 0, 0,
      h_session_use_org, CMD_ARGS(args_session_use_org) },
};

const struct cmd_table g_cmd_table_auth = {
    g_cmd_auth, sizeof g_cmd_auth / sizeof g_cmd_auth[0]
};