auth.cpp

#include "auth.h"
#include "config.h"
#include "azure_msi.h"
#include "nlohmann/json.hpp"

#include <iostream>
#include <fstream>
#include <filesystem>
#include <cstdlib>
#include <chrono>
#include <iomanip>
#include <sstream>
#include <sys/stat.h>

using json = nlohmann::json;

// =============================================================================
// KeyStore factory
// =============================================================================

std::unique_ptr<KeyStore> KeyStore::create(const std::string& uri) {
    if (uri.empty()) {
        return std::make_unique<NoneKeyStore>();
    }

    // Parse URI scheme
    auto scheme_end = uri.find("://");
    if (scheme_end == std::string::npos) {
        throw std::runtime_error("Invalid apikey_store URI: " + uri + " (expected scheme://...)");
    }

    std::string scheme = uri.substr(0, scheme_end);
    std::string path = uri.substr(scheme_end + 3);  // Everything after ://

    if (scheme == "file") {
        return std::make_unique<JsonKeyStore>(path);
    }
    if (scheme == "postgresql" || scheme == "postgres") {
#ifdef ENABLE_POSTGRESQL
        // PGKeyStore created with full URI (needs connection string)
        return create_pg_keystore(uri);
#else
        throw std::runtime_error("PostgreSQL support not compiled. Rebuild with -DENABLE_POSTGRESQL=ON");
#endif
    }
    if (scheme == "msi") {
        return std::make_unique<MsiKeyStore>();
    }

    throw std::runtime_error("Unknown apikey_store scheme: " + scheme + " (valid: file, postgresql, msi)");
}

// =============================================================================
// JsonKeyStore implementation
// =============================================================================

JsonKeyStore::JsonKeyStore(const std::string& path) {
    if (!path.empty()) {
        custom_path = path;
    }
    keys = load_keys();
}

std::string JsonKeyStore::get_keys_file_path() {
    if (!custom_path.empty()) {
        return custom_path;
    }
    // Use XDG base directory (same pattern as config.cpp)
    std::string config_home;
    const char* xdg_config = getenv("XDG_CONFIG_HOME");
    if (xdg_config && xdg_config[0] != '\0') {
        config_home = xdg_config;
    } else {
        config_home = Config::get_home_directory() + "/.config";
    }
    return config_home + "/shepherd/api_keys.json";
}

std::map<std::string, ApiKeyEntry> JsonKeyStore::load_keys() {
    std::map<std::string, ApiKeyEntry> result;
    std::string path = get_keys_file_path();

    if (!std::filesystem::exists(path)) {
        return result;  // No keys file = empty map
    }

    std::ifstream file(path);
    if (!file.is_open()) {
        return result;
    }

    try {
        json j;
        file >> j;

        for (auto& [key, value] : j.items()) {
            ApiKeyEntry entry;
            entry.name = value.value("name", "");
            entry.notes = value.value("notes", "");
            entry.created = value.value("created", "");
            entry.permissions = value.value("permissions", json::object());
            result[key] = entry;
        }
    } catch (const json::exception&) {
        throw std::runtime_error("Failed to parse api_keys.json");
    }

    return result;
}

void JsonKeyStore::save_keys(const std::map<std::string, ApiKeyEntry>& keys) {
    std::string path = get_keys_file_path();

    // Ensure directory exists
    std::filesystem::path dir = std::filesystem::path(path).parent_path();
    if (!std::filesystem::exists(dir)) {
        std::filesystem::create_directories(dir);
    }

    // Build JSON
    json j;
    for (const auto& [key, entry] : keys) {
        j[key] = {
            {"name", entry.name},
            {"notes", entry.notes},
            {"created", entry.created},
            {"permissions", entry.permissions}
        };
    }

    // Write to file
    std::ofstream file(path);
    if (!file.is_open()) {
        throw std::runtime_error("Failed to write api_keys.json");
    }
    file << j.dump(4) << std::endl;
    file.close();

    // Set file permissions to 0600 (owner read/write only)
    chmod(path.c_str(), S_IRUSR | S_IWUSR);
}

