bttoxin-pipeline/scripts/pixi_runner.py

#!/usr/bin/env python3
"""PixiRunner - 在 pixi 环境中执行 BtToxin_Digger 的运行器。

此模块提供 PixiRunner 类，用于替代 DockerContainerManager，
在本地 pixi 环境中运行 BtToxin_Digger 分析。

Example:
    runner = PixiRunner()
    result = runner.run_bttoxin_digger(
        input_dir=Path("input"),
        output_dir=Path("output"),
        log_dir=Path("logs"),
    )
"""
from __future__ import annotations

import json
import logging
import shutil
import subprocess
from pathlib import Path
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)


class PixiRunnerError(Exception):
    """PixiRunner 相关错误的基类。"""
    pass


class PixiNotInstalledError(PixiRunnerError):
    """pixi 未安装时抛出的错误。"""
    pass


class EnvironmentNotFoundError(PixiRunnerError):
    """pixi 环境不存在时抛出的错误。"""
    pass


class BtToxinDiggerNotFoundError(PixiRunnerError):
    """BtToxin_Digger 未找到时抛出的错误。"""
    pass


class PixiRunner:
    """在 pixi 环境中执行 BtToxin_Digger 的运行器。

    此类提供与 DockerContainerManager 兼容的接口，
    但使用本地 pixi 环境而非 Docker 容器执行命令。
    
    Attributes:
        pixi_project_dir: pixi.toml 所在目录
        env_name: pixi 环境名称，默认为 "digger"
    """

    def __init__(
        self,
        pixi_project_dir: Optional[Path] = None,
        env_name: str = "digger",
    ) -> None:
        """初始化运行器，验证 pixi 和环境可用性。
        
        Args:
            pixi_project_dir: pixi.toml 所在目录，默认为项目根目录
            env_name: pixi 环境名称，默认为 "digger"
            
        Raises:
            PixiNotInstalledError: pixi 未安装
            EnvironmentNotFoundError: pixi 环境不存在
            BtToxinDiggerNotFoundError: BtToxin_Digger 未找到
        """
        self.env_name = env_name
        
        # 确定 pixi 项目目录
        if pixi_project_dir is None:
            # 默认为脚本所在目录的父目录（项目根目录）
            self.pixi_project_dir = Path(__file__).resolve().parents[1]
        else:
            self.pixi_project_dir = Path(pixi_project_dir).resolve()
        
        # 验证 pixi.toml 存在
        pixi_toml = self.pixi_project_dir / "pixi.toml"
        if not pixi_toml.exists():
            raise PixiRunnerError(
                f"pixi.toml 配置文件不存在: {pixi_toml}。"
                "请确保在项目根目录运行，或指定正确的 pixi_project_dir。"
            )

    def check_environment(self) -> Dict[str, Any]:
        """检查 pixi 环境和 BtToxin_Digger 可用性。
        
        Returns:
            包含环境状态的字典:
            - pixi_installed: bool, pixi 是否已安装
            - env_exists: bool, 指定环境是否存在
            - bttoxin_available: bool, BtToxin_Digger 是否可用
            - error: str | None, 错误信息（包含可操作的指导）
        """
        result: Dict[str, Any] = {
            "pixi_installed": False,
            "env_exists": False,
            "bttoxin_available": False,
            "error": None,
        }
        
        # 检查 pixi 是否安装
        try:
            proc = subprocess.run(
                ["pixi", "--version"],
                capture_output=True,
                text=True,
                timeout=10,
            )
            if proc.returncode == 0:
                result["pixi_installed"] = True
            else:
                result["error"] = (
                    "pixi 未安装。请访问 https://pixi.sh 安装 pixi，"
                    "或运行: curl -fsSL https://pixi.sh/install.sh | bash"
                )
                return result
        except FileNotFoundError:
            result["error"] = (
                "pixi 未安装。请访问 https://pixi.sh 安装 pixi，"
                "或运行: curl -fsSL https://pixi.sh/install.sh | bash"
            )
            return result
        except subprocess.TimeoutExpired:
            result["error"] = "pixi 命令超时。请检查 pixi 安装是否正确。"
            return result
        except Exception as e:
            result["error"] = f"检查 pixi 时出错: {e}。请确保 pixi 已正确安装。"
            return result
        
        # 检查环境是否存在
        try:
            proc = subprocess.run(
                ["pixi", "info", "--json"],
                capture_output=True,
                text=True,
                timeout=30,
                cwd=str(self.pixi_project_dir),
            )
            if proc.returncode == 0:
                try:
                    info = json.loads(proc.stdout)
                    envs = info.get("environments_info", [])
                    env_names = [e.get("name") for e in envs if e.get("name")]
                    if self.env_name in env_names:
                        result["env_exists"] = True
                    else:
                        result["error"] = (
                            f"{self.env_name} 环境不存在。"
                            f"请在项目根目录运行 'pixi install' 安装依赖。"
                        )
                        return result
                except json.JSONDecodeError:
                    # 回退到简单检查
                    result["env_exists"] = True
            else:
                result["error"] = (
                    f"{self.env_name} 环境不存在。"
                    f"请在项目根目录运行 'pixi install' 安装依赖。"
                )
                return result
        except Exception as e:
            result["error"] = f"检查环境时出错: {e}。请运行 'pixi install' 安装依赖。"
            return result
        
        # 检查 BtToxin_Digger 是否可用
        try:
            proc = subprocess.run(
                ["pixi", "run", "-e", self.env_name, "BtToxin_Digger", "--help"],
                capture_output=True,
                text=True,
                timeout=60,
                cwd=str(self.pixi_project_dir),
            )
            output = proc.stdout + proc.stderr
            if proc.returncode == 0 or "BtToxin_Digger" in output or "Usage" in output:
                result["bttoxin_available"] = True
            else:
                result["error"] = (
                    "BtToxin_Digger 未找到。请确保 pixi install 已完成，"
                    "并且 pixi.toml 中包含 bttoxin_digger 依赖。"
                )
                return result
        except FileNotFoundError:
            result["error"] = (
                "pixi 命令未找到。请确保 pixi 已正确安装并在 PATH 中。"
            )
            return result
        except subprocess.TimeoutExpired:
            result["error"] = (
                "检查 BtToxin_Digger 超时。请运行 'pixi install' 确保环境已正确安装。"
            )
            return result
        except Exception as e:
            result["error"] = (
                f"检查 BtToxin_Digger 时出错: {e}。"
                "请运行 'pixi install' 安装依赖。"
            )
            return result
        
        return result


    def build_digger_command(
        self,
        input_dir: Path,
        sequence_type: str = "nucl",
        scaf_suffix: str = ".fna",
        threads: int = 4,
        **kwargs: Any,
    ) -> List[str]:
        """构建 BtToxin_Digger 命令列表。
        
        Args:
            input_dir: 输入文件目录
            sequence_type: 序列类型 (nucl/orfs/prot/reads)
            scaf_suffix: scaffold 文件后缀
            threads: 线程数
            **kwargs: 其他 BtToxin_Digger 参数
            
        Returns:
            命令参数列表，以 'pixi run -e digger BtToxin_Digger' 开头
        """
        cmd: List[str] = [
            "pixi", "run", "-e", self.env_name,
            "BtToxin_Digger",
            "--SeqPath", str(input_dir),
            "--SequenceType", sequence_type,
            "--threads", str(threads),
        ]
        
        if sequence_type == "nucl":
            cmd += ["--Scaf_suffix", scaf_suffix]
        elif sequence_type == "orfs":
            cmd += ["--orfs_suffix", kwargs.get("orfs_suffix", ".ffn")]
        elif sequence_type == "prot":
            cmd += ["--prot_suffix", kwargs.get("prot_suffix", ".faa")]
        elif sequence_type == "reads":
            platform = kwargs.get("platform", "illumina")
            cmd += ["--platform", platform]
            if platform == "illumina":
                r1 = kwargs.get("reads1_suffix", "_R1.fastq.gz")
                r2 = kwargs.get("reads2_suffix", "_R2.fastq.gz")
                sfx = kwargs.get("suffix_len") or len(r1)
                cmd += ["--reads1", r1, "--reads2", r2, "--suffix_len", str(sfx)]
            elif platform in ("pacbio", "oxford"):
                r = kwargs.get("reads1_suffix", ".fastq.gz")
                gsize = kwargs.get("genome_size", "6.07m")
                sfx = kwargs.get("suffix_len") or len(r)
                cmd += ["--reads1", r, "--genomeSize", gsize, "--suffix_len", str(sfx)]
            elif platform == "hybrid":
                short1 = kwargs.get("short1")
                short2 = kwargs.get("short2")
                long_reads = kwargs.get("long")
                if short1:
                    cmd += ["--short1", short1]
                if short2:
                    cmd += ["--short2", short2]
                if long_reads:
                    cmd += ["--long", long_reads]
                hout = kwargs.get("hout")
                if hout:
                    cmd += ["--hout", hout]
        
        if kwargs.get("assemble_only"):
            cmd.append("--assemble_only")
        
        return cmd

    def run_bttoxin_digger(
        self,
        input_dir: Path,
        output_dir: Path,
        log_dir: Path,
        sequence_type: str = "nucl",
        scaf_suffix: str = ".fna",
        threads: int = 4,
        bttoxin_db_dir: Optional[Path] = None,
        **kwargs: Any,
    ) -> Dict[str, Any]:
        """执行 BtToxin_Digger 分析。
        
        Args:
            input_dir: 输入文件目录
            output_dir: 输出目录
            log_dir: 日志目录
            sequence_type: 序列类型 (nucl/orfs/prot/reads)
            scaf_suffix: scaffold 文件后缀
            threads: 线程数
            bttoxin_db_dir: 外部数据库目录（可选）
            **kwargs: 其他 BtToxin_Digger 参数
            
        Returns:
            结果字典，包含:
            - success: bool, 执行是否成功
            - exit_code: int, 进程退出码
            - logs: str, 执行日志
            - status: str, 状态 (completed/failed/error)
            - error: str (可选), 错误信息
            - output_files: int (可选), 输出文件数量
        """
        # 确保目录存在
        output_dir = Path(output_dir).resolve()
        log_dir = Path(log_dir).resolve()
        input_dir = Path(input_dir).resolve()
        
        output_dir.mkdir(parents=True, exist_ok=True)
        log_dir.mkdir(parents=True, exist_ok=True)
        
        # 准备工作目录结构
        work_input_dir = output_dir / "input_files"
        work_input_dir.mkdir(parents=True, exist_ok=True)
        
        # 复制输入文件
        if sequence_type == "nucl":
            pattern = f"*{scaf_suffix}"
        elif sequence_type == "orfs":
            pattern = f"*{kwargs.get('orfs_suffix', '.ffn')}"
        elif sequence_type == "prot":
            pattern = f"*{kwargs.get('prot_suffix', '.faa')}"
        else:
            pattern = "*"
        
        copied_files = 0
        for f in input_dir.glob(pattern):
            if f.is_file():
                shutil.copy2(f, work_input_dir / f.name)
                copied_files += 1
        
        logger.info(f"已复制 {copied_files} 个输入文件到 {work_input_dir}")
        
        # 设置外部数据库（如果提供）
        if bttoxin_db_dir is not None:
            db_setup_ok = self._setup_external_db(bttoxin_db_dir)
            if not db_setup_ok:
                logger.warning("外部数据库设置失败，将使用默认数据库")
        
        # 构建命令
        cmd = self.build_digger_command(
            input_dir=work_input_dir,
            sequence_type=sequence_type,
            scaf_suffix=scaf_suffix,
            threads=threads,
            **kwargs,
        )
        
        logger.info(f"执行命令: {' '.join(cmd)}")
        
        # 执行命令
        try:
            proc = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                cwd=str(output_dir),
            )
            
            logs = f"=== STDOUT ===\n{proc.stdout}\n=== STDERR ===\n{proc.stderr}"
            
            # 保存日志
            log_file = log_dir / "digger_execution.log"
            log_file.write_text(logs, encoding="utf-8")
            logger.info(f"日志已保存: {log_file}")
            
            # 构建结果
            result: Dict[str, Any] = {
                "success": proc.returncode == 0,
                "exit_code": proc.returncode,
                "logs": logs,
                "status": "completed" if proc.returncode == 0 else "failed",
            }
            
            if proc.returncode != 0:
                result["error"] = (
                    f"BtToxin_Digger 执行失败 (exit={proc.returncode}): "
                    f"{proc.stderr[:500] if proc.stderr else 'No error output'}"
                )
            
            # 统计输出文件
            results_dir = output_dir / "Results"
            if results_dir.exists():
                files = [f for f in results_dir.rglob("*") if f.is_file()]
                result["output_files"] = len(files)
            else:
                result["output_files"] = 0
            
            return result
            
        except FileNotFoundError as e:
            error_msg = (
                f"命令执行失败: {e}。"
                "请确保 pixi 已正确安装并在 PATH 中。"
                "安装说明: https://pixi.sh"
            )
            logger.error(error_msg)
            return {
                "success": False,
                "exit_code": -1,
                "logs": "",
                "status": "error",
                "error": error_msg,
            }
        except subprocess.TimeoutExpired as e:
            error_msg = f"命令执行超时: {e}"
            logger.error(error_msg)
            return {
                "success": False,
                "exit_code": -1,
                "logs": "",
                "status": "error",
                "error": error_msg,
            }
        except Exception as e:
            error_msg = f"执行 BtToxin_Digger 时出错: {e}"
            logger.error(error_msg, exc_info=True)
            return {
                "success": False,
                "exit_code": -1,
                "logs": "",
                "status": "error",
                "error": error_msg,
            }


    def _setup_external_db(self, bttoxin_db_dir: Path) -> bool:
        """设置外部数据库符号链接。
        
        将 external_dbs/bt_toxin 链接到 pixi 环境的 BTTCMP_db/bt_toxin。
        
        Args:
            bttoxin_db_dir: 外部数据库目录路径
            
        Returns:
            bool: 设置是否成功
        """
        db_path = Path(bttoxin_db_dir).resolve()
        
        # 验证数据库目录结构
        if not db_path.exists():
            logger.warning(f"外部数据库目录不存在: {db_path}")
            return False
        
        if not (db_path / "db").exists():
            logger.warning(f"外部数据库目录结构不完整，缺少 db 子目录: {db_path}")
            return False
        
        # 获取 pixi 环境路径
        try:
            proc = subprocess.run(
                ["pixi", "info", "--json"],
                capture_output=True,
                text=True,
                timeout=30,
                cwd=str(self.pixi_project_dir),
            )
            if proc.returncode != 0:
                logger.warning("无法获取 pixi 环境信息")
                return False
            
            info = json.loads(proc.stdout)
            envs_info = info.get("environments_info", [])
            
            env_prefix = None
            for env in envs_info:
                if env.get("name") == self.env_name:
                    env_prefix = env.get("prefix")
                    break
            
            if not env_prefix:
                logger.warning(f"未找到 {self.env_name} 环境的路径")
                return False
            
            # 创建符号链接
            # BtToxin_Digger 期望数据库在 BTTCMP_db/bt_toxin
            env_path = Path(env_prefix)
            bttcmp_db_dir = env_path / "bin" / "BTTCMP_db"
            target_link = bttcmp_db_dir / "bt_toxin"
            
            # 确保 BTTCMP_db 目录存在
            bttcmp_db_dir.mkdir(parents=True, exist_ok=True)
            
            # 如果已存在符号链接或目录，先删除
            if target_link.exists() or target_link.is_symlink():
                if target_link.is_symlink():
                    target_link.unlink()
                elif target_link.is_dir():
                    shutil.rmtree(target_link)
            
            # 创建符号链接
            target_link.symlink_to(db_path)
            logger.info(f"已创建数据库符号链接: {target_link} -> {db_path}")
            return True
            
        except json.JSONDecodeError as e:
            logger.warning(f"解析 pixi info 输出失败: {e}")
            return False
        except PermissionError as e:
            logger.warning(f"无法创建数据库符号链接（权限不足）: {e}")
            return False
        except Exception as e:
            logger.warning(f"设置外部数据库时出错: {e}")
            return False


