from . import (
    __NEXT_OPERATION__,
    __FLOW_COMPLETE__,
    authentication_base,
    _auth_api_request,
    FORCE_PASSWORD_PROMPT,
    throw_if_request_message_is_missing_key,
    AuthStorage,
    STORE_PASSWORD_IN_MEMORY,
    CLIENT_GET_REQUEST_RESULT,
)
from .native import _authenticate_native

import getpass
import jsonpatch
import jsonpointer
import logging
import sys

# Constants defining the states and operations for the pam_interactive authentication flow
AUTH_CLIENT_AUTH_REQUEST = "pam_auth_client_request"
AUTH_CLIENT_AUTH_RESPONSE = "pam_auth_response"
PERFORM_RUNNING = "running"
PERFORM_READY = "ready"
PERFORM_NEXT = "next"
PERFORM_RESPONSE = "response"
PERFORM_WAITING = "waiting"
PERFORM_WAITING_PW = "waiting_pw"
PERFORM_ERROR = "error"
PERFORM_TIMEOUT = "timeout"
PERFORM_AUTHENTICATED = "authenticated"
PERFORM_NOT_AUTHENTICATED = "not_authenticated"
PAM_INTERACTIVE_SCHEME = "pam_interactive"
PERFORM_NATIVE_AUTH = "native_auth"

AUTH_AGENT_AUTH_REQUEST = "auth_agent_auth_request"
AUTH_AGENT_AUTH_RESPONSE = "auth_agent_auth_response"

_logger = logging.getLogger(__name__)


def login(conn, **extra_opt):
    """The entry point for the pam_interactive authentication scheme."""

    # The AuthStorage object holds the token generated by the server for the native auth step
    depot = AuthStorage.create_temp_pw_storage(conn)

    auth_client_object = _pam_interactive_ClientAuthState(conn, depot, scheme=PAM_INTERACTIVE_SCHEME)
    auth_client_object.authenticate_client(initial_request=extra_opt)


