# Copyright (c) Facebook, Inc. and its affiliates.

include(FBCMakeParseArgs)
include(FBPythonBinary)

# ---------------------------------------------------------------------------
# thrift-python runtime discovery
# ---------------------------------------------------------------------------
# thrift_types.py (generated by mstch_python) needs two C-extension packages
# at runtime:
#   1. thrift.python  – from the pip-installed fbthrift-python wheel
#   2. folly          – from folly-python site-packages
#
# Because these are C extensions (.so), they CANNOT be loaded from inside a
# Python zip archive.  Any executable that transitively depends on a
# thrift-python library MUST use TYPE dir (not the default zipapp).
#
# Call fb_find_thrift_python_runtime() once at configure time to locate the
# site-packages directories.  Then use add_fb_thrift_python_executable()
# instead of add_fb_python_executable() for targets that import thrift_types.
# ---------------------------------------------------------------------------

macro(fb_find_thrift_python_runtime)
  if(NOT DEFINED _FB_THRIFT_PYTHON_RUNTIME_FOUND)
    # Ensure Python3 is found so we can use ${Python3_EXECUTABLE}
    if(NOT Python3_EXECUTABLE)
      find_package(Python3 COMPONENTS Interpreter REQUIRED)
    endif()

    # Derive the getdeps install root (where all dependencies are installed).
    # getdeps sets GETDEPS_INSTALL_DIR; fall back to CMAKE_INSTALL_PREFIX/..
    if(DEFINED ENV{GETDEPS_INSTALL_DIR})
      set(_getdeps_root "$ENV{GETDEPS_INSTALL_DIR}")
    else()
      get_filename_component(_getdeps_root "${CMAKE_INSTALL_PREFIX}/.." ABSOLUTE)
    endif()
    message(STATUS "getdeps install root: ${_getdeps_root}")

    # --- thrift.python (from pip-installed fbthrift-python wheel) ---
    # Detect via import thrift.python.types — this module only exists in
    # fbthrift-python, not in Apache Thrift, so it's an unambiguous check.
    # If not available, auto-install the wheel from getdeps build output.
    execute_process(
      COMMAND "${Python3_EXECUTABLE}" -c
        "import thrift.python.types; import thrift, os; print(os.path.dirname(os.path.dirname(thrift.__file__)))"
      OUTPUT_VARIABLE THRIFT_PYTHON_SITE_PACKAGES
      OUTPUT_STRIP_TRAILING_WHITESPACE
      ERROR_QUIET
      RESULT_VARIABLE _thrift_result
    )
    if(NOT _thrift_result EQUAL 0 OR NOT THRIFT_PYTHON_SITE_PACKAGES)
      # Not pip-installed yet — auto-install from getdeps build output
      file(GLOB _fbthrift_wheels
        "${_getdeps_root}/fbthrift-python/share/thrift/wheels/*.whl"
      )
      if(_fbthrift_wheels)
        list(GET _fbthrift_wheels 0 _fbthrift_whl)
        message(STATUS "Auto-installing fbthrift-python wheel: ${_fbthrift_whl}")
        execute_process(
          COMMAND "${Python3_EXECUTABLE}" -m pip install --no-deps "${_fbthrift_whl}"
          RESULT_VARIABLE _pip_result
        )
        if(_pip_result EQUAL 0)
          execute_process(
            COMMAND "${Python3_EXECUTABLE}" -c
              "import thrift.python.types; import thrift, os; print(os.path.dirname(os.path.dirname(thrift.__file__)))"
            OUTPUT_VARIABLE THRIFT_PYTHON_SITE_PACKAGES
            OUTPUT_STRIP_TRAILING_WHITESPACE
            ERROR_QUIET
            RESULT_VARIABLE _thrift_result
          )
        endif()
      endif()
    endif()

    if(THRIFT_PYTHON_SITE_PACKAGES)
      message(STATUS "Found thrift-python runtime: ${THRIFT_PYTHON_SITE_PACKAGES}")
    endif()
    if(NOT THRIFT_PYTHON_SITE_PACKAGES)
      message(WARNING
        "thrift-python runtime not found.  "
        "Searched: ${_getdeps_root}/fbthrift-python/share/thrift/wheels/*.whl  "
        "pip install the fbthrift-python wheel before configuring the project.")
    endif()

    # --- folly (from folly-python site-packages) ---
    # folly-python doesn't produce a wheel; it installs directly to
    # a site-packages directory under its getdeps prefix.
    execute_process(
      COMMAND "${Python3_EXECUTABLE}" -c
        "import folly, os; print(os.path.dirname(os.path.dirname(folly.__file__)))"
      OUTPUT_VARIABLE FOLLY_PYTHON_SITE_PACKAGES
      OUTPUT_STRIP_TRAILING_WHITESPACE
      ERROR_QUIET
      RESULT_VARIABLE _folly_result
    )
    if(NOT _folly_result EQUAL 0)
      # Search the getdeps install prefix for folly-python site-packages
      file(GLOB _folly_sp_dirs
        "${_getdeps_root}/folly-python/lib64/python*/site-packages"
        "${_getdeps_root}/folly-python/lib/python*/site-packages"
      )
      foreach(_sp IN LISTS _folly_sp_dirs)
        if(EXISTS "${_sp}/folly/__init__.py")
          set(FOLLY_PYTHON_SITE_PACKAGES "${_sp}")
          break()
        endif()
      endforeach()
      if(FOLLY_PYTHON_SITE_PACKAGES)
        message(STATUS "Found folly-python runtime (via prefix): ${FOLLY_PYTHON_SITE_PACKAGES}")
      else()
        message(WARNING
          "folly-python runtime not found.  "
          "Searched: ${_getdeps_root}/folly-python/lib*/python*/site-packages  "
          "Ensure folly-python is built.")
      endif()
    else()
      message(STATUS "Found folly-python runtime: ${FOLLY_PYTHON_SITE_PACKAGES}")
    endif()

    set(_FB_THRIFT_PYTHON_RUNTIME_FOUND TRUE CACHE INTERNAL "thrift-python runtime discovery completed")
    set(THRIFT_PYTHON_SITE_PACKAGES "${THRIFT_PYTHON_SITE_PACKAGES}" CACHE INTERNAL "thrift-python site-packages path")
    set(FOLLY_PYTHON_SITE_PACKAGES "${FOLLY_PYTHON_SITE_PACKAGES}" CACHE INTERNAL "folly-python site-packages path")
  endif()