def build_shotter_command(
    pixi_project_dir: Path,
    script_path: Path,
    toxicity_csv: Path,
    all_toxins: Path,
    output_dir: Path,
    min_identity: float = 0.0,
    min_coverage: float = 0.0,
    allow_unknown_families: bool = True,
    require_index_hit: bool = False,
) -> List[str]:
    """构建 Shotter 命令列表（使用 pipeline 环境）。
    
    Args:
        pixi_project_dir: pixi.toml 所在目录
        script_path: bttoxin_shoter.py 脚本路径
        toxicity_csv: 毒性数据 CSV 文件路径
        all_toxins: All_Toxins.txt 文件路径
        output_dir: 输出目录
        min_identity: 最小 identity 阈值
        min_coverage: 最小 coverage 阈值
        allow_unknown_families: 是否允许未知家族
        require_index_hit: 是否要求索引命中
        
    Returns:
        命令参数列表，以 'pixi run -e pipeline python' 开头
    """
    cmd: List[str] = [
        "pixi", "run", "-e", "pipeline",
        "python", str(script_path),
        "--toxicity_csv", str(toxicity_csv),
        "--all_toxins", str(all_toxins),
        "--output_dir", str(output_dir),
    ]
    
    if min_identity and min_identity > 0:
        cmd += ["--min_identity", str(min_identity)]
    if min_coverage and min_coverage > 0:
        cmd += ["--min_coverage", str(min_coverage)]
    if not allow_unknown_families:
        cmd += ["--disallow_unknown_families"]
    if require_index_hit:
        cmd += ["--require_index_hit"]
    
    return cmd