class _pam_interactive_ClientAuthState(authentication_base):
    def __init__(self, conn, depot, *_, **_kw):
        super().__init__(conn, *_, **_kw)
        self.depot = depot
        self._list_for_request_result_return = None

    def auth_client_start(self, request):
        self._list_for_request_result_return = request.pop(CLIENT_GET_REQUEST_RESULT, None)

        resp = request.copy()

        # "pstate" stores persistent values across multiple authentication steps, allowing the server
        # to recall previous inputs through JSON pointers. "pdirty" flags if pstate has changed and needs syncing.
        # The server side implementation can be found here: https://github.com/irods/irods_auth_plugin_pam_interactive
        # The plugin is built on the authentication framework described here:
        # https://github.com/irods-contrib/irods_working_group_authentication/tree/e83e84df8ea4a732e5de894fb28aae281c3b3d29/development

        resp["pstate"] = resp.get("pstate", {})
        resp["pdirty"] = resp.get("pdirty", False)

        resp['user_name'] = self.conn.account.proxy_user
        resp['zone_name'] = self.conn.account.proxy_zone

        # If not forcing a prompt, check for existing credentials (.irodsA) to attempt native auth directly
        if not resp.get(FORCE_PASSWORD_PROMPT, False):
            if self.conn.account.password and self.conn.account.derived_auth_file:
                resp[__NEXT_OPERATION__] = PERFORM_NATIVE_AUTH
                return resp

        # Otherwise, begin the full interactive flow
        resp[__NEXT_OPERATION__] = AUTH_CLIENT_AUTH_REQUEST
        return resp

    def pam_auth_client_request(self, request):
        server_req = request.copy()
        server_req[__NEXT_OPERATION__] = AUTH_AGENT_AUTH_REQUEST

        resp = _auth_api_request(self.conn, server_req)
        resp[__NEXT_OPERATION__] = AUTH_CLIENT_AUTH_RESPONSE

        return resp

    def pam_auth_response(self, request):
        throw_if_request_message_is_missing_key(request, ["user_name", "zone_name"])

        server_req = request.copy()
        # The "next_operation" key here instructs the server which of its own operations to run
        server_req[__NEXT_OPERATION__] = AUTH_AGENT_AUTH_RESPONSE

        resp = _auth_api_request(self.conn, server_req)

        # The server's response contains a "next_operation" key which dictates the client's next state.
        return resp

    def _get_default_value(self, request):
        """Looks for a default value in pstate based on a path from the server."""

        default_path = request.get("msg", {}).get("default_path", "")
        if default_path:
            jptr = jsonpointer.JsonPointer(default_path)
            return str(jptr.get(request.get("pstate", {}), ""))
        return ""

    def _patch_state(self, req):
        """Applies server patch instructions to the client's pstate."""

        patch_ops = req.get("msg", {}).get("patch")
        if not patch_ops:
            return

        resp = req.get("resp", "")

        # If the patch operation is an add or replace without a value, use the response value (following json patch RFC)
        for op in patch_ops:
            if op.get("op") in ["add", "replace"] and "value" not in op:
                op["value"] = resp

        req["pstate"] = jsonpatch.apply_patch(req.get("pstate", {}), patch_ops)
        req["pdirty"] = True

        del req["msg"]["patch"]

    def _retrieve_entry(self, req):
        """Checks if the server is asking for a value already stored in pstate."""

        if "retrieve" not in req.get("msg", {}):
            return False

        retr_path = req.get("msg", {}).get("retrieve", "")
        if retr_path:
            jptr = jsonpointer.JsonPointer(retr_path)
            req["resp"] = str(jptr.get(req.get("pstate", {}), ""))
            return True

        # If no value found in pstate, set resp to empty string.
        # The server will then decide the next step based on its configuration for the current prompt.
        # It may terminate the flow if the input was required, fall back to a default value, or re-prompt the user.
        req["resp"] = ""
        return True

    def _get_input(self, server_req, is_password=False, prompt_label="Input: "):
        """Handles input from the user, either as a password or regular input."""

        # If the server asks for a value we already have, send it without prompting
        if self._retrieve_entry(server_req):
            self._patch_state(server_req)
            server_req[__NEXT_OPERATION__] = AUTH_AGENT_AUTH_RESPONSE
            return _auth_api_request(self.conn, server_req)

        prompt = server_req.get("msg", {}).get("prompt", prompt_label)
        default_value = self._get_default_value(server_req)

        display_prompt = prompt
        if default_value:
            if is_password:
                display_prompt += " [******] "
            else:
                display_prompt += f" [{default_value}] "

        if is_password:
            user_input = getpass.getpass(display_prompt)
        else:
            sys.stderr.write(display_prompt)
            sys.stderr.flush()
            user_input = sys.stdin.readline().strip()

        server_req["resp"] = user_input or default_value

        self._patch_state(server_req)
        server_req[__NEXT_OPERATION__] = AUTH_AGENT_AUTH_RESPONSE

        return _auth_api_request(self.conn, server_req)

    def waiting(self, request):
        """Handles interactive input requests from the server."""
        server_req = request.copy()
        return self._get_input(server_req, is_password=False)

    def waiting_pw(self, request):
        """Handles the case where a password is specifically requested."""
        server_req = request.copy()
        return self._get_input(server_req, is_password=True, prompt_label="Password: ")

    def authenticated(self, request):
        throw_if_request_message_is_missing_key(request, ["request_result"])
        pw = request["request_result"]  # The password token returned by the server

        if not self.depot:
            raise RuntimeError("auth storage object was either not set, or allowed to expire prematurely.")

        if request.get(STORE_PASSWORD_IN_MEMORY):
            self.depot.use_client_auth_file(None)

        self.depot.store_pw(pw)

        # Return the password token to the caller if requested
        if isinstance(self._list_for_request_result_return, list):
            self._list_for_request_result_return[:] = (pw,)

        resp = request.copy()
        resp[__NEXT_OPERATION__] = PERFORM_NATIVE_AUTH

        return resp

    def native_auth(self, request):
        resp = request.copy()

        # The native auth function will use the depot to retrieve the password token
        _authenticate_native(self.conn, request)

        resp[__NEXT_OPERATION__] = __FLOW_COMPLETE__
        self.loggedIn = 1
        return resp

    # Pass-through states like next, running, ready, and response are used when the server
    # performs internal updates without needing user input. The client applies state changes,
    # optionally shows a prompt, and immediately sends the updated context back.

    def next(self, request):
        prompt = request.get("msg", {}).get("prompt", "")
        if prompt:
            _logger.info("Server prompt: %s", prompt)

        server_req = request.copy()
        self._patch_state(server_req)
        server_req[__NEXT_OPERATION__] = AUTH_AGENT_AUTH_RESPONSE

        resp = _auth_api_request(self.conn, server_req)

        return resp

    def running(self, request):
        return self.next(request)

    def ready(self, request):
        return self.next(request)

    def response(self, request):
        return self.next(request)

    # Failure states like error, timeout, and not_authenticated occur when authentication fails.
    # The client informs the user, marks login as unsuccessful, and ends the flow.

    def _auth_failure(self, request, message):
        _logger.error(message)
        resp = request.copy()
        resp[__NEXT_OPERATION__] = __FLOW_COMPLETE__
        self.loggedIn = 0
        return resp

    def error(self, request):
        return self._auth_failure(request, "Authentication error.")

    def timeout(self, request):
        return self._auth_failure(request, "Authentication timed out.")

    def not_authenticated(self, request):
        return self._auth_failure(request, "Authentication failed possibly due to incorrect credentials.")