endmacro()

# ---------------------------------------------------------------------------
# add_fbthrift_python_library()
# ---------------------------------------------------------------------------
# Generate a thrift-python library from a thrift file.
# Analogous to add_fbthrift_py_library() but uses the mstch_python generator
# (--gen mstch_python) which produces thrift_types.py instead of ttypes.py.
#
# Unlike add_fbthrift_py_library(), this does NOT depend on FBThrift::thrift_py
# (the legacy pure-Python runtime).  The thrift-python runtime (thrift.python
# and folly) are C extensions resolved at runtime via symlinks — see
# add_fb_thrift_python_executable().
# ---------------------------------------------------------------------------
function(add_fbthrift_python_library LIB_NAME THRIFT_FILE)
  # Parse the arguments
  set(one_value_args NAMESPACE THRIFT_INCLUDE_DIR)
  set(multi_value_args SERVICES DEPENDS OPTIONS)
  fb_cmake_parse_args(
    ARG "" "${one_value_args}" "${multi_value_args}" "${ARGN}"
  )

  if(NOT DEFINED ARG_THRIFT_INCLUDE_DIR)
    set(ARG_THRIFT_INCLUDE_DIR "include/thrift-files")
  endif()

  get_filename_component(base ${THRIFT_FILE} NAME_WE)
  set(output_dir "${CMAKE_CURRENT_BINARY_DIR}/${THRIFT_FILE}-python")

  # Parse the namespace value
  if (NOT DEFINED ARG_NAMESPACE)
    set(ARG_NAMESPACE "")
  endif()

  # mstch_python outputs to gen-python/<namespace_py3>/<basename>/
  # NAMESPACE should match the "namespace py3" value in the thrift file
  # (consistent with BUCK's py3_namespace).  The basename is appended
  # automatically by this function, just as the mstch_python generator does.
  string(REPLACE "." "/" namespace_dir "${ARG_NAMESPACE}")
  if(namespace_dir STREQUAL "")
    set(py_output_dir "${output_dir}/gen-python/${base}")
  else()
    set(py_output_dir "${output_dir}/gen-python/${namespace_dir}/${base}")
  endif()
  list(APPEND generated_sources
    "${py_output_dir}/thrift_types.py"
    "${py_output_dir}/thrift_enums.py"
    "${py_output_dir}/thrift_metadata.py"
    "${py_output_dir}/thrift_abstract_types.py"
    "${py_output_dir}/thrift_mutable_types.py"
  )
  foreach(service IN LISTS ARG_SERVICES)
    list(APPEND generated_sources
      "${py_output_dir}/thrift_clients.py"
      "${py_output_dir}/thrift_mutable_clients.py"
      "${py_output_dir}/thrift_services.py"
      "${py_output_dir}/thrift_mutable_services.py"
    )
    break()  # Service files are per-module, not per-service
  endforeach()

  # Define a dummy interface library to help propagate the thrift include
  # directories between dependencies.
  add_library("${LIB_NAME}.thrift_includes" INTERFACE)
  target_include_directories(
    "${LIB_NAME}.thrift_includes"
    INTERFACE
      "$<BUILD_INTERFACE:${CMAKE_SOURCE_DIR}>"
      "$<INSTALL_INTERFACE:${ARG_THRIFT_INCLUDE_DIR}>"
  )
  foreach(dep IN LISTS ARG_DEPENDS)
    target_link_libraries(
      "${LIB_NAME}.thrift_includes"
      INTERFACE "${dep}.thrift_includes"
    )
  endforeach()

  # This generator expression gets the list of include directories required
  # for all of our dependencies.
  if (NOT POLICY CMP0067)
    message(FATAL_ERROR "add_fbthrift_python_library() requires CMake 3.8+")
  endif()
  set(
    thrift_include_options
    "-I;$<JOIN:$<TARGET_PROPERTY:${LIB_NAME}.thrift_includes,INTERFACE_INCLUDE_DIRECTORIES>,;-I;>"
  )

  string(REPLACE ";" "," GEN_ARG_STR "${ARG_OPTIONS}")

  # mstch_python generates output at gen-python/<namespace_py3>/<basename>/
  # The CMake NAMESPACE parameter should match the "namespace py3" directive
  # in the thrift file (consistent with BUCK's py3_namespace).  The basename
  # is appended automatically by this function.

  # Emit the rule to run the thrift compiler
  add_custom_command(
    OUTPUT
      ${generated_sources}
    COMMAND_EXPAND_LISTS
    COMMAND
      "${CMAKE_COMMAND}" -E make_directory "${output_dir}"
    COMMAND
      "${FBTHRIFT_COMPILER}"
      --gen "mstch_python:${GEN_ARG_STR}"
      "${thrift_include_options}"
      -I "${FBTHRIFT_INCLUDE_DIR}"
      -o "${output_dir}"
      "${CMAKE_CURRENT_SOURCE_DIR}/${THRIFT_FILE}"
    WORKING_DIRECTORY
      "${CMAKE_BINARY_DIR}"
    MAIN_DEPENDENCY
      "${THRIFT_FILE}"
    DEPENDS
      "${FBTHRIFT_COMPILER}"
  )

  # Register as a python library.
  # No FBThrift::thrift_py here — thrift-python does not use the legacy
  # pure-Python thrift runtime.  The C-extension runtime (thrift.python,
  # folly) is resolved at runtime by add_fb_thrift_python_executable().
  add_fb_python_library(
    "${LIB_NAME}"
    BASE_DIR "${output_dir}/gen-python"
    NAMESPACE ""
    SOURCES ${generated_sources}
    DEPENDS ${ARG_DEPENDS}
  )
