#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
环境构建SDK工具：固定Python 3.10.19，生成uv标准的pyproject.toml源配置
"""
import os
import sys
from typing import Dict, List, Tuple

# 固定Python版本（核心配置）
FIXED_PYTHON_VERSION = "3.10.19"
# 兼容的Python版本范围（用于提示）
COMPATIBLE_PYTHON_RANGE = ">=3.10,<3.11"

# 预定义框架版本映射（新增uv源名称/url映射）
FRAMEWORK_VERSIONS = {
    "pytorch": {
        "2.2.1": {
            "cpu": {
                "deps": {
                    "torch": "2.2.1",
                    "torchvision": "0.17.1",
                    "torchaudio": "2.2.1"
                },
                "install_cmd": "pip install torch==2.2.1+cpu torchvision==0.17.1+cpu torchaudio==2.2.1+cpu -f https://download.pytorch.org/whl/cpu/torch_stable.html",
                "uv_index_name": "pytorch-official",
                "uv_index_url": "https://download.pytorch.org/whl/cpu",
                "cuda_alias": "cpu"
            },
            "gpu": {
                "cuda_version": "12.1",
                "deps": {
                    "torch": "2.2.1",
                    "torchvision": "0.17.1",
                    "torchaudio": "2.2.1"
                },
                "install_cmd": "pip install torch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 --index-url https://download.pytorch.org/whl/cu121",
                "uv_index_name": "pytorch-official",
                "uv_index_url": "https://download.pytorch.org/whl/cu121",
                "cuda_alias": "cu121"
            }
        },
        "2.5.1": {
            "cpu": {
                "deps": {
                    "torch": "2.5.1",
                    "torchvision": "0.20.1",
                    "torchaudio": "2.5.1"
                },
                "install_cmd": "pip install torch==2.5.1+cpu torchvision==0.20.1+cpu torchaudio==2.5.1+cpu -f https://download.pytorch.org/whl/cpu/torch_stable.html",
                "uv_index_name": "pytorch-official",
                "uv_index_url": "https://download.pytorch.org/whl/cpu",
                "cuda_alias": "cpu"
            },
            "gpu": {
                "cuda_version": "12.1",
                "deps": {
                    "torch": "2.5.1",
                    "torchvision": "0.20.1",
                    "torchaudio": "2.5.1"
                },
                "install_cmd": "conda install pytorch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 pytorch-cuda=12.1 -c pytorch -c nvidia",
                "uv_index_name": "pytorch-official",
                "uv_index_url": "https://download.pytorch.org/whl/cu121",
                "cuda_alias": "cu121",
                "conda_channels": ["pytorch", "nvidia"]
            }
        }
    },
    "paddle": {
        "2.5.2": {
            "cpu": {
                "deps": {
                    "paddlepaddle": "2.5.2",
                    "paddlenlp": "2.5.2"
                },
                "install_cmd": "pip install paddlepaddle==2.5.2 -i https://pypi.tuna.tsinghua.edu.cn/simple",
                "uv_index_name": "paddle-official",
                "uv_index_url": "https://pypi.tuna.tsinghua.edu.cn/simple",
                "cuda_alias": "cpu"
            },
            "gpu": {
                "cuda_version": "11.7",
                "deps": {
                    "paddlepaddle-gpu": "2.5.2",
                    "paddlenlp": "2.5.2"
                },
                "install_cmd": "pip install paddlepaddle-gpu==2.5.2.post117 -f https://www.paddlepaddle.org.cn/whl/linux/mkl/avx/stable.html",
                "uv_index_name": "paddle-official",
                "uv_index_url": "https://www.paddlepaddle.org.cn/whl/linux/mkl/avx/stable.html",
                "cuda_alias": "cu117"
            }
        }
    }
}

# 强制内置依赖
MANDATORY_DEPS = {
    "ray[serve]": "2.54.0"
}


class EnvBuilderSDK:
    def __init__(self):
        self.selected_framework = None
        self.selected_version = None
        self.selected_device = None  # cpu/gpu
        self.framework_deps = {}
        self.requirements_deps = {}
        self.all_deps = {}
        self.install_config = {}  # 存储uv源配置、conda通道等

    def check_python_version(self):
        """检查当前Python版本，给出兼容性提示"""
        current_ver = ".".join(map(str, sys.version_info[:3]))
        print(f"\n📌 本工具要求Python版本：{FIXED_PYTHON_VERSION}")
        print(f"   当前Python版本：{current_ver}")
        if current_ver != FIXED_PYTHON_VERSION:
            print(f"⚠️  版本不匹配！建议使用conda创建指定版本环境：")
            print(f"   conda create -n farmland python={FIXED_PYTHON_VERSION}")
            print(f"   conda activate farmland")
        else:
            print("✅ Python版本符合要求！")

    def select_framework(self) -> str:
        """交互式选择框架"""
        print("\n===== 选择深度学习框架 =====")
        frameworks = list(FRAMEWORK_VERSIONS.keys())
        for idx, fw in enumerate(frameworks, 1):
            print(f"{idx}. {fw.upper()}")

        while True:
            try:
                choice = int(input(f"\n请输入框架序号(1-{len(frameworks)}): "))
                if 1 <= choice <= len(frameworks):
                    self.selected_framework = frameworks[choice - 1]
                    print(f"已选择框架：{self.selected_framework.upper()}")
                    return self.selected_framework
                else:
                    print(f"请输入1-{len(frameworks)}之间的数字！")
            except ValueError:
                print("请输入有效的数字！")

    def select_device_type(self) -> str:
        """选择CPU/GPU版本"""
        print(f"\n===== 选择{self.selected_framework.upper()}运行设备 =====")
        devices = ["cpu", "gpu"]
        for idx, dev in enumerate(devices, 1):
            print(f"{idx}. {dev.upper()}")

        while True:
            try:
                choice = int(input(f"\n请输入设备序号(1-{len(devices)}): "))
                if 1 <= choice <= len(devices):
                    self.selected_device = devices[choice - 1]
                    print(f"已选择设备：{self.selected_device.upper()}")
                    # 提示CUDA版本（GPU时）
                    if self.selected_device == "gpu":
                        cuda_ver = FRAMEWORK_VERSIONS[self.selected_framework][self.selected_version][
                            self.selected_device].get("cuda_version")
                        if cuda_ver:
                            print(f"⚠️  该版本需要CUDA {cuda_ver}，请确保系统已安装对应版本CUDA")
                    return self.selected_device
                else:
                    print(f"请输入1-{len(devices)}之间的数字！")
            except ValueError:
                print("请输入有效的数字！")

    def select_framework_version(self) -> str:
        """交互式选择框架版本（先选版本，再选CPU/GPU）"""
        print(f"\n===== 选择{self.selected_framework.upper()}版本 =====")
        versions = list(FRAMEWORK_VERSIONS[self.selected_framework].keys())
        for idx, ver in enumerate(versions, 1):
            print(f"{idx}. 版本{ver}")

        while True:
            try:
                choice = int(input(f"\n请输入版本序号(1-{len(versions)}): "))
                if 1 <= choice <= len(versions):
                    self.selected_version = versions[choice - 1]
                    print(f"已选择{self.selected_framework.upper()}版本：{self.selected_version}")

                    # 选择CPU/GPU
                    self.select_device_type()

                    # 获取该版本+设备的依赖和安装配置
                    self.framework_deps = \
                    FRAMEWORK_VERSIONS[self.selected_framework][self.selected_version][self.selected_device]["deps"]
                    self.install_config = FRAMEWORK_VERSIONS[self.selected_framework][self.selected_version][
                        self.selected_device]

                    # 打印安装命令参考
                    print(f"\n📝 参考安装命令：{self.install_config['install_cmd']}")
                    return self.selected_version
                else:
                    print(f"请输入1-{len(versions)}之间的数字！")
            except ValueError:
                print("请输入有效的数字！")

    def read_requirements(self, file_path: str = "requirements.txt") -> Dict[str, str]:
        """读取requirements.txt文件，解析依赖"""
        self.requirements_deps = {}
        if not os.path.exists(file_path):
            print(f"\n⚠️  未找到{file_path}文件，跳过读取第三方依赖")
            return self.requirements_deps

        print(f"\n===== 读取{file_path}依赖 =====")
        with open(file_path, "r", encoding="utf-8") as f:
            lines = f.readlines()

        for line in lines:
            line = line.strip()
            # 跳过注释和空行
            if not line or line.startswith("#"):
                continue
            # 解析包名和版本（支持==/>=/<=等）
            if "==" in line:
                pkg, ver = line.split("==", 1)
                self.requirements_deps[pkg.strip()] = ver.strip()
            else:
                # 无版本号时保留原字符串
                self.requirements_deps[line.strip()] = ""

        print(f"成功读取{len(self.requirements_deps)}个依赖：{list(self.requirements_deps.keys())}")
        return self.requirements_deps

    def merge_dependencies(self) -> Dict[str, str]:
        """合并所有依赖：框架依赖 + requirements依赖 + 强制依赖"""
        self.all_deps = {}
        # 1. 框架依赖（优先级最高，避免版本冲突）
        self.all_deps.update(self.framework_deps)
        # 2. requirements依赖（若与框架依赖冲突，以框架为准）
        for pkg, ver in self.requirements_deps.items():
            if pkg not in self.all_deps:
                self.all_deps[pkg] = ver
        # 3. 强制依赖（ray[serve]，优先级最高）
        self.all_deps.update(MANDATORY_DEPS)

        print(f"\n===== 合并后总依赖 =====")
        for pkg, ver in self.all_deps.items():
            print(f"{pkg}: {ver if ver else '未指定版本'}")
        return self.all_deps

    def generate_pyproject_toml(self, output_path: str = "pyproject.toml") -> None:
        """生成pyproject.toml文件（固定Python 3.10.19 + uv源配置）"""
        # 基础配置
        pyproject_content = f"""
