"""
汇聚器模块 (assembler)
模块: 服务端 / 汇聚器
输入: 已收卷集合
输出: 完整密文文件

两道门: 卷齐才合并 (缺卷返回列表让客户端补); 合并后总哈希必须匹配才进解包。
"""

import hashlib
import shutil
from pathlib import Path

from task_manager import TaskManager


class AssemblerError(RuntimeError):
    """汇聚失败 (卷未齐/总哈希不一致)"""


class Assembler:
    def __init__(self, tasks: TaskManager, tmp_root: str | Path) -> None:
        self.tasks = tasks
        self.tmp_root = Path(tmp_root)

    def assemble(self, transfer_id: str) -> Path:
        """按 idx 顺序合并全部卷 -> 返回合并后密文路径"""
        transfer = self.tasks.get(transfer_id)
        if transfer is None:
            raise AssemblerError(f"任务不存在: {transfer_id}")
        total = transfer["chunk_count"]

        received = self.tasks.received(transfer_id)
        if len(received) < total:
            missing = sorted(set(range(1, total + 1)) - received)
            raise AssemblerError(f"缺卷 {missing}, 等齐再合并")

        chunk_dir = self.tmp_root / transfer_id
        merged_path = chunk_dir / "merged.bin"
        with open(merged_path, "wb") as out:
            for idx in range(1, total + 1):
                chunk_path = chunk_dir / f"chunk_{idx:04d}"
                with open(chunk_path, "rb") as f:
                    shutil.copyfileobj(f, out)

        # 总哈希校验
        h = hashlib.sha256()
        with open(merged_path, "rb") as f:
            for chunk in iter(lambda: f.read(1 << 20), b""):
                h.update(chunk)
        if h.hexdigest() != transfer["total_sha256"]:
            merged_path.unlink(missing_ok=True)
            raise AssemblerError("合并后总 SHA-256 不一致, 任务 failed (需整体重传)")

        return merged_path