endfunction()

# ---------------------------------------------------------------------------
# add_fb_thrift_python_executable()
# ---------------------------------------------------------------------------
# Drop-in replacement for add_fb_python_executable() for targets that
# transitively import thrift_types.py.
#
# Differences from add_fb_python_executable():
#   1. Forces TYPE dir (C extensions cannot load from zip archives)
#   2. Post-build: symlinks thrift/python/ from the pip-installed
#      fbthrift-python wheel into the bundled thrift/ directory so that
#      both thrift.Thrift (legacy) and thrift.python (new) coexist
#   3. Post-build: symlinks folly/ from folly-python site-packages
#
# Usage is identical to add_fb_python_executable():
#   add_fb_thrift_python_executable(
#       my_target
#       MAIN_MODULE my.module:main
#       SOURCES ${SRCS}
#       DEPENDS some_py_lib FBThrift::thrift_py
#   )
# ---------------------------------------------------------------------------
function(add_fb_thrift_python_executable TARGET)
  fb_find_thrift_python_runtime()

  # Forward all arguments, injecting TYPE dir
  add_fb_python_executable(
    ${TARGET}
    TYPE dir
    ${ARGN}
  )

  set(target_dir "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}")

  # Post-build symlinks resolve C extension imports at runtime.
  # Expected directory layout after build:
  #   <target>/
  #   ├── __main__.py
  #   ├── thrift/
  #   │   ├── __init__.py
  #   │   └── python/ → <pip-site-packages>/thrift/python/
  #   ├── folly/ → <getdeps-site-packages>/folly/
  #   ├── <your_module>/...            (Python sources)
  #   └── <namespace>/<module>/thrift_types.py  (generated thrift-python)
  #
  # Since we no longer depend on FBThrift::thrift_py (the legacy pure-Python
  # runtime), there may be no bundled thrift/ directory yet.  Create it
  # with an __init__.py so that "import thrift.python" works.
  if(THRIFT_PYTHON_SITE_PACKAGES AND EXISTS "${THRIFT_PYTHON_SITE_PACKAGES}/thrift/python")
    add_custom_command(
      TARGET ${TARGET}.GEN_PY_EXE POST_BUILD
      COMMAND ${CMAKE_COMMAND} -E make_directory "${target_dir}/thrift"
      COMMAND ${CMAKE_COMMAND} -E touch "${target_dir}/thrift/__init__.py"
      COMMAND ${CMAKE_COMMAND} -E create_symlink
        "${THRIFT_PYTHON_SITE_PACKAGES}/thrift/python"
        "${target_dir}/thrift/python"
      COMMENT "Symlinking thrift.python C extensions into ${TARGET}"
    )
  else()
    message(WARNING
      "${TARGET}: thrift.python not found at "
      "${THRIFT_PYTHON_SITE_PACKAGES}/thrift/python — "
      "thrift_types.py imports will fail at runtime")
  endif()

  # Symlink folly C extensions (folly.iobuf etc.)
  if(FOLLY_PYTHON_SITE_PACKAGES AND EXISTS "${FOLLY_PYTHON_SITE_PACKAGES}/folly")
    add_custom_command(
      TARGET ${TARGET}.GEN_PY_EXE POST_BUILD
      COMMAND ${CMAKE_COMMAND} -E create_symlink
        "${FOLLY_PYTHON_SITE_PACKAGES}/folly"
        "${target_dir}/folly"
      COMMENT "Symlinking folly-python C extensions into ${TARGET}"
    )
  else()
    message(WARNING
      "${TARGET}: folly-python not found — "
      "thrift_types.py imports will fail at runtime")
  endif()
endfunction()