[project]
name = "tool"
version = "0.1.0"
description = "（Python {FIXED_PYTHON_VERSION} | {self.selected_framework.upper()} {self.selected_version} {self.selected_device.upper()}）"
authors = [{{ name="Your Name", email="your@email.com" }}]
requires-python = "{COMPATIBLE_PYTHON_RANGE}"  # 强制Python 3.10.x，精确版本{FIXED_PYTHON_VERSION}
license = {{ file="LICENSE" }}

# 核心依赖（包含{self.selected_framework} {self.selected_device}版本 + 第三方依赖）
dependencies = [
"""
        # 添加所有依赖到dependencies列表
        dep_lines = []
        for pkg, ver in self.all_deps.items():
            if ver:
                dep_lines.append(f'    "{pkg}=={ver}",')
            else:
                dep_lines.append(f'    "{pkg}",')

        pyproject_content += "\n".join(dep_lines)
        pyproject_content += """
]

"""
        # ========== uv标准的源配置 ==========
        pyproject_content += f"""# 新增 {self.selected_framework} 官方源（{self.install_config['cuda_alias']} 版本）
[[tool.uv.index]]
name = "{self.install_config['uv_index_name']}"
url = "{self.install_config['uv_index_url']}"
explicit = true

"""
        # 针对PyTorch：强制torch/torchaudio/torchvision从官方源下载
        if self.selected_framework == "pytorch":
            pyproject_content += f"""# 关键：指定PyTorch相关包强制从 {self.install_config['uv_index_name']} 源下载
