#include "auth_plugin.h"
#include <security/pam_appl.h>
#include <pwd.h>
#include <grp.h>
#include <string>
#include <cstring>
#include <memory>

class PAMAuthPlugin : public Auth {
private:
	std::string service_name;
	uid_t user_uid;
	gid_t user_gid;

	// PAM Conversation function
	static int pam_conv_func(int num_msg, const struct pam_message **msg,
							struct pam_response **resp, void *appdata_ptr) {
		const char* password = static_cast<const char*>(appdata_ptr);
		if (!password) return PAM_CONV_ERR;

		// Allocate memory for responses
		*resp = static_cast<pam_response*>(calloc(num_msg, sizeof(struct pam_response)));
		if (*resp == nullptr) return PAM_CONV_ERR;

		// Handle messages
		for (int i = 0; i < num_msg; i++) {
			if (msg[i]->msg_style == PAM_PROMPT_ECHO_OFF) {
				(*resp)[i].resp = strdup(password);
				if ((*resp)[i].resp == nullptr) {
					for (int j = 0; j < i; j++) {
						free((*resp)[j].resp);
					}
					free(*resp);
					*resp = nullptr;
					return PAM_CONV_ERR;
				}
				(*resp)[i].resp_retcode = 0;
			} else {
				(*resp)[i].resp = nullptr;
				(*resp)[i].resp_retcode = 0;
			}
		}

		return PAM_SUCCESS;
	}

public:
	PAMAuthPlugin() : service_name("ftpd") {}

	virtual bool initialize(const std::map<std::string, std::string>& config) override {
		// Get PAM service name with default
		auto service_it = config.find("pam_service");
		service_name = (service_it != config.end()) ? service_it->second : "ftpd";
		
		auto chroot_it = config.find("chroot");
		this->setChroot((chroot_it != config.end()) ? 
					 (chroot_it->second == "on" || chroot_it->second == "true" || chroot_it->second == "yes") : 
					 true);
		
		return true;
	}

	virtual bool authenticate(ClientAuthDetails* auth_data) override {
		if (!auth_data || !auth_data->username[0] || !auth_data->password[0]) {
			logger->print(LOGLEVEL_ERROR, "auth_pam: Cannot use empty auth data");
			return false;
		}

		struct pam_conv conv;
		conv.conv = pam_conv_func;
		conv.appdata_ptr = static_cast<void*>(auth_data->password);

		pam_handle_t* pamh = nullptr;

		// Start PAM session
		int retval = pam_start(service_name.c_str(), auth_data->username, &conv, &pamh);
		if (retval != PAM_SUCCESS) {
			logger->print(LOGLEVEL_ERROR, "auth_pam: Failed to start PAM: %s", 
						 pamh ? pam_strerror(pamh, retval) : "Unknown error");
			if (pamh) pam_end(pamh, retval);
			return false;
		}

		// Authenticate user
		retval = pam_authenticate(pamh, 0);
		if (retval != PAM_SUCCESS) {
			logger->print(LOGLEVEL_ERROR, "auth_pam: Authentication failed: %s", 
						 pam_strerror(pamh, retval));
			pam_end(pamh, retval);
			return false;
		}

		// Check account validity
		retval = pam_acct_mgmt(pamh, 0);
		if (retval != PAM_SUCCESS) {
			logger->print(LOGLEVEL_ERROR, "auth_pam: Account validation failed: %s", 
						 pam_strerror(pamh, retval));
			pam_end(pamh, retval);
			return false;
		}

		// End PAM session
		pam_end(pamh, PAM_SUCCESS);

		// Get user info
		struct passwd* pw = getpwnam(auth_data->username);
		if (!pw) {
			logger->print(LOGLEVEL_ERROR, "auth_pam: Failed to get user info for %s", 
						 auth_data->username);
			return false;
		}

		// Store user info
		user_uid = pw->pw_uid;
		user_gid = pw->pw_gid;

		// Set home directory
		if (pw->pw_dir) {
			strncpy(auth_data->home_dir, pw->pw_dir, sizeof(auth_data->home_dir) - 1);
			auth_data->home_dir[sizeof(auth_data->home_dir) - 1] = '\0';
		}

		return true;
	}

	virtual bool isPasswordRequired() override {
		return true;
	}
};

IMPLEMENT_AUTH_PLUGIN(PAMAuthPlugin, "pam", "PAM-based local authentication", "1.0.0")