def build_plot_command(
    pixi_project_dir: Path,
    script_path: Path,
    strain_scores: Path,
    toxin_support: Path,
    species_scores: Path,
    out_dir: Path,
    merge_unresolved: bool = True,
    report_mode: str = "paper",
    lang: str = "zh",
    per_hit_strain: Optional[str] = None,
) -> List[str]:
    """构建 Plot 命令列表（使用 pipeline 环境）。
    
    Args:
        pixi_project_dir: pixi.toml 所在目录
        script_path: plot_shotter.py 脚本路径
        strain_scores: strain_target_scores.tsv 文件路径
        toxin_support: toxin_support.tsv 文件路径
        species_scores: strain_target_species_scores.tsv 文件路径
        out_dir: 输出目录
        merge_unresolved: 是否合并 unresolved 列
        report_mode: 报告模式 (paper/summary)
        lang: 报告语言 (zh/en)
        per_hit_strain: 可选的菌株名称用于绘制 per-hit 热图
        
    Returns:
        命令参数列表，以 'pixi run -e pipeline python' 开头
    """
    cmd: List[str] = [
        "pixi", "run", "-e", "pipeline",
        "python", str(script_path),
        "--strain_scores", str(strain_scores),
        "--toxin_support", str(toxin_support),
        "--species_scores", str(species_scores),
        "--out_dir", str(out_dir),
        "--report_mode", report_mode,
        "--lang", lang,
    ]
    
    if merge_unresolved:
        cmd.append("--merge_unresolved")
    
    if per_hit_strain:
        cmd += ["--per_hit_strain", per_hit_strain]
    
    return cmd