std::string JsonKeyStore::generate_key() {
    static const char charset[] =
        "abcdefghijklmnopqrstuvwxyz"
        "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        "0123456789";

    std::string key = "sk-";

    // Use /dev/urandom for cryptographic randomness
    std::ifstream urandom("/dev/urandom", std::ios::binary);
    if (!urandom) {
        throw std::runtime_error("Failed to open /dev/urandom");
    }

    for (int i = 0; i < 32; i++) {
        unsigned char byte;
        urandom.read(reinterpret_cast<char*>(&byte), 1);
        key += charset[byte % (sizeof(charset) - 1)];
    }

    return key;
}

bool JsonKeyStore::validate_key(const std::string& key) {
    if (key.empty() || keys.empty()) {
        return false;
    }

    // Constant-time comparison to prevent timing attacks
    // Check against all keys, accumulating result
    bool found = false;
    for (const auto& [stored_key, entry] : keys) {
        if (stored_key.length() != key.length()) {
            continue;
        }

        // Constant-time comparison
        int result = 0;
        for (size_t i = 0; i < stored_key.length(); i++) {
            result |= stored_key[i] ^ key[i];
        }

        if (result == 0) {
            found = true;
            // Don't break early - continue checking to maintain constant time
        }
    }

    return found;
}

bool JsonKeyStore::is_enabled() {
    return !keys.empty();
}

const ApiKeyEntry* JsonKeyStore::get_entry(const std::string& key) const {
    auto it = keys.find(key);
    return (it != keys.end()) ? &it->second : nullptr;
}

std::map<std::string, ApiKeyEntry> JsonKeyStore::list_keys() {
    return load_keys();
}

bool JsonKeyStore::add_key(const std::string& key, const ApiKeyEntry& entry) {
    auto all = load_keys();
    all[key] = entry;
    save_keys(all);
    keys = all;
    return true;
}

bool JsonKeyStore::remove_key(const std::string& key) {
    auto all = load_keys();
    if (all.erase(key) == 0) return false;
    save_keys(all);
    keys = all;
    return true;
}

bool JsonKeyStore::update_key(const std::string& key, const ApiKeyEntry& entry) {
    auto all = load_keys();
    auto it = all.find(key);
    if (it == all.end()) return false;
    it->second = entry;
    save_keys(all);
    keys = all;
    return true;
}

std::string JsonKeyStore::store_location() {
    return get_keys_file_path();
}

// =============================================================================
// MsiKeyStore implementation
// =============================================================================

extern std::unique_ptr<Config> config;  // Global config

MsiKeyStore::MsiKeyStore() {
    // Keys loaded lazily on first use
}

void MsiKeyStore::ensure_loaded() {
    if (loaded) return;
    loaded = true;

    // Get vault name from config
    if (!config || config->keyvault_name.empty()) {
        error_message = "MsiKeyStore requires --config msi --kv <vault>";
        std::cerr << "Auth error: " << error_message << std::endl;
        return;
    }

    // Fetch keys from Key Vault
    auto secret = azure::get_keyvault_secret(config->keyvault_name, "shepherd-keys");
    if (!secret) {
        error_message = "Failed to fetch 'shepherd-keys' from Key Vault";
        std::cerr << "Auth error: " << error_message << std::endl;
        return;
    }

    // Parse JSON (same format as local api_keys.json)
    try {
        json j = json::parse(*secret);

        for (auto& [key, value] : j.items()) {
            ApiKeyEntry entry;
            entry.name = value.value("name", "");
            entry.notes = value.value("notes", "");
            entry.created = value.value("created", "");
            entry.permissions = value.value("permissions", json::object());
            keys[key] = entry;
        }

        std::cout << "Loaded " << keys.size() << " API keys from Key Vault" << std::endl;
    } catch (const json::exception& e) {
        error_message = "Failed to parse shepherd-keys JSON: " + std::string(e.what());
        std::cerr << "Auth error: " << error_message << std::endl;
    }
}

bool MsiKeyStore::validate_key(const std::string& key) {
    ensure_loaded();

    if (key.empty() || keys.empty()) {
        return false;
    }

    // Constant-time comparison to prevent timing attacks
    bool found = false;
    for (const auto& [stored_key, entry] : keys) {
        if (stored_key.length() != key.length()) {
            continue;
        }

        int result = 0;
        for (size_t i = 0; i < stored_key.length(); i++) {
            result |= stored_key[i] ^ key[i];
        }

        if (result == 0) {
            found = true;
        }
    }

    return found;
}

