# Copyright 2026 gRPC authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from cpython.bytes cimport PyBytes_FromStringAndSize
from libcpp.utility cimport move
from cpython cimport PyObject, PyBytes_AsString, PyBytes_GET_SIZE

import threading


cdef StatusOr[string] MakeInternalError(string message):
    return StatusOr[string](Status(AbslStatusCode.kUnknown, message))

cdef StatusOr[string] MakeStringResult(string result):
  return StatusOr[string](result)

cdef class OnCompleteWrapper:
  cdef weak_ptr[PrivateKeySignerPyWrapper.CompletionContext] completion_context

  # Makes this class callable
  def __call__(self, result):
    cdef StatusOr[string] cpp_result
    cdef string cpp_string
    cdef shared_ptr[PrivateKeySignerPyWrapper.CompletionContext] locked_completion_context 
    cdef PrivateKeySignerPyWrapper.CompletionContext* local_completion_context
    if isinstance(result, bytes):
      # We got a signature
      cpp_string = MakeStringForCython(PyBytes_AsString(result), PyBytes_GET_SIZE(result))
      cpp_result = MakeStringResult(cpp_string)
    elif isinstance(result, Exception):
      # If python returns an exception, convert to absl::Status
      cpp_string = MakeStringForCython(PyBytes_AsString(str(result).encode('utf-8')))
      cpp_result = MakeInternalError(cpp_string)
    else:
      # Any other return type is not valid
      cpp_string = MakeStringForCython(PyBytes_AsString(f"Invalid result type: {type(result)}".encode('utf-8')))
      cpp_result = MakeInternalError(cpp_string)
    locked_completion_context = self.completion_context.lock()
    local_completion_context = locked_completion_context.get()
    with nogil:
      if local_completion_context != NULL:
        local_completion_context.OnComplete(cpp_result)

cdef PrivateKeySignerPyWrapper.PrivateKeySignerPyWrapperResult async_sign_wrapper(string_view inp, CSignatureAlgorithm algorithm, PyObject* py_user_sign_fn, weak_ptr[PrivateKeySignerPyWrapper.CompletionContext] completion_context) noexcept nogil:
  cdef string cpp_string
  cdef const char* data
  cdef size_t size
  cdef PrivateKeySignerPyWrapper.PrivateKeySignerPyWrapperResult cpp_result
  with gil:
    # Cast the PyObject* pointer holding the user's python sign impl
    py_user_func = <object>py_user_sign_fn

    py_on_complete_wrapper = OnCompleteWrapper()
    py_on_complete_wrapper.completion_context = move(completion_context)

    # Call the user's Python function and handle results
    py_result = None
    # Mark sync by default, change if the result is async
    cpp_result.is_sync = True
    try:
      data = inp.data()
      size = inp.length()
      py_bytes = PyBytes_FromStringAndSize(data, size)
      py_result = py_user_func(py_bytes, algorithm, py_on_complete_wrapper)
      if isinstance(py_result, bytes):
        # We got a signature
        cpp_string = MakeStringForCython(PyBytes_AsString(py_result), PyBytes_GET_SIZE(py_result))
        cpp_result.sync_result = MakeStringResult(cpp_string)
      elif callable(py_result):
        # Cancellation func
        cpp_result.is_sync = False
        Py_INCREF(py_result)
        cpp_result.async_result.py_user_cancel_fn = <PyObject*> py_result
        cpp_result.async_result.cancel_wrapper = cancel_wrapper
      else:
        # Any other return type is not valid
        cpp_string = MakeStringForCython(PyBytes_AsString(f"Invalid result type: {type(py_result)}".encode('utf-8')))
        cpp_result.sync_result = MakeInternalError(cpp_string)
      return cpp_result
    except Exception as e:
      # If Python raises an exception, make it an error status
      cpp_result.sync_result = MakeInternalError(MakeStringForCython(PyBytes_AsString(f"Exception in user function: {e}".encode('utf-8'))))
      return cpp_result

cdef void cancel_wrapper(PyObject* py_cancel_user_fn) noexcept nogil:
  with gil:
    try:
      py_cancel_func = <object>py_cancel_user_fn
      py_cancel_func()
    except Exception as e:
      # Exceptions in cancellation
      _LOGGER.warning('Exception in user-provided cancellation function: %s', e)


# To be called from the python layer when the user provides a signer function.
cdef shared_ptr[PrivateKeySigner] build_private_key_signer(py_user_func):
  # Run a non-daemon thread that waits for this event to be set. This ensures
  # that the Python interpreter will stay alive long enough for these objects
  # passed into C++ to be properly dereferenced and cleaned up.
  destroy_event = threading.Event()
  destroy_lambda = lambda event=destroy_event: event.wait()
  threading.Thread(target=destroy_lambda, daemon=False).start()
  py_private_key_signer = PrivateKeySignerPyWrapper.Create(async_sign_wrapper, <PyObject*>py_user_func, <PyObject*>destroy_event)
  return py_private_key_signer
