# Copyright 2026 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================

"""
Repository rule to set python version and wheel name.

Version can be set via build parameter "--repo_env=TF_PYTHON_VERSION=3.10"
Defaults to 3.10.

To set wheel name, add "--repo_env=WHEEL_NAME=tensorflow_cpu"
"""

VERSIONS = ["3.10", "3.11", "3.12", "3.13", "3.14"]
DEFAULT_VERSION = "3.11"
WARNING = """
TF_PYTHON_VERSION environment variable was not set correctly; using Python {}.

To set Python version, run:
export TF_PYTHON_VERSION=3.11
""".format(DEFAULT_VERSION)

content = """
TF_PYTHON_VERSION = "{}"
HERMETIC_PYTHON_VERSION = "{}"
WHEEL_NAME = "{}"
WHEEL_COLLAB = "{}"
"""

def _python_repository_impl(repository_ctx):
    repository_ctx.file("BUILD", "")
    version = repository_ctx.os.environ.get("TF_PYTHON_VERSION", "")
    wheel_name = repository_ctx.os.environ.get("WHEEL_NAME", "tensorflow")
    wheel_collab = repository_ctx.os.environ.get("WHEEL_COLLAB", False)
    if version not in VERSIONS:
        print(WARNING)  # buildifier: disable=print
        version = DEFAULT_VERSION
    repository_ctx.file(
        "py_version.bzl",
        content.format(version, version, wheel_name, wheel_collab),
    )

python_repository = repository_rule(
    implementation = _python_repository_impl,
    environ = ["TF_PYTHON_VERSION", "WHEEL_NAME", "WHEEL_COLLAB"],
)