from conan import ConanFile
from conan.errors import ConanInvalidConfiguration
from conan.tools.build import check_min_cppstd
from conan.tools.cmake import CMake, CMakeToolchain, CMakeDeps, cmake_layout
from conan.tools.env import VirtualBuildEnv
from conan.tools.files import get, save, copy, export_conandata_patches, apply_conandata_patches
from conan.tools.scm import Version
from os.path import join
import textwrap

required_conan_version = ">=1.53.0"


class TensorflowLiteConan(ConanFile):
    name = "tensorflow-lite"
    user = "aurora"
    license = "Apache-2.0"
    homepage = "https://www.tensorflow.org/lite/guide"
    url = "https://gitverse.ru/auroraos/mirror_conan-center-index"
    description = ("TensorFlow Lite is a set of tools that enables on-device machine learning "
                   "by helping developers run their models on mobile, embedded, and IoT devices.")
    topics = ("machine-learning", "neural-networks", "deep-learning")
    package_type = "library"
    settings = "os", "arch", "compiler", "build_type"
    options = {
        "shared": [True, False],
        "with_gpu": [True, False]
    }
    default_options = {
        "shared": False,
        "with_gpu": True
    }

    short_paths = True

    @property
    def _min_cppstd(self):
        return "17"

    @property
    def _compilers_minimum_version(self):
        return {
            "gcc": "8",
            "Visual Studio": "15.8",
            "msvc": "191",
            "clang": "5",
            "apple-clang": "5.1",
        }

    def export_sources(self):
        export_conandata_patches(self)

    def layout(self):
        cmake_layout(self, src_folder="src")

    def requirements(self):
        self.requires("abseil/20240116.1@aurora")
        self.requires("eigen/3.4.0@aurora")
        self.requires("farmhash/cci.20190513@aurora")
        self.requires("fft/cci.20061228@aurora")
        if Version(self.version) < "2.15.0":
            self.requires("flatbuffers/23.3.3@aurora", transitive_headers=True)
        else:
            self.requires("flatbuffers/23.5.26@aurora", transitive_headers=True)
        self.requires("gemmlowp/cci.20210928@aurora")
        self.requires("ruy/cci.20231129@aurora")
        if self.settings.arch in ("x86", "x86_64"):
            self.requires("intel-neon2sse/cci.20210225@aurora")
        if Version(self.version) >= "2.12.0":
            self.requires("pthreadpool/cci.20231129@aurora")
            self.requires("fxdiv/cci.20200417@aurora")
        if self.options.with_gpu:
            self.requires("vulkan-headers/1.3.290.0@aurora")
            self.requires("opencl-headers/2023.12.14@aurora")
            self.requires("fp16/cci.20210320@aurora")

    def validate(self):
        if self.settings.get_safe("compiler.cppstd"):
            check_min_cppstd(self, self._min_cppstd)

        minimum_version = self._compilers_minimum_version.get(str(self.settings.compiler), False)
        if minimum_version and Version(self.settings.compiler.version) < minimum_version:
            raise ConanInvalidConfiguration(
                f"{self.ref} requires C++{self._min_cppstd}, which your compiler does not support."
            )

    def source(self):
        get(self, **self.conan_data["sources"][self.version], strip_root=True)

    def generate(self):
        env = VirtualBuildEnv(self)
        env.generate()
        tc = CMakeToolchain(self)
        tc.variables.update({
            "TFLITE_ENABLE_RUY": True,
            "TFLITE_ENABLE_NNAPI": False,
            "TFLITE_ENABLE_GPU": self.options.with_gpu,
            "TFLITE_ENABLE_XNNPACK": False,
            "TFLITE_ENABLE_MMAP": False,
            "FETCHCONTENT_FULLY_DISCONNECTED": True,
            "SYSTEM_PTHREADPOOL": True,
            "SYSTEM_FARMHASH": True,
            "clog_POPULATED": True,
            "CMAKE_C_FLAGS": "-funsafe-math-optimizations -DCL_TARGET_OPENCL_VERSION=200 -DCL_DEVICE_PREFERRED_WORK_GROUP_SIZE_MULTIPLE=1",
            "CMAKE_CXX_FLAGS": "-funsafe-math-optimizations -DCL_TARGET_OPENCL_VERSION=200 -DCL_DEVICE_PREFERRED_WORK_GROUP_SIZE_MULTIPLE=1",
            "CMAKE_FIND_PACKAGE_PREFER_CONFIG": True,
            "OVERRIDABLE_FETCH_CONTENT_LICENSE_CHECK": False,
        })
        if self.settings.arch == "armv8":
            # Not defined by Conan for Apple Silicon. See https://github.com/conan-io/conan/pull/8026
            tc.variables["CMAKE_SYSTEM_PROCESSOR"] = "arm64"
        tc.generate()
        deps = CMakeDeps(self)
        deps.generate()

    def build(self):
        apply_conandata_patches(self)
        cmake = CMake(self)
        cmake.configure(build_script_folder=join("tensorflow", "lite"))
        cmake.build()

    @staticmethod
    def _create_cmake_module_alias_target(conanfile, module_file):
        aliased = "tensorflowlite::tensorflowlite"
        alias = "tensorflow::tensorflowlite"
        content = textwrap.dedent(f"""\
                if(TARGET {aliased} AND NOT TARGET {alias})
                    add_library({alias} INTERFACE IMPORTED)
                    set_property(TARGET {alias} PROPERTY INTERFACE_LINK_LIBRARIES {aliased})
                endif()
            """)
        save(conanfile, module_file, content)

    @property
    def _module_file(self):
        return join("lib", "cmake", f"conan-official-{self.name}-targets.cmake")

    def package(self):
        copy(self, "LICENSE", self.source_folder, join(self.package_folder, "licenses"))
        copy(self, "*.h", join(self.source_folder, "tensorflow", "lite"), join(self.package_folder, "include", "tensorflow", "lite"))
        copy(self, "version.h", join(self.source_folder, "tensorflow", "core", "public"), join(self.package_folder, "include", "tensorflow", "core", "public"))
        copy(self, "*.a", self.build_folder, join(self.package_folder, "lib"), keep_path=False)
        copy(self, "*.so", self.build_folder, join(self.package_folder, "lib"), keep_path=False)
        self._create_cmake_module_alias_target(self, join(self.package_folder, self._module_file))

    def package_info(self):
        self.cpp_info.set_property("cmake_file_name", "tensorflowlite")
        self.cpp_info.set_property("cmake_target_name", "tensorflow::tensorflowlite")

        self.cpp_info.names["cmake_find_package"] = "tensorflowlite"
        self.cpp_info.names["cmake_find_package_multi"] = "tensorflowlite"
        self.cpp_info.build_modules["cmake_find_package"] = [self._module_file]
        self.cpp_info.build_modules["cmake_find_package_multi"] = [self._module_file]

        defines = []
        if not self.options.shared:
            defines.append("TFL_STATIC_LIBRARY_BUILD")

        self.cpp_info.defines = defines
        self.cpp_info.libs = ["tensorflow-lite"]
        if self.settings.os in ["Linux", "FreeBSD"]:
            self.cpp_info.system_libs.append("dl")