[tool.uv.sources]
torch = {{ index = "{self.install_config['uv_index_name']}" }}
torchaudio = {{ index = "{self.install_config['uv_index_name']}" }}
torchvision = {{ index = "{self.install_config['uv_index_name']}" }}

"""
        # 针对Paddle：强制Paddle相关包从官方源下载
        elif self.selected_framework == "paddle":
            pyproject_content += f"""# 关键：指定Paddle相关包强制从 {self.install_config['uv_index_name']} 源下载
[tool.uv.sources]
paddlepaddle = {{ index = "{self.install_config['uv_index_name']}" }}
paddlepaddle-gpu = {{ index = "{self.install_config['uv_index_name']}" }}
paddlenlp = {{ index = "{self.install_config['uv_index_name']}" }}

"""
        # 可选开发依赖
        pyproject_content += """[project.optional-dependencies]
dev = [
    "pytest>=7.0",
    "black>=23.0",
    "flake8>=6.0"
]
"""

        # 写入文件
        with open(output_path, "w", encoding="utf-8") as f:
            f.write(pyproject_content)

        print(f"\n✅ 成功生成{output_path}文件！")
        print(f"文件路径：{os.path.abspath(output_path)}")
        # 提示关键信息
        print(f"📌 强制指定Python版本：{FIXED_PYTHON_VERSION}")
        if self.selected_device == "gpu":
            cuda_ver = self.install_config.get("cuda_version")
            print(f"⚠️  注意：该配置需要CUDA {cuda_ver}，请确保系统已安装对应版本CUDA")

    def run(self):
        """SDK主执行流程"""
        print("===== 环境构建SDK（Python 3.10.19固定版）=====")
        # 前置步骤：检查Python版本
        self.check_python_version()
        # 步骤1：选择框架
        self.select_framework()
        # 步骤2：选择框架版本 + CPU/GPU
        self.select_framework_version()
        # 步骤3：读取requirements.txt
        self.read_requirements()
        # 步骤4：合并依赖
        self.merge_dependencies()
        # 步骤5：生成pyproject.toml
        self.generate_pyproject_toml()
        print("\n===== SDK执行完成 =====")

# 核心新增：定义main函数，供raykit.build调用
def main():
    """环境构建的入口函数（供模块导入调用）"""
    sdk = EnvBuilderSDK()
    sdk.run()

# 保留原有直接运行的逻辑
if __name__ == "__main__":
    main()