#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.
"""
Generates pyproject.toml and poetry.lock files from requirements.txt

Black Duck requires poetry lock files to perform security scans. TensorRT uses requirements.txt to define
Python dependencies.
This script parses through the requirements.txt files in the project and generates the poetry lock files
required to perform the security scans.

Pip install requirements:
pip3 install poetry

To generate pyproject.toml and poetry.lock recursively for all requirements.txt in the project:
python3 scripts/generate_lock_files.py

To generate pyproject.toml and poetry.lock for a single requirements.txt:
python3 scripts/generate_lock_files.py --path <path>/requirements.txt
"""

import argparse
import json
import os
import re
import shutil
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path

sys.path.insert(0, os.getcwd())

FOLDER_SECURITY_SCANNING = "security_scanning"

url_mapping = {}

target_env = {
    "python_version": 310,
    "platform_system": "",
    "platform_machine": "x86_64",
    "sys_platform": "linux",
}


def get_project_info(path: str):
    path_project = re.sub(rf"^{os.getcwd()}\/?", "", path)
    name = "unknown-package"
    version = "0.1.0"
    if not path_project:
        name = "tensorrt-llm"
        import importlib.util

        # get trtllm version from tensorrt_llm/version.py
        module_path = os.path.join("tensorrt_llm", "version.py")
        spec = importlib.util.spec_from_file_location("trtllm_version",
                                                      module_path)
        if spec and spec.loader:
            version_module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(version_module)
            version = version_module.__version__
    else:
        matches = re.match(r"^(?:([\w\-]+)?\/)?([\w\-]+)$", path_project)
        if matches:
            if matches.group(1):
                name = f"{matches.group(2)}-{matches.group(1)}"
            else:
                name = matches.group(2)
    return {"name": name, "version": version}


def generate_metadata_json():
    try:
        commit_hash = subprocess.check_output(["git", "rev-parse", "HEAD"],
                                              text=True).strip()
    except subprocess.CalledProcessError as e:
        print(f"Error retrieving git commit hash: {e}")
        raise

    data = {
        "commit_hash": commit_hash,
        "timestamp": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    }
    with open(f"{FOLDER_SECURITY_SCANNING}/metadata.json",
              "w",
              encoding="utf-8") as f:
        json.dump(data, f, indent=2)
        f.write("\n")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Lock files generator",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("-p", "--path", help="Path to requirements.txt")
    args, _ = parser.parse_known_args()

    if args.path:
        _, filename = os.path.split(args.path)
        assert filename == 'requirements.txt'
        realpath = Path(args.path).resolve()
        paths = [realpath]
    else:
        # get paths to all files names requirements.txt
        paths = Path.cwd().rglob('requirements.txt')

    if os.path.exists(FOLDER_SECURITY_SCANNING):
        shutil.rmtree(FOLDER_SECURITY_SCANNING)
    os.mkdir(FOLDER_SECURITY_SCANNING)
    generate_metadata_json()

    # generate pyproject.toml and poetry.lock files in the same location
    for path in paths:
        file_path, file_name = os.path.split(path)
        curr_path = Path.cwd()
        if "3rdparty" in file_path:
            continue

        # init poetry
        save_path = os.path.join(FOLDER_SECURITY_SCANNING,
                                 Path(file_path).relative_to(curr_path))
        os.makedirs(save_path, exist_ok=True)
        print(f"Initializing PyProject.toml in {file_path}")
        project_info = get_project_info(file_path)
        name = project_info["name"]
        author = '"TensorRT LLM [90828364+tensorrt-cicd@users.noreply.github.com]"'
        version = project_info["version"]
        py_version = '">=3.10,<3.13"'
        poetry_init_cmd = f'poetry init --no-interaction --name {name} --author {author} --python {py_version}'
        if name == "tensorrt-llm":
            poetry_init_cmd += " -l Apache-2.0"
        subprocess.run(poetry_init_cmd, shell=True, cwd=save_path)
        if version != "0.1.0":
            subprocess.run(f"poetry version {version}",
                           shell=True,
                           cwd=file_path)
        output = subprocess.run(f'cat {path}',
                                shell=True,
                                capture_output=True,
                                text=True).stdout
        packages = output.split('\n')

        if packages[-1] == '':  # last entry is newline
            packages = packages[:-1]

        for package in packages:
            # WAR: ignore lines with "-f": No tool exists to parse complex requirements.txt
            if '-f' in package or \
                "#" in package or \
                package.startswith('--'):
                continue

            curr_env = None
            if ';' in package:
                curr_env = package.split(';')[1]
                curr_env = curr_env.replace("sys.", "sys_")

                # WAR for "3.8" < "3.10" evaluating to False:
                # convert to int and remove decimal 38 < 310 is True
                while '.' in curr_env:
                    py_version_str = curr_env.split(
                        '.')[0][-1] + '.' + curr_env.split('.')[1][0]
                    if curr_env.split('.')[1][1] != '"' and curr_env.split(
                            '.')[1][1] != "'":
                        py_version_str += curr_env.split('.')[1][1]

                    py_version_int = py_version_str.replace('.', '')

                    curr_env = curr_env.replace(f'"{py_version_str}"',
                                                py_version_int)

            if curr_env is None or eval(curr_env, target_env):
                data = package.split(';')
                package_name = data[0].replace(" ", "")
                if (not package_name.startswith("tensorrt_llm")
                        and not package_name.startswith("3rdparty")):
                    poetry_cmd = f"poetry add '{package_name}'"
                    if package_name in url_mapping:
                        poetry_cmd = f"poetry add '{url_mapping[package_name]}'"
                    subprocess.run(poetry_cmd, shell=True, cwd=save_path)