#!/usr/bin/env python3
"""
RNA-Seq Alignment阶段：将clean reads比对到参考基因组

输入：
  - 来自RNA_QC阶段的clean reads
  - 参考基因组FASTA
  - HISAT2索引

输出：
  - 排序后的BAM文件
  - 比对报告和manifest
"""

import argparse
import json
import subprocess
import shutil
import sys
import time
import yaml
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple
import hashlib
import os

PROJECT_ROOT = Path(__file__).resolve().parents[2]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))
import re

from gfa_core.rnaseq_context import (
    apply_rnaseq_context,
    build_rnaseq_context,
    find_rnaseq_hisat2_index,
    find_rnaseq_reference_file,
    hisat2_rna_strandness_args,
    require_rnaseq_library_strandness,
    RnaseqLibraryStrandnessError,
    resolve_rnaseq_project_context,
)
from cli.rnaseq_env import prepend_path, rnaseq_bin

RNASEQ_BIN = str(rnaseq_bin())
prepend_path(RNASEQ_BIN)


def resolve_thread_plan(cpu: int, sample_parallel_jobs: int = 1, threads_per_sample: Optional[int] = None) -> Dict[str, int]:
    """Resolve conservative RNA-seq sample/thread settings from CLI args."""
    total_cpu = max(1, int(cpu or 1))
    jobs = max(1, int(sample_parallel_jobs or 1))
    if threads_per_sample is None or int(threads_per_sample) < 1:
        threads = max(1, total_cpu // jobs)
    else:
        threads = max(1, int(threads_per_sample))
    if jobs * threads > total_cpu:
        threads = max(1, total_cpu // jobs)
    return {"cpu": total_cpu, "sample_parallel_jobs": jobs, "threads_per_sample": threads}


def species_token_from_prefix(prefix: str) -> str:
    """Extract species token from prefixes like Os, Os.WRKY, or Os_WRKY."""
    prefix = (prefix or "").strip()
    if not prefix:
        return ""
    return re.split(r"[._]", prefix, maxsplit=1)[0].strip()


def find_canonical_genome_fasta(ref_dir: Path, species_token: str) -> Optional[Path]:
    """Find the canonical genome FASTA: {species}.genome.fasta/fa/fna."""
    if not species_token:
        return None
    for suffix in ("fasta", "fa", "fna"):
        candidate = ref_dir / f"{species_token}.genome.{suffix}"
        if candidate.is_file():
            return candidate
    return None


def canonical_genome_index_prefix(ref_dir: Path, species_token: str) -> Optional[Path]:
    """Return exact {species}.genome HISAT2 index prefix if it exists."""
    if not species_token:
        return None
    prefix = ref_dir / f"{species_token}.genome"
    if Path(f"{prefix}.1.ht2").is_file():
        return prefix
    return None


def build_hisat2_index(genome_fasta: Path, index_prefix: Path, cpu: int = 4) -> bool:
    """Build HISAT2 index from the locked genome FASTA."""
    index_prefix.parent.mkdir(parents=True, exist_ok=True)
    cmd = [
        f"{RNASEQ_BIN}/hisat2-build",
        str(genome_fasta),
        str(index_prefix),
        "-p", str(cpu),
    ]
    try:
        subprocess.run(cmd, check=True, capture_output=True, text=True)
        return Path(f"{index_prefix}.1.ht2").exists()
    except subprocess.CalledProcessError as e:
        print(f"hisat2-build failed: {e.stderr}", file=sys.stderr)
        return False


def resolve_locked_hisat2_index(
    requested_index: Path,
    input_dir: Path,
    align_dir: Path,
    prefix: str,
    sample_config: Optional[str],
    cpu: int = 4,
) -> Path:
    """Lock RNA-seq alignment reference to {species}.genome.*."""
    sample_dir = Path(sample_config).parent if sample_config else None
    context = resolve_rnaseq_project_context(
        input_dir=sample_dir or requested_index.parent,
        final_dir=input_dir,
        requested_prefix=prefix,
        sample_config=sample_config,
    )
    species_token = context.get("species_code") or species_token_from_prefix(context.get("output_prefix") or prefix)
    reference_id = context.get("reference_id", "")
    if not species_token or species_token.lower() in {"unknown", "genome", "genome_prepare", "rnaseq"}:
        raise ValueError(
            "RNA-seq Align requires a real species code from samples.yaml/config.xlsx, "
            "for example species.code: Os and rnaseq.outputPrefix: Os"
        )

    ref_dirs = []
    if requested_index:
        ref_dirs.append(requested_index.parent)
    if sample_dir:
        ref_dirs.append(sample_dir)
    ref_dirs.append(input_dir)

    seen = set()
    unique_ref_dirs = []
    for ref_dir in ref_dirs:
        resolved = ref_dir.resolve()
        if resolved not in seen:
            unique_ref_dirs.append(ref_dir)
            seen.add(resolved)

    for ref_dir in unique_ref_dirs:
        exact_index = find_rnaseq_hisat2_index(ref_dir, species_token, reference_id)
        if exact_index:
            print(f"[INFO] Locked RNA-seq genome index: {exact_index}")
            return exact_index

    for ref_dir in unique_ref_dirs:
        genome_fasta = find_rnaseq_reference_file(ref_dir, "genome", species_token, reference_id)
        if genome_fasta is None:
            continue
        build_prefix = align_dir / "index" / f"{species_token}.genome"
        exact_built = canonical_genome_index_prefix(build_prefix.parent, species_token)
        if exact_built:
            print(f"[INFO] Locked RNA-seq genome index: {exact_built}")
            return exact_built
        print(f"[INFO] Building locked HISAT2 genome index from {genome_fasta}: {build_prefix}")
        if not build_hisat2_index(genome_fasta, build_prefix, cpu):
            raise RuntimeError(f"Failed to build HISAT2 index from {genome_fasta}")
        return build_prefix

    searched = ", ".join(str(p) for p in unique_ref_dirs)
    raise FileNotFoundError(
        f"RNA-seq Align locked reference lookup failed: resolved species={species_token!r}, "
        f"reference_id={reference_id!r}; expected {species_token}.genome.fasta "
        f"(or .fa/.fna) or {species_token}.genome.1.ht2 under: {searched}. "
        f"Protein/CDS indexes such as {species_token}.pep.*.ht2 or *.cds.*.ht2 are not valid for RNA-seq alignment."
    )

def compute_file_hash(filepath: Path) -> str:
    """计算文件的SHA256哈希值"""
    sha256_hash = hashlib.sha256()
    try:
        with open(filepath, "rb") as f:
            for byte_block in iter(lambda: f.read(4096), b""):
                sha256_hash.update(byte_block)
        return sha256_hash.hexdigest()
    except Exception:
        return ""


def parse_annotation_attributes(attr_text: str) -> Dict[str, str]:
    attrs: Dict[str, str] = {}
    for item in attr_text.split(";"):
        item = item.strip()
        if not item:
            continue
        if "=" in item:
            key, value = item.split("=", 1)
            attrs[key.strip()] = value.strip().strip('"')
        elif " " in item:
            key, value = item.split(" ", 1)
            attrs[key.strip()] = value.strip().strip('"')
    return attrs


def annotation_transcript_id(feature_type: str, attr_text: str) -> Optional[str]:
    attrs = parse_annotation_attributes(attr_text)
    feature = feature_type.lower()
    if feature in {"mrna", "transcript"}:
        return attrs.get("ID") or attrs.get("transcript_id") or attrs.get("Name")
    if feature == "exon":
        parent = attrs.get("Parent") or attrs.get("transcript_id")
        if parent:
            return parent.split(",")[0]
        return attrs.get("ID")
    return None


def extract_hisat2_splice_sites(annotation: Path, output_file: Path) -> Dict[str, Any]:
    """Extract HISAT2 known splice sites from exon records in GTF/GFF3."""
    transcripts: Dict[str, List[Tuple[str, int, int, str]]] = {}
    with open(annotation, "r", encoding="utf-8", errors="ignore") as handle:
        for line in handle:
            if not line.strip() or line.startswith("#"):
                continue
            parts = line.rstrip("\n").split("\t")
            if len(parts) < 9 or parts[2].lower() != "exon":
                continue
            transcript_id = annotation_transcript_id(parts[2], parts[8])
            if not transcript_id:
                continue
            try:
                start = int(parts[3])
                end = int(parts[4])
            except ValueError:
                continue
            transcripts.setdefault(transcript_id, []).append((parts[0], start, end, parts[6] or "."))

    splice_sites = set()
    for exons in transcripts.values():
        exons = sorted(exons, key=lambda row: (row[0], row[1], row[2]))
        for left, right in zip(exons, exons[1:]):
            if left[0] != right[0]:
                continue
            intron_start = left[2] + 1
            intron_end = right[1] - 1
            if intron_start <= intron_end:
                splice_sites.add((left[0], intron_start, intron_end, left[3]))

    output_file.parent.mkdir(parents=True, exist_ok=True)
    with open(output_file, "w", encoding="utf-8") as out:
        for chrom, start, end, strand in sorted(splice_sites):
            out.write(f"{chrom}\t{start}\t{end}\t{strand}\n")

    return {
        "status": "generated" if splice_sites else "empty",
        "source": str(annotation),
        "path": str(output_file),
        "count": len(splice_sites),
    }


def resolve_known_splice_site(
    source_mode: str,
    known_splice_site: Optional[str],
    final_dir: Path,
    prefix: str,
) -> Tuple[Optional[Path], Dict[str, Any]]:
    mode = (source_mode or "auto").strip().lower()
    if mode not in {"auto", "off", "manual", "annotation"}:
        mode = "auto"
    info: Dict[str, Any] = {
        "mode": mode,
        "status": "disabled" if mode == "off" else "not_found",
        "source": "",
        "path": "",
        "count": 0,
    }
    if mode == "off":
        return None, info

    source_path = Path(known_splice_site).expanduser() if known_splice_site else None
    if not source_path or not source_path.exists():
        if mode == "manual":
            raise FileNotFoundError(f"--known-splice-site does not exist: {known_splice_site}")
        return None, info

    suffixes = "".join(source_path.suffixes).lower()
    if source_path.suffix.lower() in {".ss", ".txt", ".tsv"} and mode != "annotation":
        info.update({
            "status": "provided",
            "source": str(source_path),
            "path": str(source_path),
            "count": sum(1 for line in source_path.read_text(errors="ignore").splitlines() if line.strip() and not line.startswith("#")),
        })
        return source_path, info

    if source_path.suffix.lower() in {".gtf", ".gff", ".gff3"} or any(token in suffixes for token in (".gtf", ".gff")):
        out = final_dir / f"{prefix}.hisat2.splice_sites.ss"
        generated = extract_hisat2_splice_sites(source_path, out)
        info.update(generated)
        info["mode"] = mode
        return (out if out.exists() and out.stat().st_size > 0 else None), info

    if mode in {"manual", "annotation"}:
        raise ValueError(f"Unsupported known splice site source format: {source_path}")
    return None, info


def get_tool_versions() -> Dict[str, str]:
    """获取所需工具版本信息"""
    versions = {}
    
    # hisat2
    try:
        result = subprocess.run(
            [f"{RNASEQ_BIN}/hisat2", "--version"],
            capture_output=True,
            text=True,
            check=False,
            timeout=5
        )
        output = result.stdout.strip() or result.stderr.strip()
        versions["hisat2"] = output.split('\n')[0] if output else "unknown"
    except Exception as e:
        versions["hisat2"] = f"not_found: {e}"
    
    # samtools
    try:
        result = subprocess.run(
            [f"{RNASEQ_BIN}/samtools", "--version"],
            capture_output=True,
            text=True,
            check=False,
            timeout=5
        )
        output = result.stdout.strip() or result.stderr.strip()
        versions["samtools"] = output.split('\n')[0] if output else "unknown"
    except Exception as e:
        versions["samtools"] = f"not_found: {e}"

    # MultiQC
    try:
        result = subprocess.run(
            [f"{RNASEQ_BIN}/multiqc", "--version"],
            capture_output=True,
            text=True,
            check=False,
            timeout=10
        )
        output = result.stdout.strip() or result.stderr.strip()
        versions["multiqc"] = output.split('\n')[0] if output else "unknown"
    except Exception as e:
        versions["multiqc"] = f"not_found: {e}"
    
    return versions


def run_hisat2_sort(
    hisat2_idx: Path,
    r1: Path,
    r2: Path,
    sorted_bam: Path,
    log_out: Path,
    cpu: int = 4,
    library_strandness: str = "non_stranded",
    known_splice_site: Optional[Path] = None,
) -> bool:
    """Run HISAT2 and stream SAM directly into samtools sort."""
    idx_prefix = str(hisat2_idx).rsplit('.', 1)[0] if hisat2_idx.suffix == '.ht2' else str(hisat2_idx)
    hisat2_cmd = [
        f"{RNASEQ_BIN}/hisat2",
        "-x", idx_prefix,
        "-1", str(r1),
        "-2", str(r2),
        "-p", str(cpu),
        "--new-summary",
    ]
    hisat2_cmd.extend(hisat2_rna_strandness_args(library_strandness))
    if known_splice_site:
        hisat2_cmd.extend(["--known-splicesite-infile", str(known_splice_site)])
    sort_cmd = [
        f"{RNASEQ_BIN}/samtools",
        "sort",
        "-@", str(cpu),
        "-o", str(sorted_bam),
        "-",
    ]

    sorted_bam.parent.mkdir(parents=True, exist_ok=True)
    with open(log_out, "w") as log:
        log.write("Command: " + " ".join(hisat2_cmd) + " | " + " ".join(sort_cmd) + "\n\n")
        hisat2_proc = subprocess.Popen(hisat2_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        assert hisat2_proc.stdout is not None
        sort_proc = subprocess.Popen(sort_cmd, stdin=hisat2_proc.stdout, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        hisat2_proc.stdout.close()
        sort_stdout, sort_stderr = sort_proc.communicate()
        hisat2_stderr = hisat2_proc.stderr.read() if hisat2_proc.stderr else b""
        hisat2_return = hisat2_proc.wait()

        log.write("--- HISAT2 STDERR ---\n")
        log.write(hisat2_stderr.decode("utf-8", errors="replace"))
        log.write("\n--- SAMTOOLS SORT STDOUT ---\n")
        log.write((sort_stdout or b"").decode("utf-8", errors="replace"))
        log.write("\n--- SAMTOOLS SORT STDERR ---\n")
        log.write((sort_stderr or b"").decode("utf-8", errors="replace"))

    return hisat2_return == 0 and sort_proc.returncode == 0 and sorted_bam.exists() and sorted_bam.stat().st_size > 0


def index_bam(bam: Path) -> bool:
    """索引BAM"""
    cmd = [f"{RNASEQ_BIN}/samtools", "index", str(bam)]
    try:
        subprocess.run(cmd, check=True, capture_output=True)
        return True
    except subprocess.CalledProcessError as e:
        print(f"samtools index failed: {e.stderr}", file=sys.stderr)
        return False


def run_samtools_qc(bam: Path, flagstat_out: Path, stats_out: Path, cpu: int = 4) -> bool:
    """Generate samtools flagstat and stats files for an aligned BAM."""
    ok = True
    try:
        with open(flagstat_out, "w") as out:
            subprocess.run(
                [f"{RNASEQ_BIN}/samtools", "flagstat", "-@", str(cpu), str(bam)],
                check=True,
                stdout=out,
                stderr=subprocess.PIPE,
                text=True,
            )
    except subprocess.CalledProcessError as e:
        print(f"samtools flagstat failed for {bam.name}: {e.stderr}", file=sys.stderr)
        ok = False
    try:
        with open(stats_out, "w") as out:
            subprocess.run(
                [f"{RNASEQ_BIN}/samtools", "stats", "-@", str(cpu), str(bam)],
                check=True,
                stdout=out,
                stderr=subprocess.PIPE,
                text=True,
            )
    except subprocess.CalledProcessError as e:
        print(f"samtools stats failed for {bam.name}: {e.stderr}", file=sys.stderr)
        ok = False
    return ok


def parse_samtools_flagstat(flagstat_file: Path) -> Dict[str, Any]:
    """Parse a small subset of samtools flagstat metrics."""
    stats: Dict[str, Any] = {}
    if not flagstat_file.exists():
        return stats
    content = flagstat_file.read_text(errors="replace")
    patterns = {
        "flagstat_total": r"^(\d+) \+ \d+ in total",
        "flagstat_mapped": r"^(\d+) \+ \d+ mapped ",
        "flagstat_properly_paired": r"^(\d+) \+ \d+ properly paired ",
        "flagstat_duplicates": r"^(\d+) \+ \d+ duplicates",
    }
    for key, pattern in patterns.items():
        match = re.search(pattern, content, re.MULTILINE)
        if match:
            stats[key] = int(match.group(1))
    match = re.search(r"\((\d+(?:\.\d+)?)% : N/A\)", content)
    if match:
        stats["flagstat_mapped_rate"] = float(match.group(1))
    return stats


def run_multiqc(search_paths: Sequence[Path], output_dir: Path, report_name: str) -> Optional[Path]:
    """Run MultiQC over selected stage outputs and return the report path."""
    multiqc = Path(f"{RNASEQ_BIN}/multiqc")
    existing = [str(path) for path in search_paths if path.exists()]
    if not multiqc.exists() or not existing:
        return None
    cmd = [str(multiqc), *existing, "-o", str(output_dir), "-n", report_name, "-f"]
    try:
        subprocess.run(cmd, check=True, capture_output=True, text=True)
        report = output_dir / report_name
        return report if report.exists() else None
    except subprocess.CalledProcessError as e:
        print(f"MultiQC failed: {e.stderr}", file=sys.stderr)
        return None

def parse_hisat2_log(log_file: Path) -> Dict[str, Any]:
    """解析HISAT2日志获取统计信息"""
    stats = {
        "total_reads": 0,
        "aligned_reads": 0,
        "unaligned_reads": 0,
        "alignment_rate": 0.0,
        "unique_alignments": 0,
        "multiple_alignments": 0
    }
    
    if not log_file.exists():
        return stats
    
    with open(log_file, 'r') as f:
        content = f.read()
    
    # 解析总reads数
    match = re.search(r'(\d+) reads; of these:', content)
    if match:
        stats["total_reads"] = int(match.group(1))
    else:
        match = re.search(r'Total pairs:\s+(\d+)', content)
        if match:
            stats["total_reads"] = int(match.group(1))
    
    # 解析比对率
    match = re.search(r'([\d.]+)% overall alignment rate', content)
    if match:
        stats["alignment_rate"] = float(match.group(1))
    else:
        match = re.search(r'Overall alignment rate:\s+([\d.]+)%', content)
        if match:
            stats["alignment_rate"] = float(match.group(1))
    
    # 解析唯一比对和多比对
    match = re.search(r'(\d+) \(.*?\) aligned concordantly exactly 1 time', content)
    if match:
        stats["unique_alignments"] = int(match.group(1))
    else:
        match = re.search(r'Aligned concordantly 1 time:\s+(\d+)', content)
        if match:
            stats["unique_alignments"] = int(match.group(1))
    
    match = re.search(r'(\d+) \(.*?\) aligned concordantly >1 times', content)
    if match:
        stats["multiple_alignments"] = int(match.group(1))
    else:
        match = re.search(r'Aligned concordantly >1 times:\s+(\d+)', content)
        if match:
            stats["multiple_alignments"] = int(match.group(1))
        match = re.search(r'Aligned discordantly 1 time:\s+(\d+)', content)
        if match:
            stats["multiple_alignments"] += int(match.group(1))
    
    # 计算未比对
    if stats["total_reads"] > 0:
        aligned = stats["unique_alignments"] + stats["multiple_alignments"]
        if "Total pairs:" in content:
            stats["aligned_reads"] = aligned
            stats["unaligned_reads"] = max(stats["total_reads"] - aligned, 0)
        else:
            stats["aligned_reads"] = aligned * 2  # 双端reads
            stats["unaligned_reads"] = stats["total_reads"] * 2 - stats["aligned_reads"]
    
    return stats


def resolve_clean_fastq(input_dir: Path, sample_name: str, original_path: Path, mate: int) -> Optional[Path]:
    """Resolve clean FASTQ generated by rnaseq_qc, including prefixed final/ names."""
    candidates = []

    original_stem = original_path.name
    for suffix in (".fastq.gz", ".fq.gz", ".fastq", ".fq", ".gz"):
        if original_stem.endswith(suffix):
            original_stem = original_stem[: -len(suffix)]
            break
    candidates.extend(
        [
            input_dir / f"{original_stem}.clean.fastq.gz",
            input_dir / f"{sample_name}_{mate}.clean.fastq.gz",
            input_dir / f"{sample_name}_R{mate}.clean.fastq.gz",
        ]
    )
    candidates.extend(sorted(input_dir.glob(f"*.{sample_name}_{mate}.clean.fastq.gz")))
    candidates.extend(sorted(input_dir.glob(f"*{sample_name}*_{mate}.clean.fastq.gz")))
    candidates.extend(sorted(input_dir.glob(f"*{sample_name}*_R{mate}*.clean.fastq.gz")))

    if original_path.exists():
        candidates.append(original_path)

    for candidate in candidates:
        if candidate.exists() and candidate.stat().st_size > 0:
            return candidate
    return None


def generate_report(
    output_dir: Path,
    sample_stats: Dict[str, Dict[str, Any]],
    versions: Dict[str, str],
    start_time: datetime,
    end_time: datetime,
    prefix: str = "",
    cpu: int = 4,
    sample_parallel_jobs: int = 1,
    threads_per_sample: int = 4,
    sample_config: Optional[str] = None,
) -> Path:
    """生成比对报告"""
    report_name = f"{prefix}.rnaseq_align_report.md" if prefix else "rnaseq_align_report.md"
    report_file = output_dir / report_name
    
    elapsed = (end_time - start_time).total_seconds()
    
    content = f"""# RNA-Seq Alignment Report

## Summary

- **Analysis Date**: {end_time.strftime("%Y-%m-%d %H:%M:%S")}
- **Total Samples**: {len(sample_stats)}
- **Total Runtime**: {elapsed:.1f} seconds

## Tool Versions

"""
    for tool, version in versions.items():
        content += f"- **{tool}**: {version}\n"
    
    content += "\n## Sample Alignment Statistics\n\n"
    content += "| Sample | Total Reads | Aligned | Unaligned | Alignment Rate |\n"
    content += "|--------|-------------|---------|-----------|----------------|\n"
    
    for sample, stats in sample_stats.items():
        content += f"| {sample} | {stats.get('total_reads', 'N/A')} | {stats.get('aligned_reads', 'N/A')} | {stats.get('unaligned_reads', 'N/A')} | {stats.get('alignment_rate', 'N/A')}% |\n"
    
    with open(report_file, 'w') as f:
        f.write(content)
    
    return report_file


def generate_manifest(
    output_dir: Path,
    sample_stats: Dict[str, Dict[str, Any]],
    versions: Dict[str, str],
    start_time: datetime,
    end_time: datetime,
    cmd: List[str],
    input_dir: Path,
    workspace_dir: Path,
    job_id: str,
    run_id: str,
    prefix: str = "",
    cpu: int = 4,
    sample_parallel_jobs: int = 1,
    threads_per_sample: int = 4,
    sample_config: Optional[str] = None,
    library_strandness: Optional[Dict[str, Any]] = None,
    splice_site_info: Optional[Dict[str, Any]] = None,
) -> Path:
    """生成manifest文件"""
    manifest_name = f"{prefix}.rnaseq_align.manifest.json" if prefix else "rnaseq_align.manifest.json"
    manifest_file = output_dir / manifest_name
    
    total_aligned = sum(s.get('aligned_reads', 0) for s in sample_stats.values())
    total_reads = sum(s.get('total_reads', 0) for s in sample_stats.values())
    
    avg_rate = (total_aligned / total_reads * 100) if total_reads > 0 else 0
    
    splice_info = splice_site_info or {}
    splice_param = ""
    if splice_info.get("path"):
        splice_param = f" --known-splicesite-infile {splice_info.get('path')}"
    manifest = {
        "manifest_version": "1.0",
        "stage": "rnaseq_align",
        "status": "success",
        "job_id": job_id,
        "run_id": run_id,
        "timestamp": end_time.isoformat(),
        "start_time": start_time.isoformat(),
        "end_time": end_time.isoformat(),
        "elapsed_sec": (end_time - start_time).total_seconds(),
        "cmd": cmd,
        "workspace_root": str(output_dir.parent.parent),
        "workspace": str(output_dir),
        "input_dir": str(input_dir),
        "metrics": {
            "total_samples": len(sample_stats),
            "total_reads": total_reads,
            "aligned_reads": total_aligned,
            "average_alignment_rate": avg_rate,
            "samples": sample_stats
        },
        "parameters": {
            "cpu": cpu,
            "sample_parallel_jobs": sample_parallel_jobs,
            "threads_per_sample": threads_per_sample,
            "hisat2_params": f"-p {threads_per_sample} --new-summary{splice_param} | samtools sort",
            "samtools_params": f"sort/index/flagstat/stats -@ {threads_per_sample}",
            "splice_site_source": splice_info.get("mode", "auto"),
            "known_splice_site_status": splice_info.get("status", "not_found"),
            "known_splice_site": splice_info.get("path", ""),
            "known_splice_site_input": splice_info.get("source", ""),
            "known_splice_site_count": splice_info.get("count", 0),
            "library_strandness_by_sample": (library_strandness or {}).get("library_strandness_by_sample", {}),
            "library_strandness_values": (library_strandness or {}).get("library_strandness_values", []),
            "library_strandness_signature": (library_strandness or {}).get("library_strandness_signature", ""),
        },
        "artifacts": {},
        "tools": versions
    }
    context = build_rnaseq_context(
        prefix=prefix,
        sample_config=sample_config,
        foreground_mode="global",
        stage="rnaseq_align",
    )
    apply_rnaseq_context(manifest, context)
    if library_strandness:
        manifest["rnaseq_library_strandness_by_sample"] = library_strandness.get("library_strandness_by_sample", {})
        manifest["rnaseq_library_strandness_values"] = library_strandness.get("library_strandness_values", [])
        manifest["rnaseq_library_strandness_signature"] = library_strandness.get("library_strandness_signature", "")
        manifest["rnaseq_library_strandness_mixed"] = library_strandness.get("library_strandness_mixed", False)
    if splice_info:
        manifest["rnaseq_known_splice_site"] = splice_info
    
    # 添加artifacts（参考rnaseq_qc契约，包含完整字段）
    def add_artifact(name: str, filepath: Path, artifact_type: str, **extra):
        """添加 artifact 到 manifest"""
        if filepath.exists():
            manifest["artifacts"][name] = {
                "artifact": name,
                "path": f"final/{filepath.name}",
                "full_path": str(filepath.absolute()),
                "type": artifact_type,
                "size": filepath.stat().st_size,
                "sha256": compute_file_hash(filepath),
                **extra
            }
    
    for sample in sample_stats.keys():
        # 排序后的BAM文件 - 在output_dir (即final_dir)
        sorted_bam = output_dir / f"{sample}.sorted.bam"
        add_artifact(f"sorted_bam_{sample}", sorted_bam, "bam", format="bam")
        
        # BAM索引 - 在output_dir
        bam_index = output_dir / f"{sample}.sorted.bam.bai"
        add_artifact(f"bam_index_{sample}", bam_index, "bai", format="bai")
        
        # HISAT2日志和 samtools QC 文件 - 在output_dir
        log_file = output_dir / f"{sample}_hisat2.log"
        add_artifact(f"alignment_log_{sample}", log_file, "log", format="txt")
        flagstat_file = output_dir / f"{sample}.flagstat.txt"
        stats_file = output_dir / f"{sample}.samtools.stats.txt"
        add_artifact(f"samtools_flagstat_{sample}", flagstat_file, "text", format="txt")
        add_artifact(f"samtools_stats_{sample}", stats_file, "text", format="txt")
    
    # 报告、MultiQC和manifest - 在output_dir
    report_name = f"{prefix}.rnaseq_align_report.md" if prefix else "rnaseq_align_report.md"
    report_file = output_dir / report_name
    add_artifact("report", report_file, "markdown", format="md")
    generated_splice = output_dir / f"{prefix}.hisat2.splice_sites.ss"
    add_artifact("known_splice_site", generated_splice, "tsv", format="hisat2_splice_site")
    multiqc_name = f"{prefix}.rnaseq_align_multiqc_report.html" if prefix else "rnaseq_align_multiqc_report.html"
    multiqc_file = output_dir / multiqc_name
    add_artifact("multiqc_report", multiqc_file, "html", format="html")
    multiqc_data = output_dir / multiqc_name.replace(".html", "_data")
    if multiqc_data.exists():
        manifest["artifacts"]["multiqc_data"] = {
            "artifact": "multiqc_data",
            "path": f"final/{multiqc_data.name}",
            "full_path": str(multiqc_data.absolute()),
            "type": "directory",
            "format": "multiqc_data"
        }
    
    # manifest自身
    manifest["artifacts"]["manifest"] = {
        "artifact": "manifest",
        "path": f"final/{manifest_name}",
        "type": "json",
        "format": "json"
    }
    
    with open(manifest_file, 'w') as f:
        json.dump(manifest, f, indent=2)
    
    return manifest_file


def main():
    parser = argparse.ArgumentParser(description="RNA-Seq Alignment")
    parser.add_argument("-i", "--input", required=True, help="Input directory with clean FASTQ files")
    parser.add_argument("-x", "--hisat2-index", required=True, help="HISAT2 index prefix")
    parser.add_argument("-o", "--output", required=True, help="Output directory")
    parser.add_argument("-j", "--job-id", default="unknown", help="Job ID")
    parser.add_argument("-r", "--run-id", default="unknown", help="Run ID")
    parser.add_argument("-t", "--cpu", type=int, default=4, help="Total CPU threads for this stage")
    parser.add_argument("--sample-parallel-jobs", type=int, default=1, help="Number of samples to process concurrently")
    parser.add_argument("--threads-per-sample", type=int, help="Threads used by each per-sample tool invocation")
    parser.add_argument("-p", "--prefix", default="", help="Output file prefix")
    parser.add_argument("--sample-config", help="Sample config YAML for sample list")
    parser.add_argument(
        "--splice-site-source",
        default="auto",
        choices=["auto", "off", "manual", "annotation"],
        help="HISAT2 known splice site source. auto/manual use --known-splice-site when provided; annotation extracts .ss from GTF/GFF.",
    )
    parser.add_argument(
        "--known-splice-site",
        help="HISAT2 splice site file (.ss/.txt) or annotation GTF/GFF used to generate a splice-site file.",
    )
    
    args = parser.parse_args()
    thread_plan = resolve_thread_plan(args.cpu, args.sample_parallel_jobs, args.threads_per_sample)
    threads_per_sample = thread_plan["threads_per_sample"]
    
    start_time = datetime.now()
    input_dir = Path(args.input)
    output_dir = Path(args.output)
    hisat2_idx = Path(args.hisat2_index)
    context_for_prefix = resolve_rnaseq_project_context(
        input_dir=Path(args.sample_config).parent if args.sample_config else hisat2_idx.parent,
        final_dir=input_dir,
        requested_prefix=args.prefix,
        sample_config=args.sample_config,
    )
    resolved_prefix = context_for_prefix.get("output_prefix")
    if resolved_prefix and resolved_prefix != args.prefix:
        print(
            f"[INFO] Using RNA-seq output prefix {resolved_prefix} "
            f"instead of requested {args.prefix or '<empty>'} "
            f"(source={context_for_prefix.get('source', '')})"
        )
        args.prefix = resolved_prefix
    
    # 创建工作目录
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # 最终结果目录: workspace/<job_id>/final/
    # 根据契约，最终产物必须位于 workspace/<job_id>/final/
    workspace_dir = output_dir.parent  # job_id 目录
    final_dir = workspace_dir / "final"
    final_dir.mkdir(exist_ok=True)
    
    # 中间文件目录: rnaseq_align/analysis/
    align_dir = output_dir / "analysis"
    align_dir.mkdir(exist_ok=True)
    
    # 确定样本列表
    samples = {}
    if args.sample_config:
        import yaml
        with open(args.sample_config, 'r') as f:
            config = yaml.safe_load(f)
        samples = config.get("samples", {})
        if not samples:
            raise ValueError(f"Config file {args.sample_config} must contain 'samples' field")
    else:
        # 从输入目录自动检测样本
        for f in sorted(input_dir.glob("*_1.clean.fastq.gz")):
            sample_name = f.stem.replace("_1.clean", "")
            r1 = f
            r2 = input_dir / f"{sample_name}_2.clean.fastq.gz"
            if r2.exists():
                samples[sample_name] = (str(r1), str(r2))
    
    print(f"Processing {len(samples)} samples...")
    sample_names_for_strandness = list(map(str, samples.keys()))
    try:
        strandness_info = require_rnaseq_library_strandness(args.sample_config, sample_names_for_strandness)
    except RnaseqLibraryStrandnessError as exc:
        print(f"Error: {exc}", file=sys.stderr)
        sys.exit(1)
    strandness_by_sample = strandness_info.get("library_strandness_by_sample", {})
    print(f"Library strandness: {strandness_by_sample}")
    
    # 获取工具版本
    versions = get_tool_versions()
    print(f"Tool versions: {versions}")
    
    try:
        hisat2_idx = resolve_locked_hisat2_index(
            requested_index=hisat2_idx,
            input_dir=input_dir,
            align_dir=align_dir,
            prefix=args.prefix,
            sample_config=args.sample_config,
            cpu=threads_per_sample,
        )
        print(f"Using locked HISAT2 genome index prefix: {hisat2_idx}")
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)

    try:
        known_splice_site, splice_site_info = resolve_known_splice_site(
            source_mode=args.splice_site_source,
            known_splice_site=args.known_splice_site,
            final_dir=final_dir,
            prefix=args.prefix,
        )
    except Exception as e:
        print(f"Error: failed to prepare HISAT2 known splice sites: {e}", file=sys.stderr)
        sys.exit(1)
    if known_splice_site:
        print(
            f"Using HISAT2 known splice sites: {known_splice_site} "
            f"(source={splice_site_info.get('source')}, count={splice_site_info.get('count')})"
        )
    else:
        print(f"HISAT2 known splice sites: {splice_site_info.get('status', 'not_found')}")
    
    sample_stats = {}

    def align_one_sample(sample_name: str, paths: Any) -> tuple[str, Optional[Dict[str, Any]], Optional[str]]:
        # 支持两种格式：列表 [r1, r2] 或字典 {"r1": ..., "r2": ...}
        if isinstance(paths, list):
            if len(paths) < 2:
                return sample_name, None, f"Sample {sample_name} has insufficient paths"
            r1_path, r2_path = paths[0], paths[1]
        elif isinstance(paths, dict):
            r1_path = paths.get("r1") or paths.get("read1")
            r2_path = paths.get("r2") or paths.get("read2")
            if not r1_path or not r2_path:
                return sample_name, None, f"Sample {sample_name} missing r1/r2 paths"
        else:
            return sample_name, None, f"Unknown format for sample {sample_name}"

        r1 = Path(r1_path)
        r2 = Path(r2_path)

        # 如果路径是相对路径，从input_dir查找
        if not r1.is_absolute():
            r1 = input_dir / r1.name
        if not r2.is_absolute():
            r2 = input_dir / r2.name

        r1 = resolve_clean_fastq(input_dir, sample_name, r1, 1) or r1
        r2 = resolve_clean_fastq(input_dir, sample_name, r2, 2) or r2

        if not r1.exists() or not r2.exists():
            return sample_name, None, f"FASTQ files not found for sample {sample_name}: R1={r1}, R2={r2}"

        print(f"Aligning sample: {sample_name}")
        print(f"  R1: {r1}")
        print(f"  R2: {r2}")
        sample_strandness = strandness_by_sample.get(str(sample_name), "")
        print(f"  library_strandness: {sample_strandness}")

        # HISAT2 比对直接 pipe 到 samtools sort，减少 SAM/BAM 中间落盘
        log_file = align_dir / f"{sample_name}_hisat2.log"
        sorted_bam = final_dir / f"{sample_name}.sorted.bam"
        if not run_hisat2_sort(
            hisat2_idx,
            r1,
            r2,
            sorted_bam,
            log_file,
            threads_per_sample,
            sample_strandness,
            known_splice_site=known_splice_site,
        ):
            return sample_name, None, f"HISAT2/samtools sort failed for sample {sample_name}"

        # 索引BAM、生成 samtools QC 文件
        index_bam(sorted_bam)
        flagstat_file = final_dir / f"{sample_name}.flagstat.txt"
        samtools_stats_file = final_dir / f"{sample_name}.samtools.stats.txt"
        run_samtools_qc(sorted_bam, flagstat_file, samtools_stats_file, threads_per_sample)

        # 复制日志到final_dir
        final_log = final_dir / f"{sample_name}_hisat2.log"
        shutil.copy2(str(log_file), str(final_log))

        # 解析日志和 samtools flagstat 获取统计
        stats = parse_hisat2_log(log_file)
        stats.update(parse_samtools_flagstat(flagstat_file))
        return sample_name, stats, None

    print(
        f"Thread plan: sample_parallel_jobs={thread_plan['sample_parallel_jobs']}, "
        f"threads_per_sample={threads_per_sample}, total_cpu={thread_plan['cpu']}"
    )
    max_workers = max(1, int(thread_plan["sample_parallel_jobs"]))
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_map = {
            executor.submit(align_one_sample, sample_name, paths): sample_name
            for sample_name, paths in samples.items()
        }
        for future in as_completed(future_map):
            sample_name = future_map[future]
            try:
                completed_name, stats, error = future.result()
            except Exception as exc:
                print(f"Warning: sample {sample_name} failed with exception: {exc}", file=sys.stderr)
                continue
            if error:
                print(f"Warning: {error}", file=sys.stderr)
                continue
            if stats is not None:
                sample_stats[completed_name] = stats

    # 生成 MultiQC 报告
    multiqc_report = run_multiqc([align_dir, final_dir], final_dir, f"{args.prefix}.rnaseq_align_multiqc_report.html" if args.prefix else "rnaseq_align_multiqc_report.html")
    if multiqc_report:
        print(f"MultiQC report: {multiqc_report}")

    # 生成报告
    if not sample_stats:
        print("Error: no samples were aligned. Check clean FASTQ names and samples.yaml.", file=sys.stderr)
        sys.exit(1)

    end_time = datetime.now()
    generate_report(final_dir, sample_stats, versions, start_time, end_time, args.prefix)
    
    # 生成manifest
    cmd = ["python", __file__, "-i", str(input_dir), "-x", str(hisat2_idx),
           "-o", str(output_dir), "-j", args.job_id, "-r", args.run_id, "-p", args.prefix,
           "-t", str(thread_plan["cpu"]),
           "--sample-parallel-jobs", str(thread_plan["sample_parallel_jobs"]),
           "--threads-per-sample", str(thread_plan["threads_per_sample"])]
    if args.sample_config:
        cmd.extend(["--sample-config", str(args.sample_config)])
    cmd.extend(["--splice-site-source", args.splice_site_source])
    if args.known_splice_site:
        cmd.extend(["--known-splice-site", str(args.known_splice_site)])
    generate_manifest(
        final_dir, sample_stats, versions, start_time, end_time,
        cmd, input_dir, output_dir, args.job_id, args.run_id, args.prefix,
        cpu=args.cpu,
        sample_parallel_jobs=thread_plan["sample_parallel_jobs"],
        threads_per_sample=threads_per_sample,
        sample_config=args.sample_config,
        library_strandness=strandness_info,
        splice_site_info=splice_site_info,
    )
    
    print(f"RNA-ALN completed. Results in {output_dir}")


if __name__ == "__main__":
    main()