def create_pipeline_bundle(
    bundle_path: Path,
    digger_dir: Path,
    shotter_dir: Path,
) -> bool:
    """创建 pipeline 结果打包文件。
    
    Args:
        bundle_path: 输出的 tar.gz 文件路径
        digger_dir: digger 输出目录
        shotter_dir: shotter 输出目录
        
    Returns:
        bool: 打包是否成功
    """
    import tarfile
    
    try:
        bundle_path.parent.mkdir(parents=True, exist_ok=True)
        with tarfile.open(bundle_path, "w:gz") as tar:
            if digger_dir.exists():
                tar.add(digger_dir, arcname="digger")
            if shotter_dir.exists():
                tar.add(shotter_dir, arcname="shotter")
        return True
    except Exception as e:
        logger.error(f"创建打包文件失败: {e}")
        return False


def verify_bundle_contents(bundle_path: Path) -> Dict[str, Any]:
    """验证打包文件内容。
    
    Args:
        bundle_path: tar.gz 文件路径
        
    Returns:
        包含验证结果的字典:
        - valid: bool, 打包是否有效
        - has_digger: bool, 是否包含 digger 目录
        - has_shotter: bool, 是否包含 shotter 目录
        - members: List[str], 所有成员名称
    """
    import tarfile
    
    result: Dict[str, Any] = {
        "valid": False,
        "has_digger": False,
        "has_shotter": False,
        "members": [],
    }
    
    try:
        with tarfile.open(bundle_path, "r:gz") as tar:
            members = tar.getnames()
            result["members"] = members
            result["has_digger"] = any(m.startswith("digger") for m in members)
            result["has_shotter"] = any(m.startswith("shotter") for m in members)
            result["valid"] = result["has_digger"] or result["has_shotter"]
    except Exception as e:
        logger.error(f"验证打包文件失败: {e
}")
    
    return result