import os
import subprocess
import shutil
from conan import ConanFile
from conan.tools.cmake import CMakeToolchain, CMakeDeps, cmake_layout
from conan.tools.files import copy, rmdir, patch
from conan.tools.scm import Git
from conan.errors import ConanInvalidConfiguration, ConanException

class HesaiLidarDriver(ConanFile):
    name = "hesai-lidar-ros2-jazzy-driver"
    description = "ROS2 (Jazzy) Driver for Hesai LiDAR sensor"
    license = "BSD-3-Clause"
    url = "https://github.com/HesaiTechnology/HesaiLidar_ROS_2.0"
    author = "thommyho1988@gmail.com"
    topics = ("ros2", "lidar", "hesai", "driver")

    settings = "os", "compiler", "build_type", "arch"

    options = {
        "with_cuda": [True, False],
    }
    default_options = {
        "with_cuda": False,
    }

    def export_sources(self):
        copy(self, "patches/**", self.recipe_folder, self.export_sources_folder)
        copy(self, "LICENSE", self.recipe_folder, self.export_sources_folder)
        copy(self, "README.md", self.recipe_folder, self.export_sources_folder)

    def layout(self):
        cmake_layout(self)

    def build_requirements(self):
        self.tool_requires("make/4.4.1")
        self.tool_requires("cmake/3.31.9")
        self.tool_requires("ros2-jazzy-toolchain/latest")

    def requirements(self):
        self.requires("ros2-jazzy-python/latest")
        self.requires("ros2-jazzy-toolchain/latest")
        self.requires("openssl/1.1.1w")
        self.requires("yaml-cpp/0.8.0")
        self.requires("tinyxml2/7.1.0")
        self.requires("boost/1.74.0", options={
            "without_thread": False,
            "without_system": False,
            "without_chrono": False,
        })

    def validate(self):
        if self.options.with_cuda and self.settings.arch not in ["x86_64", "armv8"]:
            raise ConanInvalidConfiguration("CUDA is only supported on x86_64 and armv8")

    def source(self):
        if self.version not in self.conan_data.get("sources", {}):
            raise ConanInvalidConfiguration(f"Version '{self.version}' not found in conandata.yml")

        data = self.conan_data["sources"][self.version]
        tmp_src = "src_tmp"

        try:
            self.output.info(f"Cloning {data['url']} into temporary directory...")
            git = Git(self)
            git.clone(url=data["url"], target=tmp_src, args=["--recursive"])

            if "revision" in data:
                self.output.info(f"Checking out revision {data['revision']}")
                git_tmp = Git(self, folder=tmp_src)
                git_tmp.checkout(commit=data["revision"])
                git_tmp.run("submodule update --init --recursive")

            self.output.info("Moving sources to recipe root...")
            copy(self, "*", src=tmp_src, dst=".")

        except Exception as e:
            raise ConanException(f"Source retrieval failed: {e}")
        finally:
            if os.path.exists(tmp_src):
                rmdir(self, tmp_src)

        patch_dir = os.path.join(self.source_folder, "patches", self.version)
        if os.path.exists(patch_dir):
            patches = sorted([p for p in os.listdir(patch_dir) if p.endswith(".patch")])
            for p in patches:
                self.output.info(f"Applying patch: {p}")
                patch(self, patch_file=os.path.join(patch_dir, p))

    def _get_numpy_include(self, python_exe):
        try:
            cmd = [python_exe, "-c", "import numpy; print(numpy.get_include())"]
            return subprocess.check_output(cmd, text=True).strip()
        except (subprocess.CalledProcessError, FileNotFoundError, OSError) as e:
            self.output.warning(f"Failed to determine NumPy include path: {e}")
            return ""

    def generate(self):
        deps = CMakeDeps(self)
        deps.generate()

        tc = CMakeToolchain(self)
        python_dep = self.dependencies["ros2-jazzy-python"]
        python_exe = python_dep.conf_info.get("user.ros2:python_interpreter", check_type=str)

        if python_exe:
            numpy_include = self._get_numpy_include(python_exe)
            tc.variables["Python3_EXECUTABLE"] = python_exe
            tc.variables["Python_EXECUTABLE"] = python_exe
            if numpy_include:
                tc.variables["Python3_NumPy_INCLUDE_DIR"] = numpy_include

        self._configure_special_flags(tc, python_dep)
        tc.generate()

    def _configure_special_flags(self, tc, python_dep):
        cross_blob = False
        if "cross_blob" in python_dep.options:
            cross_blob = bool(python_dep.options.cross_blob)

        if self.options.with_cuda:
            if cross_blob:
                raise ConanInvalidConfiguration("CUDA not available as conan package due to NVIDIA license.")
            self.output.info("Enabling CUDA Support")
            tc.variables["FIND_CUDA"] = "ON"

        if cross_blob:
            self.output.info("Enabling Cross-Blob workarounds (No LTO)")
            tc.variables["CMAKE_INTERPROCEDURAL_OPTIMIZATION"] = "OFF"
            tc.cache_variables["CMAKE_C_FLAGS"] = "-fno-lto"
            tc.cache_variables["CMAKE_CXX_FLAGS"] = "-fno-lto"

    def build(self):
        # 1. Resolve Setup Script
        build_dep = self.dependencies.get("ros2-jazzy-toolchain")
        setup_script = None

        if build_dep:
            env_vars = build_dep.buildenv_info.vars(self)
            raw_cmd = env_vars.get("CONAN_ROS2_SOURCE_CMD", "")
            if "source " in raw_cmd:
                setup_script = raw_cmd.split("source ")[1].strip().split(" ")[0]
            elif raw_cmd.strip().endswith((".sh", ".bash")):
                setup_script = raw_cmd.strip()

        # 2. Paths
        tc_file = os.path.join(self.generators_folder, "conan_toolchain.cmake")
        abs_build_base = os.path.join(self.build_folder, "colcon_build")
        abs_install_base = os.path.join(self.build_folder, "install")

        # 3. Colcon Command
        colcon_cmd = (
            f"colcon build --merge-install "
            f"--build-base '{abs_build_base}' "
            f"--install-base '{abs_install_base}' "
            f"--event-handlers console_direct+ "
            f"--cmake-args -DCMAKE_TOOLCHAIN_FILE='{tc_file}' "
            f"-DCMAKE_BUILD_TYPE={self.settings.build_type}"
        )

        # 4. Execute (Bash wrapper for 'source')
        if setup_script:
            full_cmd = f'/bin/bash -c "source {setup_script} && {colcon_cmd}"'
        else:
            full_cmd = colcon_cmd

        self.output.info(f"Building with: {full_cmd}")
        self.run(full_cmd, shell=True, cwd=self.source_folder)

    def package(self):
        install_dir = os.path.join(self.build_folder, "install")
        if not os.path.exists(install_dir):
            raise ConanException(f"Build failed to produce install directory: {install_dir}")

        # 1. Define Subdirectory
        ros_package_dir = os.path.join(self.package_folder, "hesai_ros_driver")

        # 2. Copy ROS Artifacts to Subdirectory
        copy(self, "*", src=install_dir, dst=ros_package_dir)

        # 3. Copy Metadata to Root
        copy(self, "LICENSE", src=self.source_folder, dst=self.package_folder)
        copy(self, "README.md", src=self.source_folder, dst=self.package_folder)

        # 4. Cleanup
        rmdir(self, os.path.join(ros_package_dir, "share", "doc"))
        rmdir(self, os.path.join(ros_package_dir, "colcon_build"))

    def package_info(self):
        # 1. Define the ROS Root inside the package
        ros_root = os.path.join(self.package_folder, "hesai_ros_driver")

        # 2. Cpp Info (Relative to package_folder)
        self.cpp_info.libs = ["hesai_ros_driver"]
        self.cpp_info.libdirs = [os.path.join("hesai_ros_driver", "lib")]
        self.cpp_info.includedirs = [os.path.join("hesai_ros_driver", "include")]
        self.cpp_info.bindirs = [os.path.join("hesai_ros_driver", "bin")]

        # 3. Environment Variables (Absolute Paths)
        # Critical: These ensure ROS 2 finds the package in the subdirectory
        self.runenv_info.prepend_path("AMENT_PREFIX_PATH", ros_root)
        self.buildenv_info.prepend_path("AMENT_PREFIX_PATH", ros_root)

        self.runenv_info.prepend_path("CMAKE_PREFIX_PATH", ros_root)
        self.buildenv_info.prepend_path("CMAKE_PREFIX_PATH", ros_root)

        self.runenv_info.prepend_path("PATH", os.path.join(ros_root, "bin"))
        self.runenv_info.prepend_path("LD_LIBRARY_PATH", os.path.join(ros_root, "lib"))

        # 4. Python Path Calculation
        lib_dir = os.path.join(ros_root, "lib")
        if os.path.exists(lib_dir):
            for item in os.listdir(lib_dir):
                if item.startswith("python"):
                    site_packages = os.path.join(lib_dir, item, "site-packages")
                    if os.path.exists(site_packages):
                        self.runenv_info.prepend_path("PYTHONPATH", site_packages)
                        self.buildenv_info.prepend_path("PYTHONPATH", site_packages)
                        break