bool MsiKeyStore::is_enabled() {
    ensure_loaded();
    return !keys.empty();
}

const ApiKeyEntry* MsiKeyStore::get_entry(const std::string& key) const {
    auto it = keys.find(key);
    return (it != keys.end()) ? &it->second : nullptr;
}

// =============================================================================
// CLI handler: shepherd apikey
// =============================================================================

static std::string get_iso_timestamp() {
    auto now = std::chrono::system_clock::now();
    auto time = std::chrono::system_clock::to_time_t(now);
    std::tm tm = *std::gmtime(&time);
    std::ostringstream oss;
    oss << std::put_time(&tm, "%Y-%m-%dT%H:%M:%SZ");
    return oss.str();
}

static std::string mask_key(const std::string& key) {
    // Show first 7 chars (sk-XXXX) and last 4 chars
    if (key.length() < 12) return key;
    return key.substr(0, 7) + "..." + key.substr(key.length() - 4);
}

int handle_apikey_args(const std::vector<std::string>& args,
                       std::function<void(const std::string&)> callback,
                       const std::string& apikey_store_uri) {
    // Create the appropriate key store from configured URI
    // Empty URI or no URI defaults to file:// (local JSON)
    std::string uri = apikey_store_uri.empty() ? "file://" : apikey_store_uri;
    std::unique_ptr<KeyStore> store;
    try {
        store = KeyStore::create(uri);
    } catch (const std::exception& e) {
        callback("Error creating key store: " + std::string(e.what()) + "\n");
        return 1;
    }

    // Helper: list keys in table format
    auto print_key_list = [&](const std::map<std::string, ApiKeyEntry>& keys) {
        callback("NAME           KEY               CREATED      NOTES\n");
        for (const auto& [key, entry] : keys) {
            std::string line = entry.name;
            while (line.length() < 14) line += " ";
            line += " " + mask_key(key);
            while (line.length() < 32) line += " ";
            std::string date = entry.created.length() >= 10 ? entry.created.substr(0, 10) : entry.created;
            line += " " + date;
            while (line.length() < 45) line += " ";
            line += " " + entry.notes + "\n";
            callback(line);
        }
    };

    // No args shows list
    if (args.empty()) {
        auto keys = store->list_keys();
        if (keys.empty()) {
            callback("No API keys configured.\n");
            callback("Create one with: shepherd apikey create <name>\n");
            return 0;
        }
        print_key_list(keys);
        return 0;
    }

    // Determine if first arg is an action (action-first) or a name (name-first)
    std::string first_arg = args[0];
    bool action_first = (first_arg == "list" || first_arg == "create" || first_arg == "add" ||
                         first_arg == "help" || first_arg == "--help" || first_arg == "-h");

    std::string name;
    std::string subcmd;

    if (action_first) {
        subcmd = first_arg;
        name = (args.size() >= 2) ? args[1] : "";
    } else {
        // Name-first: apikey NAME [action]
        name = first_arg;
        subcmd = (args.size() >= 2) ? args[1] : "help";  // No action = show help
    }

    if (subcmd == "help" || subcmd == "--help" || subcmd == "-h") {
        if (!name.empty()) {
            callback("Usage: shepherd apikey " + name + " <action>\n");
            callback("\nActions:\n");
            callback("  show     - Show full details of this key\n");
            callback("  set      - Update notes/permissions\n");
            callback("  remove   - Remove this key\n");
        } else {
            callback("Usage: shepherd apikey <name> <action>\n");
            callback("\nActions (after name):\n");
            callback("  show     - Show full details\n");
            callback("  set [--notes <notes>] [--perms <json>] - Update\n");
            callback("  remove   - Remove key\n");
            callback("\nOther commands:\n");
            callback("  list     - List all keys\n");
            callback("  create <name> [--notes <notes>] [--perms <json>]\n");
        }
        return 0;
    }

    // List keys
    if (subcmd == "list") {
        auto keys = store->list_keys();
        if (keys.empty()) {
            callback("No API keys configured.\n");
            callback("Create one with: shepherd apikey create --name <name>\n");
            return 0;
        }
        print_key_list(keys);
        return 0;
    }

    // Show key details
    if (subcmd == "show") {
        if (name.empty()) {
            callback("Usage: shepherd apikey <name> show\n");
            return 1;
        }

        auto keys = store->list_keys();

        // Find key by name
        for (const auto& [key, entry] : keys) {
            if (entry.name == name) {
                callback("Name:        " + entry.name + "\n");
                callback("Key:         " + mask_key(key) + "\n");
                callback("Created:     " + entry.created + "\n");
                callback("Notes:       " + (entry.notes.empty() ? "(none)" : entry.notes) + "\n");
                callback("Permissions: " + entry.permissions.dump() + "\n");
                return 0;
            }
        }

        callback("Error: Key '" + name + "' not found\n");
        return 1;
    }

    // Set key (update notes/perms)
    if (subcmd == "set") {
        if (name.empty()) {
            callback("Usage: shepherd apikey <name> set [--notes <notes>] [--perms <json>]\n");
            return 1;
        }

        std::string notes;
        json permissions;
        bool has_notes = false;
        bool has_perms = false;

        // Parse optional arguments (start at args[2] in name-first mode)
        for (size_t i = 2; i < args.size(); i++) {
            if (args[i] == "--notes" && i + 1 < args.size()) {
                notes = args[++i];
                has_notes = true;
            } else if (args[i] == "--perms" && i + 1 < args.size()) {
                try {
                    permissions = json::parse(args[++i]);
                    has_perms = true;
                } catch (const json::exception& e) {
                    callback("Error: Invalid JSON for --perms: " + std::string(e.what()) + "\n");
                    return 1;
                }
            }
        }

        if (!has_notes && !has_perms) {
            callback("Error: Specify --notes or --perms to update\n");
            return 1;
        }

        auto keys = store->list_keys();

        // Find and update key
        for (auto& [key, entry] : keys) {
            if (entry.name == name) {
                if (has_notes) entry.notes = notes;
                if (has_perms) entry.permissions = permissions;
                store->update_key(key, entry);
                callback("Key '" + name + "' updated.\n");
                return 0;
            }
        }

        callback("Error: Key '" + name + "' not found\n");
        return 1;
    }

    // Remove key
    if (subcmd == "remove") {
        if (name.empty()) {
            callback("Usage: shepherd apikey <name> remove\n");
            return 1;
        }

        auto keys = store->list_keys();

        // Find key by name
        std::string key_to_remove;
        for (const auto& [key, entry] : keys) {
            if (entry.name == name) {
                key_to_remove = key;
                break;
            }
        }

        if (key_to_remove.empty()) {
            callback("Error: Key '" + name + "' not found\n");
            return 1;
        }

        store->remove_key(key_to_remove);
        callback("Key '" + name + "' removed.\n");
        return 0;
    }

    // Create new key
    if (subcmd == "create" || subcmd == "add") {
        if (args.size() < 2) {
            callback("Error: Missing key name\n");
            callback("Usage: shepherd apikey create <name> [--notes <notes>] [--perms <json>]\n");
            return 1;
        }

        std::string name = args[1];
        std::string notes;
        json permissions = json::object();

        // Parse optional arguments
        for (size_t i = 2; i < args.size(); i++) {
            if (args[i] == "--notes" && i + 1 < args.size()) {
                notes = args[++i];
            } else if (args[i] == "--perms" && i + 1 < args.size()) {
                try {
                    permissions = json::parse(args[++i]);
                } catch (const json::exception& e) {
                    callback("Error: Invalid JSON for --perms: " + std::string(e.what()) + "\n");
                    return 1;
                }
            }
        }

        // Check for duplicate name
        auto keys = store->list_keys();
        for (const auto& [key, entry] : keys) {
            if (entry.name == name) {
                callback("Error: Key with name '" + name + "' already exists\n");
                return 1;
            }
        }

        // Generate key
        std::string new_key = JsonKeyStore::generate_key();

        // Create entry
        ApiKeyEntry entry;
        entry.name = name;
        entry.notes = notes;
        entry.created = get_iso_timestamp();
        entry.permissions = permissions;

        // Save
        store->add_key(new_key, entry);

        callback("API key created successfully.\n\n");
        callback("Key: " + new_key + "\n\n");
        callback("Keys are stored in: " + store->store_location() + "\n");
        return 0;
    }

    callback("Error: Unknown command '" + subcmd + "'\n");
    callback("Use 'shepherd apikey help' for usage.\n");
    return 1;
}