"""ImageHub — content-addressed imaging dataset versioning (milestone 1 walking skeleton).

A user (investigator/PI) creates a dataset, uploads imaging files, and snapshots versions.
Files are stored as content-addressed, globally deduped blobs in MinIO (one blob per distinct
sha256). The current working file set lives in ``imagehub_dataset_files``; a version freezes a
manifest snapshot. Admin sees every dataset (the clinical data repository); a non-admin sees
only their own (their research data). Every mutation writes an append-only audit row.

Mounted under ``/api/v1`` in main.py → routes live at ``/api/v1/datasets/*``.
"""
from __future__ import annotations

import json
import re
import unicodedata
import uuid
from datetime import datetime, timedelta, timezone
from typing import Any, Optional

from fastapi import APIRouter, Body, File, Form, Header, HTTPException, Query, UploadFile
from pydantic import BaseModel, Field
from sqlalchemy import func, or_, select
from sqlalchemy.exc import IntegrityError

from src.auth_jwt import decode_access_token_user_id, decode_bearer_token
from src.imagehub_segmentation import MaskUpload, SegmentationError, SegmentationService
from src.imagehub_task_pipeline import (
    StageInfo,
    TaskPipelineError,
    compute_finalize,
    compute_review,
    initial_transition,
    validate_set_reference,
)
from src.initiative_db.engine import get_session, is_postgres_enabled
from src.initiative_db.models import (
    ImagehubBlob,
    ImagehubDataset,
    ImagehubDatasetAudit,
    ImagehubDatasetFile,
    ImagehubDatasetMember,
    ImagehubDatasetStage,
    ImagehubTask,
    ImagehubTaskReviewEvent,
    ImagehubVersion,
    ResearchProject,
    User,
)
from src.minio.storage import StorageError, storage

router = APIRouter(prefix="/datasets", tags=["imagehub"])

_VISIBILITIES = ("private", "internal", "public")
_ROLE_ADMIN = "Quản trị viên"
_ROLE_OWNER = "Chủ sở hữu"


# --------------------------------------------------------------------------- #
# Auth (mirrors research_routes / the extracted admin routers)
# --------------------------------------------------------------------------- #
def _jwt_roles(authorization: str | None) -> list[str]:
    p = decode_bearer_token(authorization)
    if not p:
        return []
    r = p.get("roles")
    return [str(x) for x in r] if isinstance(r, list) else []


def _require_authed_uid(authorization: str | None) -> uuid.UUID:
    uid = decode_access_token_user_id(authorization)
    if uid is None:
        raise HTTPException(status_code=401, detail="Đăng nhập để thực hiện thao tác.")
    return uid


def _is_admin(authorization: str | None) -> bool:
    return "admin" in _jwt_roles(authorization)


def _require_admin_uid(authorization: str | None) -> uuid.UUID:
    uid = _require_authed_uid(authorization)
    if not _is_admin(authorization):
        raise HTTPException(status_code=403, detail="Chỉ tài khoản quản trị mới thực hiện được.")
    return uid


def _require_db() -> None:
    if not is_postgres_enabled():
        raise HTTPException(status_code=503, detail="Cơ sở dữ liệu chưa sẵn sàng.")


def _role_label(authorization: str | None) -> str:
    return _ROLE_ADMIN if _is_admin(authorization) else _ROLE_OWNER


# --------------------------------------------------------------------------- #
# Helpers
# --------------------------------------------------------------------------- #
def _slugify(name: str) -> str:
    base = unicodedata.normalize("NFKD", name or "").encode("ascii", "ignore").decode("ascii")
    base = re.sub(r"[^a-zA-Z0-9]+", "-", base).strip("-").lower()
    return base[:60] or "dataset"


def _safe_logical_path(name: Optional[str]) -> str:
    """Flatten an uploaded filename to a safe logical path (basename; diacritics preserved)."""
    raw = (name or "").strip().replace("\\", "/")
    base = raw.rsplit("/", 1)[-1]
    base = re.sub(r"\s+", "_", base)
    base = "".join(ch for ch in base if unicodedata.category(ch)[0] != "C")
    return base[:200] or "file"


def _safe_folder_path(name: Optional[str]) -> str:
    """The sanitized relative directory of an uploaded path (folders kept, basename dropped).

    Mirrors _safe_logical_path but preserves slashes so a dataset can hold a real folder tree
    (e.g. the nnU-Net imagesTr/labelsTr layout). Rejects traversal, leading slashes and control
    characters. Returns '' when the upload carries no directory component.
    """
    raw = (name or "").strip().replace("\\", "/")
    head = raw.rsplit("/", 1)[0] if "/" in raw else ""
    parts: list[str] = []
    for seg in head.split("/"):
        seg = re.sub(r"\s+", "_", seg.strip())
        seg = "".join(ch for ch in seg if unicodedata.category(ch)[0] != "C")
        if seg in ("", ".", ".."):
            continue
        parts.append(seg)
    return "/".join(parts)[:400]


def _split_name_ext(name: str) -> tuple[str, str]:
    """Split a filename into (stem, ext), treating .nii.gz / .tar.gz as one extension."""
    low = name.lower()
    for double in (".nii.gz", ".tar.gz"):
        if low.endswith(double):
            return name[: -len(double)], name[-len(double):]
    dot = name.rfind(".")
    return (name, "") if dot <= 0 else (name[:dot], name[dot:])


def _case_number(stem: str) -> int | None:
    """The trailing integer of a case stem (after stripping any _NNNN channel tag), else None.
    e.g. "1"->1, "100"->100, "POLYP25_00001"->1, "10_0000"->10."""
    base = re.sub(r"_\d{4}$", "", stem)
    m = re.search(r"(\d+)$", base)
    return int(m.group(1)) if m else None


def _normalized_name(logical_path: str, prefix: str, is_label: bool) -> str | None:
    """Target name: image -> {prefix}_{NNNNN}_0000{ext}, label -> {prefix}_{NNNNN}{ext}, where
    NNNNN is the file's case number 5-digit zero-padded (so an image and its label share a case
    identifier, e.g. POLYP25_00001_0000.png + POLYP25_00001.png). Returns None when no case number
    can be derived or the name is already correct (idempotent)."""
    stem, ext = _split_name_ext(logical_path)
    num = _case_number(stem)
    if num is None:
        return None
    case_id = f"{prefix}_{num:05d}"
    new = f"{case_id}{ext}" if is_label else f"{case_id}_0000{ext}"
    return new if new != logical_path else None


def _coerce_tags(v: Any) -> list[str]:
    if isinstance(v, list):
        return [str(x).strip() for x in v if str(x).strip()]
    return []


def _sniff_imaging_meta(filename: Optional[str], data: bytes, media_type: str) -> dict[str, Any]:
    """Best-effort, synchronous imaging metadata. Never raises — returns {} on any failure.

    Heavy extraction (full tag dump, thumbnails, conversion) is deferred to the worker tier.
    """
    name = (filename or "").lower()
    # DICOM: "DICM" magic at byte 128, or a .dcm extension.
    if (len(data) > 132 and data[128:132] == b"DICM") or name.endswith(".dcm"):
        try:
            import io as _io

            import pydicom

            ds = pydicom.dcmread(_io.BytesIO(data), stop_before_pixels=True, force=True)
            out: dict[str, Any] = {"format": "dicom"}
            for tag, key in (
                ("Modality", "modality"),
                ("Rows", "rows"),
                ("Columns", "columns"),
                ("BodyPartExamined", "bodyPart"),
                ("StudyInstanceUID", "studyUid"),
                ("SeriesInstanceUID", "seriesUid"),
            ):
                val = getattr(ds, tag, None)
                if val is not None:
                    out[key] = int(val) if key in ("rows", "columns") else str(val)
            return out
        except Exception:
            return {}
    # NIfTI: .nii / .nii.gz — load via a temp file (nibabel needs a path for gz).
    if name.endswith(".nii") or name.endswith(".nii.gz"):
        import os as _os
        import tempfile as _tempfile

        suffix = ".nii.gz" if name.endswith(".nii.gz") else ".nii"
        tmp_path = ""
        try:
            import nibabel as nib

            with _tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
                tmp.write(data)
                tmp_path = tmp.name
            img = nib.load(tmp_path)
            zooms = [float(z) for z in img.header.get_zooms()]
            return {"format": "nifti", "shape": [int(s) for s in img.shape], "voxelSize": zooms}
        except Exception:
            return {}
        finally:
            if tmp_path:
                try:
                    _os.unlink(tmp_path)
                except OSError:
                    pass
    return {}


async def _actor_name(session, uid: uuid.UUID) -> str:
    u = await session.get(User, uid)
    if u is None:
        return ""
    return (u.full_name or u.email or "").strip()


async def _write_audit(
    session,
    dataset_id: uuid.UUID,
    actor_uid: Optional[uuid.UUID],
    actor_name: str,
    role_label: str,
    action: str,
    subject: str = "",
    detail: str = "",
) -> None:
    session.add(
        ImagehubDatasetAudit(
            dataset_id=dataset_id,
            actor_user_id=actor_uid,
            actor_name=actor_name or "",
            role_label=role_label or "",
            action=action,
            subject=subject or "",
            detail=detail or "",
        )
    )


async def _load_dataset(session, dataset_id: str, uid: uuid.UUID, is_admin: bool) -> ImagehubDataset:
    """Fetch a dataset enforcing owner-or-admin access (404 hides others' rows)."""
    try:
        did = uuid.UUID(dataset_id)
    except (ValueError, TypeError):
        raise HTTPException(status_code=404, detail="Không tìm thấy bộ dữ liệu.")
    row = (
        await session.execute(select(ImagehubDataset).where(ImagehubDataset.id == did))
    ).scalar_one_or_none()
    if row is None or (not is_admin and row.owner_user_id != uid):
        raise HTTPException(status_code=404, detail="Không tìm thấy bộ dữ liệu.")
    return row


# --------------------------------------------------------------------------- #
# Membership-aware access (multi-labeler). `_load_dataset` above stays owner-or-platform-admin and
# guards the management ops (create/delete/settings, upload, stage CRUD, member CRUD, generate,
# assign-to-others). The helpers below additionally admit dataset MEMBERS, for the read + task-work
# surface a labeler needs.
# --------------------------------------------------------------------------- #
_PROJECT_ADMIN_ROLES = ("owner", "admin", "project_admin")


async def _member_role(
    session, dataset: ImagehubDataset, uid: uuid.UUID, is_admin: bool
) -> Optional[str]:
    """The caller's effective role on a dataset: 'admin' (platform), 'owner', 'project_admin' or
    'member' (from the membership table), else None (no access)."""
    if is_admin:
        return "admin"
    if dataset.owner_user_id == uid:
        return "owner"
    return (
        await session.execute(
            select(ImagehubDatasetMember.role).where(
                ImagehubDatasetMember.dataset_id == dataset.id,
                ImagehubDatasetMember.user_id == uid,
            )
        )
    ).scalar_one_or_none()


async def _load_dataset_any(
    session, dataset_id: str, uid: uuid.UUID, is_admin: bool
) -> tuple[ImagehubDataset, str]:
    """Fetch a dataset enforcing owner-OR-member-OR-admin access (404 hides others). Returns
    (dataset, role) so callers can gate by role."""
    try:
        did = uuid.UUID(dataset_id)
    except (ValueError, TypeError):
        raise HTTPException(status_code=404, detail="Không tìm thấy bộ dữ liệu.")
    row = (
        await session.execute(select(ImagehubDataset).where(ImagehubDataset.id == did))
    ).scalar_one_or_none()
    if row is None:
        raise HTTPException(status_code=404, detail="Không tìm thấy bộ dữ liệu.")
    role = await _member_role(session, row, uid, is_admin)
    if role is None:
        raise HTTPException(status_code=404, detail="Không tìm thấy bộ dữ liệu.")
    return row, role


async def _load_dataset_read(session, dataset_id: str, uid: uuid.UUID, is_admin: bool) -> ImagehubDataset:
    """Owner/member/admin read access — the dataset only (drops the role)."""
    ds, _role = await _load_dataset_any(session, dataset_id, uid, is_admin)
    return ds


async def _load_dataset_admin(
    session, dataset_id: str, uid: uuid.UUID, is_admin: bool
) -> tuple[ImagehubDataset, str]:
    """Project-management access: owner, platform-admin, or a project_admin member (403 for plain
    members, 404 for non-members). Dataset-structural / destructive ops still use `_load_dataset`."""
    ds, role = await _load_dataset_any(session, dataset_id, uid, is_admin)
    if role not in _PROJECT_ADMIN_ROLES:
        raise HTTPException(status_code=403, detail="Chỉ quản trị dự án mới thực hiện được.")
    return ds, role


def _can_work_task(role: str, task_assignee: Optional[uuid.UUID], uid: uuid.UUID) -> bool:
    """A project admin/owner may work any task; a plain member only the tasks assigned to them."""
    if role in _PROJECT_ADMIN_ROLES:
        return True
    return task_assignee is not None and task_assignee == uid


async def _counts(session, dataset_ids: list[uuid.UUID]) -> tuple[dict, dict]:
    """Grouped file + version counts for a set of datasets (avoids N+1 in the list view)."""
    if not dataset_ids:
        return {}, {}
    file_rows = (
        await session.execute(
            select(ImagehubDatasetFile.dataset_id, func.count())
            .where(ImagehubDatasetFile.dataset_id.in_(dataset_ids))
            .group_by(ImagehubDatasetFile.dataset_id)
        )
    ).all()
    ver_rows = (
        await session.execute(
            select(ImagehubVersion.dataset_id, func.count())
            .where(ImagehubVersion.dataset_id.in_(dataset_ids))
            .group_by(ImagehubVersion.dataset_id)
        )
    ).all()
    return {r[0]: int(r[1]) for r in file_rows}, {r[0]: int(r[1]) for r in ver_rows}


# --------------------------------------------------------------------------- #
# Schemas
# --------------------------------------------------------------------------- #
class DatasetOut(BaseModel):
    id: str
    ownerUserId: str
    ownerEmail: Optional[str] = None
    name: str = ""
    slug: str = ""
    description: str = ""
    visibility: str = "private"
    modalityTags: list[str] = Field(default_factory=list)
    labelMap: dict[str, str] = Field(default_factory=dict)
    defaultBranch: str = "main"
    researchProjectId: Optional[str] = None
    fileCount: int = 0
    versionCount: int = 0
    createdAt: Optional[datetime] = None
    updatedAt: Optional[datetime] = None


class DatasetCreateIn(BaseModel):
    name: str = Field(default="", max_length=300)
    description: str = Field(default="", max_length=4000)
    visibility: str = "private"
    modalityTags: list[str] = Field(default_factory=list)
    researchProjectId: Optional[str] = None


class DatasetUpdateIn(BaseModel):
    name: Optional[str] = Field(default=None, max_length=300)
    description: Optional[str] = Field(default=None, max_length=4000)
    visibility: Optional[str] = None
    modalityTags: Optional[list[str]] = None
    labelMap: Optional[dict[str, str]] = None


class VersionCreateIn(BaseModel):
    message: Optional[str] = Field(default=None, max_length=2000)


class FileOut(BaseModel):
    id: str
    logicalPath: str
    folderPath: str = ""
    sha256: str
    size: int
    mediaType: str
    imagingMeta: dict[str, Any] = Field(default_factory=dict)
    fileKind: str = "image"
    parentFileId: Optional[str] = None
    organLabel: str = ""
    uploadedAt: Optional[datetime] = None
    downloadUrl: Optional[str] = None


class NormalizeIn(BaseModel):
    prefix: str  # lesion+year code, e.g. "POLYP25"


class NormalizeResultOut(BaseModel):
    ok: bool
    renamed: int
    skipped: int
    files: list[dict[str, str]]  # {id, oldPath, newPath, folderPath}


class VersionOut(BaseModel):
    id: str
    seq: int
    message: str = ""
    fileCount: int = 0
    manifest: list[dict[str, Any]] = Field(default_factory=list)
    createdAt: Optional[datetime] = None


class AuditOut(BaseModel):
    id: int
    occurredAt: Optional[datetime] = None
    actorName: str = ""
    roleLabel: str = ""
    action: str
    subject: str = ""
    detail: str = ""


def _coerce_label_map(value: Any) -> dict[str, str]:
    """Sanitize a per-dataset value->name label map: positive-int string keys, trimmed non-empty
    names (<=100 chars), bounded to 512 entries. Returns {} for anything malformed so the viewer
    falls back to the built-in TotalSegmentator names."""
    if not isinstance(value, dict):
        return {}
    out: dict[str, str] = {}
    for k, v in value.items():
        ks = str(k)
        if not (ks.isascii() and ks.isdigit()):  # plain positive-int keys only ("1".."N")
            continue
        iv = int(ks)
        if iv <= 0 or not isinstance(v, str):
            continue
        name = v.strip()
        if not name:
            continue
        out[str(iv)] = name[:100]
        if len(out) >= 512:
            break
    return out


def _ds_to_out(
    row: ImagehubDataset, file_count: int = 0, version_count: int = 0, owner_email: Optional[str] = None
) -> DatasetOut:
    return DatasetOut(
        id=str(row.id),
        ownerUserId=str(row.owner_user_id),
        ownerEmail=owner_email,
        name=row.name or "",
        slug=row.slug or "",
        description=row.description or "",
        visibility=row.visibility or "private",
        modalityTags=_coerce_tags(row.modality_tags),
        labelMap=_coerce_label_map(row.label_map),
        defaultBranch=row.default_branch or "main",
        researchProjectId=str(row.research_project_id) if row.research_project_id else None,
        fileCount=file_count,
        versionCount=version_count,
        createdAt=row.created_at,
        updatedAt=row.updated_at,
    )


def _version_to_out(row: ImagehubVersion) -> VersionOut:
    manifest = row.manifest if isinstance(row.manifest, list) else []
    return VersionOut(
        id=str(row.id),
        seq=row.seq,
        message=row.message or "",
        fileCount=len(manifest),
        manifest=manifest,
        createdAt=row.created_at,
    )


# --------------------------------------------------------------------------- #
# Endpoints — datasets
# --------------------------------------------------------------------------- #
@router.post("", response_model=DatasetOut)
async def create_dataset(
    payload: Optional[DatasetCreateIn] = Body(None),
    authorization: Optional[str] = Header(None),
) -> DatasetOut:
    """Authed: create a dataset owned by the current user."""
    _require_db()
    uid = _require_authed_uid(authorization)
    p = payload or DatasetCreateIn()
    name = (p.name or "").strip()
    if not name:
        raise HTTPException(status_code=422, detail="Cần nhập tên bộ dữ liệu.")
    visibility = p.visibility if p.visibility in _VISIBILITIES else "private"
    project_uuid: Optional[uuid.UUID] = None
    if p.researchProjectId:
        try:
            project_uuid = uuid.UUID(p.researchProjectId)
        except (ValueError, TypeError):
            raise HTTPException(status_code=422, detail="Đề tài không hợp lệ.")
    async with get_session() as session:
        # If linking to a research project ("workspace"), it must exist and be owned by the
        # caller (or the caller is a platform admin). Read BEFORE add to avoid an autoflush.
        if project_uuid is not None:
            owner = (
                await session.execute(
                    select(ResearchProject.owner_user_id).where(ResearchProject.id == project_uuid)
                )
            ).scalar_one_or_none()
            if owner is None:
                raise HTTPException(status_code=422, detail="Đề tài không tồn tại.")
            if owner != uid and not _is_admin(authorization):
                raise HTTPException(status_code=403, detail="Bạn không sở hữu đề tài này.")
        row = ImagehubDataset(
            id=uuid.uuid4(),
            owner_user_id=uid,
            name=name,
            slug=_slugify(name),
            description=(p.description or "").strip(),
            visibility=visibility,
            modality_tags=_coerce_tags(p.modalityTags),
            research_project_id=project_uuid,
        )
        session.add(row)
        await session.flush()
        await _write_audit(
            session, row.id, uid, await _actor_name(session, uid), _ROLE_OWNER,
            "Tạo bộ dữ liệu", name,
        )
        await session.commit()
        await session.refresh(row)
        return _ds_to_out(row)


@router.get("", response_model=list[DatasetOut])
async def list_datasets(
    scope: str = "mine",
    authorization: Optional[str] = Header(None),
    projectId: Optional[str] = None,
) -> list[DatasetOut]:
    """List datasets. Non-admin always sees only their own; admin may pass ?scope=all (clinical repo).
    ?projectId=<id> further restricts to datasets linked to that research project ("workspace")."""
    _require_db()
    uid = _require_authed_uid(authorization)
    is_admin = _is_admin(authorization)
    async with get_session() as session:
        stmt = (
            select(ImagehubDataset, User.email)
            .join(User, User.id == ImagehubDataset.owner_user_id)
            .order_by(ImagehubDataset.created_at.desc())
        )
        if scope != "all" or not is_admin:
            member_ds = select(ImagehubDatasetMember.dataset_id).where(
                ImagehubDatasetMember.user_id == uid
            )
            stmt = stmt.where(
                or_(ImagehubDataset.owner_user_id == uid, ImagehubDataset.id.in_(member_ds))
            )
        if projectId:
            try:
                stmt = stmt.where(ImagehubDataset.research_project_id == uuid.UUID(projectId))
            except (ValueError, TypeError):
                raise HTTPException(status_code=422, detail="Đề tài không hợp lệ.")
        rows = (await session.execute(stmt)).all()
        datasets = [r[0] for r in rows]
        emails = {r[0].id: r[1] for r in rows}
        files, versions = await _counts(session, [d.id for d in datasets])
        return [
            _ds_to_out(d, files.get(d.id, 0), versions.get(d.id, 0), emails.get(d.id))
            for d in datasets
        ]


@router.get("/{dataset_id}", response_model=DatasetOut)
async def get_dataset(dataset_id: str, authorization: Optional[str] = Header(None)) -> DatasetOut:
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        row = await _load_dataset_read(session, dataset_id, uid, _is_admin(authorization))
        files, versions = await _counts(session, [row.id])
        owner = await session.get(User, row.owner_user_id)
        return _ds_to_out(row, files.get(row.id, 0), versions.get(row.id, 0), owner.email if owner else None)


@router.put("/{dataset_id}", response_model=DatasetOut)
async def update_dataset(
    dataset_id: str,
    payload: DatasetUpdateIn = Body(...),
    authorization: Optional[str] = Header(None),
) -> DatasetOut:
    """Owner or admin: update dataset metadata."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        row = await _load_dataset(session, dataset_id, uid, _is_admin(authorization))
        if payload.name is not None:
            new_name = payload.name.strip()
            if not new_name:
                raise HTTPException(status_code=422, detail="Tên bộ dữ liệu không được để trống.")
            row.name = new_name
            row.slug = _slugify(new_name)
        if payload.description is not None:
            row.description = payload.description.strip()
        if payload.visibility is not None:
            if payload.visibility not in _VISIBILITIES:
                raise HTTPException(status_code=422, detail="Mức hiển thị không hợp lệ.")
            row.visibility = payload.visibility
        if payload.modalityTags is not None:
            row.modality_tags = _coerce_tags(payload.modalityTags)
        if payload.labelMap is not None:
            row.label_map = _coerce_label_map(payload.labelMap)
        row.updated_at = datetime.now(tz=timezone.utc)
        await _write_audit(
            session, row.id, uid, await _actor_name(session, uid), _role_label(authorization),
            "Cập nhật bộ dữ liệu", row.name,
        )
        await session.commit()
        await session.refresh(row)
        files, versions = await _counts(session, [row.id])
        return _ds_to_out(row, files.get(row.id, 0), versions.get(row.id, 0))


@router.delete("/{dataset_id}")
async def delete_dataset(dataset_id: str, authorization: Optional[str] = Header(None)) -> dict[str, Any]:
    """Owner or admin: delete a dataset. Files/versions/audit cascade; blobs are left (no GC in v1)."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        row = await _load_dataset(session, dataset_id, uid, _is_admin(authorization))
        await session.delete(row)
        await session.commit()
    return {"ok": True}


# --------------------------------------------------------------------------- #
# Endpoints — files (content-addressed upload + browse)
# --------------------------------------------------------------------------- #
@router.post("/{dataset_id}/files")
async def upload_files(
    dataset_id: str,
    files: list[UploadFile] = File(...),
    authorization: Optional[str] = Header(None),
    paths: Optional[str] = Form(None),
) -> dict[str, Any]:
    """Owner or admin: upload one or more files. Each is content-addressed + deduped; a file at an
    existing (folder, logical path) is replaced. ``paths`` is an optional JSON array of relative
    upload paths (index-aligned with ``files``) so a directory structure is preserved as folder_path.
    Best-effort imaging metadata is sniffed synchronously."""
    _require_db()
    uid = _require_authed_uid(authorization)
    is_admin = _is_admin(authorization)
    results: list[dict[str, Any]] = []
    async with get_session() as session:
        ds = await _load_dataset(session, dataset_id, uid, is_admin)
        rel_paths: list[str] = []
        if paths:
            try:
                parsed_paths = json.loads(paths)
                if isinstance(parsed_paths, list):
                    rel_paths = [str(p) for p in parsed_paths]
            except (ValueError, TypeError):
                rel_paths = []
        for i, uf in enumerate(files):
            data = await uf.read()
            if not data:
                continue
            media_type = uf.content_type or "application/octet-stream"
            try:
                blob = await storage.put_blob(data, media_type)
            except ValueError as exc:
                raise HTTPException(status_code=413, detail=str(exc)) from exc
            except StorageError as exc:
                raise HTTPException(status_code=502, detail=f"Lưu trữ thất bại: {exc}") from exc
            if await session.get(ImagehubBlob, blob["sha256"]) is None:
                session.add(
                    ImagehubBlob(
                        sha256=blob["sha256"],
                        size_bytes=blob["size"],
                        media_type=media_type,
                        storage_bucket=blob["bucket"],
                        storage_key=blob["key"],
                    )
                )
                await session.flush()
            logical_path = _safe_logical_path(uf.filename)
            rel = rel_paths[i] if i < len(rel_paths) else (uf.filename or "")
            folder_path = _safe_folder_path(rel)
            meta = _sniff_imaging_meta(uf.filename, data, media_type)
            file_row = (
                await session.execute(
                    select(ImagehubDatasetFile).where(
                        ImagehubDatasetFile.dataset_id == ds.id,
                        ImagehubDatasetFile.folder_path == folder_path,
                        ImagehubDatasetFile.logical_path == logical_path,
                    )
                )
            ).scalar_one_or_none()
            if file_row is None:
                file_row = ImagehubDatasetFile(
                    id=uuid.uuid4(), dataset_id=ds.id, folder_path=folder_path, logical_path=logical_path
                )
                session.add(file_row)
            file_row.blob_sha256 = blob["sha256"]
            file_row.size_bytes = blob["size"]
            file_row.media_type = media_type
            file_row.imaging_meta = meta
            file_row.uploaded_by = uid
            file_row.updated_at = datetime.now(tz=timezone.utc)
            results.append(
                {
                    "path": f"{folder_path}/{logical_path}" if folder_path else logical_path,
                    "sha256": blob["sha256"],
                    "size": blob["size"],
                    "deduped": blob["deduped"],
                    "imagingMeta": meta,
                }
            )
        ds.updated_at = datetime.now(tz=timezone.utc)
        await _write_audit(
            session, ds.id, uid, await _actor_name(session, uid), _role_label(authorization),
            "Tải tệp lên", f"{len(results)} tệp",
        )
        await session.commit()
    return {"ok": True, "files": results}


@router.get("/{dataset_id}/files", response_model=list[FileOut])
async def list_files(dataset_id: str, authorization: Optional[str] = Header(None)) -> list[FileOut]:
    """Owner, member or admin: the dataset's current working files, each with a presigned URL."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds = await _load_dataset_read(session, dataset_id, uid, _is_admin(authorization))
        rows = (
            await session.execute(
                select(ImagehubDatasetFile, ImagehubBlob)
                .join(ImagehubBlob, ImagehubBlob.sha256 == ImagehubDatasetFile.blob_sha256)
                .where(ImagehubDatasetFile.dataset_id == ds.id)
                .order_by(ImagehubDatasetFile.folder_path, ImagehubDatasetFile.logical_path)
            )
        ).all()
        out: list[FileOut] = []
        for f, b in rows:
            try:
                url = await storage.get_download_url(
                    b.storage_bucket, b.storage_key, filename=f.logical_path, inline=False
                )
            except Exception:
                url = None
            out.append(
                FileOut(
                    id=str(f.id),
                    logicalPath=f.logical_path,
                    folderPath=f.folder_path or "",
                    sha256=f.blob_sha256,
                    size=f.size_bytes,
                    mediaType=f.media_type,
                    imagingMeta=f.imaging_meta if isinstance(f.imaging_meta, dict) else {},
                    fileKind=f.file_kind or "image",
                    parentFileId=str(f.parent_file_id) if f.parent_file_id else None,
                    organLabel=f.organ_label or "",
                    uploadedAt=f.created_at,
                    downloadUrl=url,
                )
            )
        return out


@router.get("/{dataset_id}/files/{file_id}/download")
async def download_file(
    dataset_id: str, file_id: str, authorization: Optional[str] = Header(None)
) -> dict[str, Any]:
    """Owner, member or admin: a fresh presigned download URL for a single file."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds = await _load_dataset_read(session, dataset_id, uid, _is_admin(authorization))
        try:
            fid = uuid.UUID(file_id)
        except (ValueError, TypeError):
            raise HTTPException(status_code=404, detail="Không tìm thấy tệp.")
        row = (
            await session.execute(
                select(ImagehubDatasetFile, ImagehubBlob)
                .join(ImagehubBlob, ImagehubBlob.sha256 == ImagehubDatasetFile.blob_sha256)
                .where(ImagehubDatasetFile.id == fid, ImagehubDatasetFile.dataset_id == ds.id)
            )
        ).first()
        if row is None:
            raise HTTPException(status_code=404, detail="Không tìm thấy tệp.")
        f, b = row
        url = await storage.get_download_url(
            b.storage_bucket, b.storage_key, filename=f.logical_path, inline=False
        )
        return {"url": url, "logicalPath": f.logical_path}


@router.post("/{dataset_id}/files/{parent_file_id}/segmentations")
async def upload_segmentations(
    dataset_id: str,
    parent_file_id: str,
    masks: list[UploadFile] = File(...),
    organs: Optional[list[str]] = Form(None),
    authorization: Optional[str] = Header(None),
) -> dict[str, Any]:
    """Owner or admin: upload organ-mask files and link them to a parent image file.

    Each mask is content-addressed + stored like any file, but recorded with
    ``file_kind='segmentation'`` and ``parent_file_id`` pointing at the image it
    segments. ``organs[i]`` is the organ label for ``masks[i]`` (parallel arrays;
    falls back to the filename). Domain rules live in ``SegmentationService`` — the
    handler is just transport (read multipart → service → audit → commit)."""
    _require_db()
    uid = _require_authed_uid(authorization)
    is_admin = _is_admin(authorization)
    organ_labels = organs or []
    async with get_session() as session:
        ds = await _load_dataset(session, dataset_id, uid, is_admin)
        items: list[MaskUpload] = []
        for i, uf in enumerate(masks):
            items.append(
                MaskUpload(
                    filename=uf.filename or "mask.nii.gz",
                    data=await uf.read(),
                    media_type=uf.content_type or "application/octet-stream",
                    organ_label=organ_labels[i] if i < len(organ_labels) else "",
                )
            )
        service = SegmentationService(
            session,
            put_blob=storage.put_blob,
            sniff_meta=_sniff_imaging_meta,
            safe_name=_safe_logical_path,
        )
        try:
            linked = await service.link_masks(ds, parent_file_id, items, uid)
        except SegmentationError as exc:
            raise HTTPException(status_code=exc.status, detail=str(exc)) from exc
        except ValueError as exc:  # storage.put_blob size cap → 413
            raise HTTPException(status_code=413, detail=str(exc)) from exc
        except StorageError as exc:
            raise HTTPException(status_code=502, detail=f"Lưu trữ thất bại: {exc}") from exc
        ds.updated_at = datetime.now(tz=timezone.utc)
        await _write_audit(
            session, ds.id, uid, await _actor_name(session, uid), _role_label(authorization),
            "Tải mặt nạ phân vùng", f"{len(linked)} mặt nạ",
        )
        await session.commit()
        results = [
            {
                "id": str(r.id),
                "logicalPath": r.logical_path,
                "organLabel": r.organ_label,
                "parentFileId": str(r.parent_file_id) if r.parent_file_id else None,
            }
            for r in linked
        ]
    return {"ok": True, "masks": results}


@router.post("/{dataset_id}/files/normalize-channels", response_model=NormalizeResultOut)
async def normalize_channels(
    dataset_id: str, payload: NormalizeIn, authorization: Optional[str] = Header(None)
) -> NormalizeResultOut:
    """Owner or admin: rename files to the {prefix}_{caseID}_0000.{ext} convention — images under
    imagesTr/imagesTs get the _0000 channel suffix, labels under labelsTr/labelsTs stay channel-free,
    and caseID is the file's number 5-digit zero-padded so an image and its label share a case id
    (e.g. POLYP25_00001_0000.png + POLYP25_00001.png). Idempotent; logical rename only — no blob move."""
    _require_db()
    uid = _require_authed_uid(authorization)
    is_admin = _is_admin(authorization)
    prefix = re.sub(r"[^A-Za-z0-9]", "", payload.prefix or "").upper()
    if not prefix:
        raise HTTPException(status_code=400, detail="Thiếu mã tổn thương (tiền tố).")
    async with get_session() as session:
        ds = await _load_dataset(session, dataset_id, uid, is_admin)
        rows = (
            await session.execute(
                select(ImagehubDatasetFile).where(ImagehubDatasetFile.dataset_id == ds.id)
            )
        ).scalars().all()
        existing = {(r.folder_path, r.logical_path) for r in rows}
        # Resolve the actor + role BEFORE mutating any logical_path: a read after a mutation would
        # autoflush the pending UPDATE and raise a unique-violation outside the commit try/except.
        actor = await _actor_name(session, uid)
        role = _role_label(authorization)
        results: list[dict[str, str]] = []
        plan: list[tuple[ImagehubDatasetFile, str]] = []
        skipped = 0
        for r in rows:
            seg = (r.folder_path or "").strip("/").split("/")[-1]
            if r.file_kind == "segmentation":
                continue
            if seg in ("imagesTr", "imagesTs"):
                is_label = False
            elif seg in ("labelsTr", "labelsTs"):
                is_label = True
            else:
                continue  # root / other folders are left alone
            new = _normalized_name(r.logical_path, prefix, is_label)
            if new is None:
                skipped += 1
                continue
            if (r.folder_path, new) in existing:
                skipped += 1
                continue
            results.append(
                {
                    "id": str(r.id),
                    "oldPath": r.logical_path,
                    "newPath": new,
                    "folderPath": r.folder_path or "",
                }
            )
            plan.append((r, new))
        renamed = len(plan)
        if renamed:
            for r, new in plan:
                r.logical_path = new
                r.updated_at = datetime.now(tz=timezone.utc)
            ds.updated_at = datetime.now(tz=timezone.utc)
            await _write_audit(
                session, ds.id, uid, actor, role,
                action="Chuẩn hoá tên tệp", subject=f"{renamed} tệp",
                detail=f"Đổi tên sang định dạng {prefix}_<caseID>_0000",
            )
        try:
            await session.commit()
        except IntegrityError:
            await session.rollback()
            raise HTTPException(status_code=409, detail="Tên tệp bị trùng sau khi chuẩn hoá.")
    return NormalizeResultOut(ok=True, renamed=renamed, skipped=skipped, files=results)


# --------------------------------------------------------------------------- #
# Endpoints — versions (the snapshot spine) + audit
# --------------------------------------------------------------------------- #
@router.post("/{dataset_id}/versions", response_model=VersionOut)
async def create_version(
    dataset_id: str,
    payload: Optional[VersionCreateIn] = Body(None),
    authorization: Optional[str] = Header(None),
) -> VersionOut:
    """Owner or admin: freeze the current working files into a new version snapshot."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds = await _load_dataset(session, dataset_id, uid, _is_admin(authorization))
        files = (
            await session.execute(
                select(ImagehubDatasetFile)
                .where(ImagehubDatasetFile.dataset_id == ds.id)
                .order_by(ImagehubDatasetFile.logical_path)
            )
        ).scalars().all()
        if not files:
            raise HTTPException(status_code=422, detail="Bộ dữ liệu chưa có tệp nào để tạo phiên bản.")
        manifest = [
            {"logicalPath": f.logical_path, "blobSha256": f.blob_sha256,
             "size": f.size_bytes, "mediaType": f.media_type}
            for f in files
        ]
        max_seq = (
            await session.execute(
                select(func.coalesce(func.max(ImagehubVersion.seq), 0)).where(
                    ImagehubVersion.dataset_id == ds.id
                )
            )
        ).scalar_one()
        parent = (
            await session.execute(
                select(ImagehubVersion.id)
                .where(ImagehubVersion.dataset_id == ds.id)
                .order_by(ImagehubVersion.seq.desc())
                .limit(1)
            )
        ).scalar_one_or_none()
        msg = ((payload.message if payload else None) or "").strip()
        row = ImagehubVersion(
            id=uuid.uuid4(),
            dataset_id=ds.id,
            seq=int(max_seq) + 1,
            message=msg,
            manifest=manifest,
            parent_version_id=parent,
            author_user_id=uid,
        )
        session.add(row)
        await session.flush()
        await _write_audit(
            session, ds.id, uid, await _actor_name(session, uid), _role_label(authorization),
            "Tạo phiên bản", f"v{row.seq}", msg,
        )
        await session.commit()
        await session.refresh(row)
        return _version_to_out(row)


@router.get("/{dataset_id}/versions", response_model=list[VersionOut])
async def list_versions(dataset_id: str, authorization: Optional[str] = Header(None)) -> list[VersionOut]:
    """Owner or admin: version history, newest first."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds = await _load_dataset(session, dataset_id, uid, _is_admin(authorization))
        rows = (
            await session.execute(
                select(ImagehubVersion)
                .where(ImagehubVersion.dataset_id == ds.id)
                .order_by(ImagehubVersion.seq.desc())
            )
        ).scalars().all()
        return [_version_to_out(r) for r in rows]


@router.get("/{dataset_id}/audit", response_model=list[AuditOut])
async def list_audit(dataset_id: str, authorization: Optional[str] = Header(None)) -> list[AuditOut]:
    """Owner or admin: the append-only audit trail for a dataset, newest first."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds = await _load_dataset(session, dataset_id, uid, _is_admin(authorization))
        rows = (
            await session.execute(
                select(ImagehubDatasetAudit)
                .where(ImagehubDatasetAudit.dataset_id == ds.id)
                .order_by(ImagehubDatasetAudit.occurred_at.desc(), ImagehubDatasetAudit.id.desc())
            )
        ).scalars().all()
        return [
            AuditOut(
                id=r.id,
                occurredAt=r.occurred_at,
                actorName=r.actor_name,
                roleLabel=r.role_label,
                action=r.action,
                subject=r.subject,
                detail=r.detail,
            )
            for r in rows
        ]


# --------------------------------------------------------------------------- #
# Endpoints — labeling-pipeline stages (Label -> Review_1 -> Review_2 ...)
# --------------------------------------------------------------------------- #
_STAGE_KINDS = ("label", "review")


class StageOut(BaseModel):
    id: str
    name: str = ""
    kind: str = "label"
    seq: int = 0
    reviewPercent: Optional[int] = None
    autoAssign: bool = True


class StageCreateIn(BaseModel):
    name: str = Field(default="", max_length=200)
    kind: str = "label"
    reviewPercent: Optional[int] = Field(default=None, ge=0, le=100)
    autoAssign: bool = True


class StageUpdateIn(BaseModel):
    name: Optional[str] = Field(default=None, max_length=200)
    reviewPercent: Optional[int] = Field(default=None, ge=0, le=100)
    autoAssign: Optional[bool] = None
    seq: Optional[int] = None


def _stage_to_out(row: ImagehubDatasetStage) -> StageOut:
    return StageOut(
        id=str(row.id),
        name=row.name or "",
        kind=row.kind or "label",
        seq=row.seq,
        reviewPercent=row.review_percent,
        autoAssign=bool(row.auto_assign),
    )


async def _load_stage(session, dataset_id: uuid.UUID, stage_id: str) -> ImagehubDatasetStage:
    try:
        sid = uuid.UUID(stage_id)
    except (ValueError, TypeError):
        raise HTTPException(status_code=404, detail="Không tìm thấy giai đoạn.")
    stage = (
        await session.execute(
            select(ImagehubDatasetStage).where(
                ImagehubDatasetStage.id == sid, ImagehubDatasetStage.dataset_id == dataset_id
            )
        )
    ).scalar_one_or_none()
    if stage is None:
        raise HTTPException(status_code=404, detail="Không tìm thấy giai đoạn.")
    return stage


@router.get("/{dataset_id}/stages", response_model=list[StageOut])
async def list_stages(dataset_id: str, authorization: Optional[str] = Header(None)) -> list[StageOut]:
    """Owner, member or admin: the dataset's labeling-pipeline stages, in pipeline order."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds = await _load_dataset_read(session, dataset_id, uid, _is_admin(authorization))
        rows = (
            await session.execute(
                select(ImagehubDatasetStage)
                .where(ImagehubDatasetStage.dataset_id == ds.id)
                .order_by(ImagehubDatasetStage.seq, ImagehubDatasetStage.created_at)
            )
        ).scalars().all()
        return [_stage_to_out(r) for r in rows]


@router.post("/{dataset_id}/stages", response_model=StageOut)
async def add_stage(
    dataset_id: str,
    payload: StageCreateIn = Body(...),
    authorization: Optional[str] = Header(None),
) -> StageOut:
    """Owner or admin: append a stage to the pipeline."""
    _require_db()
    uid = _require_authed_uid(authorization)
    if payload.kind not in _STAGE_KINDS:
        raise HTTPException(status_code=422, detail="Loại giai đoạn không hợp lệ.")
    name = (payload.name or "").strip()
    if not name:
        raise HTTPException(status_code=422, detail="Tên giai đoạn không được để trống.")
    async with get_session() as session:
        ds = await _load_dataset(session, dataset_id, uid, _is_admin(authorization))
        # Resolve the actor BEFORE adding the new row: a query here would otherwise autoflush the
        # pending INSERT and raise the unique-violation outside the commit() try/except below.
        actor = await _actor_name(session, uid)
        max_seq = (
            await session.execute(
                select(func.max(ImagehubDatasetStage.seq)).where(ImagehubDatasetStage.dataset_id == ds.id)
            )
        ).scalar()
        stage = ImagehubDatasetStage(
            dataset_id=ds.id,
            name=name,
            kind=payload.kind,
            seq=(int(max_seq) + 1) if max_seq is not None else 0,
            review_percent=payload.reviewPercent if payload.kind == "review" else None,
            auto_assign=payload.autoAssign,
        )
        session.add(stage)
        await _write_audit(
            session, ds.id, uid, actor, _role_label(authorization),
            "Thêm giai đoạn", name,
        )
        try:
            await session.commit()
        except IntegrityError:
            await session.rollback()
            raise HTTPException(status_code=409, detail="Tên giai đoạn đã tồn tại trong bộ dữ liệu.")
        await session.refresh(stage)
        return _stage_to_out(stage)


@router.patch("/{dataset_id}/stages/{stage_id}", response_model=StageOut)
async def update_stage(
    dataset_id: str,
    stage_id: str,
    payload: StageUpdateIn = Body(...),
    authorization: Optional[str] = Header(None),
) -> StageOut:
    """Owner or admin: rename a stage, set its review %, toggle Automatic Task Assignment, reorder."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds = await _load_dataset(session, dataset_id, uid, _is_admin(authorization))
        stage = await _load_stage(session, ds.id, stage_id)
        # Resolve the actor before mutating the row (autoflush-before-commit guard; see add_stage).
        actor = await _actor_name(session, uid)
        if payload.name is not None:
            nm = payload.name.strip()
            if not nm:
                raise HTTPException(status_code=422, detail="Tên giai đoạn không được để trống.")
            stage.name = nm
        if payload.reviewPercent is not None:
            stage.review_percent = payload.reviewPercent
        if payload.autoAssign is not None:
            stage.auto_assign = payload.autoAssign
        if payload.seq is not None:
            stage.seq = payload.seq
        stage.updated_at = datetime.now(tz=timezone.utc)
        await _write_audit(
            session, ds.id, uid, actor, _role_label(authorization),
            "Cập nhật giai đoạn", stage.name,
        )
        try:
            await session.commit()
        except IntegrityError:
            await session.rollback()
            raise HTTPException(status_code=409, detail="Tên giai đoạn đã tồn tại trong bộ dữ liệu.")
        await session.refresh(stage)
        return _stage_to_out(stage)


@router.delete("/{dataset_id}/stages/{stage_id}")
async def delete_stage(
    dataset_id: str, stage_id: str, authorization: Optional[str] = Header(None)
) -> dict[str, Any]:
    """Owner or admin: remove a stage from the pipeline."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds = await _load_dataset(session, dataset_id, uid, _is_admin(authorization))
        stage = await _load_stage(session, ds.id, stage_id)
        actor = await _actor_name(session, uid)
        name = stage.name
        await session.delete(stage)
        await _write_audit(
            session, ds.id, uid, actor, _role_label(authorization),
            "Xóa giai đoạn", name,
        )
        await session.commit()
    return {"ok": True}


# --------------------------------------------------------------------------- #
# Endpoints — task pipeline (a file flows Label -> Review_n -> Ground Truth)
#
# Single-user MVP: a Task is one row per *image* file (imagehub_tasks), created on demand from the
# dataset's files once the pipeline has at least one stage. Access reuses the dataset owner-or-admin
# gate; multi-labeler membership/assignment is a later phase. Every write follows the
# autoflush-before-commit guard (resolve reads before add/mutate; see add_stage).
# --------------------------------------------------------------------------- #
class TaskOut(BaseModel):
    id: str
    name: str = ""
    fileId: str
    fileLogicalPath: str = ""
    currentStageId: Optional[str] = None
    currentStageName: Optional[str] = None
    pipelineState: str = "inLabel"
    queueStatus: str = "assigned"
    assigneeUserId: Optional[str] = None
    assigneeName: Optional[str] = None
    assignmentMode: str = "auto"
    priority: float = 0.0
    isReferenceStandard: bool = False
    createdAt: datetime
    updatedAt: datetime


class GenerateTasksOut(BaseModel):
    created: int = 0
    total: int = 0


class ReviewIn(BaseModel):
    decision: str = Field(..., description="accept | acceptWithCorrections | reject")
    note: Optional[str] = Field(default=None, max_length=2000)


class ReviewStatsOut(BaseModel):
    accepted: int = 0
    acceptWithCorrections: int = 0
    rejected: int = 0


class PriorityIn(BaseModel):
    priority: float = Field(..., ge=0, le=1)


class SetReferenceIn(BaseModel):
    isReferenceStandard: bool = True


class TaskDetailOut(TaskOut):
    """A single task plus its saved annotations (the AnnotationTool's working payload)."""

    annotations: list[Any] = Field(default_factory=list)


class SaveIn(BaseModel):
    annotations: list[Any] = Field(default_factory=list)
    submit: bool = False


def _assignee_label(full_name: Optional[str], email: Optional[str]) -> Optional[str]:
    name = (full_name or email or "").strip()
    return name or None


def _build_task_out(
    task: ImagehubTask,
    file_logical_path: Optional[str],
    stage_name: Optional[str],
    assignee_name: Optional[str],
) -> TaskOut:
    return TaskOut(
        id=str(task.id),
        name=task.name or "",
        fileId=str(task.dataset_file_id),
        fileLogicalPath=file_logical_path or "",
        currentStageId=str(task.current_stage_id) if task.current_stage_id else None,
        currentStageName=stage_name,
        pipelineState=task.pipeline_state,
        queueStatus=task.queue_status,
        assigneeUserId=str(task.assignee_user_id) if task.assignee_user_id else None,
        assigneeName=assignee_name,
        assignmentMode=task.assignment_mode,
        priority=float(task.priority or 0.0),
        isReferenceStandard=bool(task.is_reference_standard),
        createdAt=task.created_at,
        updatedAt=task.updated_at,
    )


async def _stage_infos(session, dataset_id: uuid.UUID) -> list[StageInfo]:
    rows = (
        await session.execute(
            select(ImagehubDatasetStage.id, ImagehubDatasetStage.kind, ImagehubDatasetStage.seq).where(
                ImagehubDatasetStage.dataset_id == dataset_id
            )
        )
    ).all()
    return [StageInfo(id=str(sid), kind=kind or "label", seq=int(seq)) for sid, kind, seq in rows]


async def _load_task(session, dataset_id: uuid.UUID, task_id: str) -> ImagehubTask:
    try:
        tid = uuid.UUID(task_id)
    except (ValueError, TypeError):
        raise HTTPException(status_code=404, detail="Không tìm thấy công việc.")
    task = (
        await session.execute(
            select(ImagehubTask).where(ImagehubTask.id == tid, ImagehubTask.dataset_id == dataset_id)
        )
    ).scalar_one_or_none()
    if task is None:
        raise HTTPException(status_code=404, detail="Không tìm thấy công việc.")
    return task


async def _task_out_after(session, task: ImagehubTask) -> TaskOut:
    """Build a TaskOut for one task, looking up its file path / stage name / assignee name."""
    f = await session.get(ImagehubDatasetFile, task.dataset_file_id)
    file_path = f.logical_path if f else ""
    stage_name: Optional[str] = None
    if task.current_stage_id:
        st = await session.get(ImagehubDatasetStage, task.current_stage_id)
        stage_name = st.name if st else None
    assignee_name: Optional[str] = None
    if task.assignee_user_id:
        u = await session.get(User, task.assignee_user_id)
        if u is not None:
            assignee_name = _assignee_label(u.full_name, u.email)
    return _build_task_out(task, file_path, stage_name, assignee_name)


async def _task_detail_after(session, task: ImagehubTask) -> TaskDetailOut:
    """A single task as TaskDetailOut (TaskOut fields + its saved annotations)."""
    base = await _task_out_after(session, task)
    return TaskDetailOut(**base.model_dump(), annotations=task.annotations or [])


@router.post("/{dataset_id}/tasks/generate", response_model=GenerateTasksOut)
async def generate_tasks(
    dataset_id: str, authorization: Optional[str] = Header(None)
) -> GenerateTasksOut:
    """Owner or admin: create one task per image file that doesn't have one yet.

    Requires the dataset to have at least one pipeline stage (409 otherwise); tasks start at the
    first stage. Idempotent — files that already have a task are skipped.
    """
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds, _role = await _load_dataset_admin(session, dataset_id, uid, _is_admin(authorization))
        # All reads BEFORE any add (autoflush-before-commit guard; see add_stage).
        stages = await _stage_infos(session, ds.id)
        if not stages:
            raise HTTPException(
                status_code=409,
                detail="Hãy thêm giai đoạn quy trình trong Cài đặt trước khi tạo công việc.",
            )
        init = initial_transition(stages)
        actor = await _actor_name(session, uid)
        files = (
            await session.execute(
                select(ImagehubDatasetFile.id, ImagehubDatasetFile.logical_path).where(
                    ImagehubDatasetFile.dataset_id == ds.id,
                    ImagehubDatasetFile.file_kind == "image",
                )
            )
        ).all()
        existing = set(
            (
                await session.execute(
                    select(ImagehubTask.dataset_file_id).where(ImagehubTask.dataset_id == ds.id)
                )
            ).scalars().all()
        )
        start_stage = uuid.UUID(init.current_stage_id) if init.current_stage_id else None
        created = 0
        for fid, lpath in files:
            if fid in existing:
                continue
            session.add(
                ImagehubTask(
                    dataset_id=ds.id,
                    dataset_file_id=fid,
                    name=lpath or "",
                    current_stage_id=start_stage,
                    pipeline_state=init.pipeline_state,
                    queue_status="assigned",
                )
            )
            created += 1
        if created:
            await _write_audit(
                session, ds.id, uid, actor, _role_label(authorization),
                "Tạo công việc", f"{created} tác vụ",
            )
        await session.commit()
        return GenerateTasksOut(created=created, total=len(files))


@router.get("/{dataset_id}/tasks", response_model=list[TaskOut])
async def list_tasks(
    dataset_id: str,
    stage: Optional[str] = Query(None),
    status: Optional[str] = Query(None),
    state: Optional[str] = Query(None),
    reference: Optional[bool] = Query(None),
    mine: Optional[bool] = Query(None),
    authorization: Optional[str] = Header(None),
) -> list[TaskOut]:
    """Owner, member or admin: the dataset's tasks, highest priority first. Optional filters
    (mine=true → only tasks assigned to the caller)."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds = await _load_dataset_read(session, dataset_id, uid, _is_admin(authorization))
        q = (
            select(
                ImagehubTask,
                ImagehubDatasetFile.logical_path,
                ImagehubDatasetStage.name,
                User.full_name,
                User.email,
            )
            .join(ImagehubDatasetFile, ImagehubTask.dataset_file_id == ImagehubDatasetFile.id)
            .outerjoin(ImagehubDatasetStage, ImagehubTask.current_stage_id == ImagehubDatasetStage.id)
            .outerjoin(User, ImagehubTask.assignee_user_id == User.id)
            .where(ImagehubTask.dataset_id == ds.id)
        )
        if stage:
            try:
                q = q.where(ImagehubTask.current_stage_id == uuid.UUID(stage))
            except (ValueError, TypeError):
                raise HTTPException(status_code=422, detail="Giai đoạn không hợp lệ.")
        if status:
            q = q.where(ImagehubTask.queue_status == status)
        if state:
            q = q.where(ImagehubTask.pipeline_state == state)
        if reference is not None:
            q = q.where(ImagehubTask.is_reference_standard == reference)
        if mine:
            q = q.where(ImagehubTask.assignee_user_id == uid)
        q = q.order_by(ImagehubTask.priority.desc(), ImagehubTask.created_at.asc())
        rows = (await session.execute(q)).all()
        return [
            _build_task_out(t, lpath, sname, _assignee_label(fn, em))
            for (t, lpath, sname, fn, em) in rows
        ]


class AssignIn(BaseModel):
    userId: Optional[str] = None


@router.post("/{dataset_id}/tasks/{task_id}/assign", response_model=TaskOut)
async def assign_task(
    dataset_id: str,
    task_id: str,
    payload: Optional[AssignIn] = Body(None),
    authorization: Optional[str] = Header(None),
) -> TaskOut:
    """Assign a task. No body (or your own id) = claim it yourself; a project admin / owner may pass
    another member's userId to assign it to them (the target must be the owner or a member)."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds, role = await _load_dataset_any(session, dataset_id, uid, _is_admin(authorization))
        task = await _load_task(session, ds.id, task_id)
        target_uid = uid
        if payload and payload.userId:
            try:
                target_uid = uuid.UUID(payload.userId)
            except (ValueError, TypeError):
                raise HTTPException(status_code=422, detail="Người dùng không hợp lệ.")
        if target_uid != uid and role not in _PROJECT_ADMIN_ROLES:
            raise HTTPException(
                status_code=403, detail="Chỉ quản trị dự án mới gán việc cho người khác."
            )
        actor = await _actor_name(session, uid)
        # The assignee must be the dataset owner or an existing member.
        if target_uid != ds.owner_user_id:
            is_member = (
                await session.execute(
                    select(ImagehubDatasetMember.id).where(
                        ImagehubDatasetMember.dataset_id == ds.id,
                        ImagehubDatasetMember.user_id == target_uid,
                    )
                )
            ).scalar_one_or_none()
            if is_member is None:
                raise HTTPException(
                    status_code=422, detail="Người được gán không phải thành viên của bộ dữ liệu."
                )
        task.assignee_user_id = target_uid
        task.assignment_mode = "manual"
        task.updated_at = datetime.now(tz=timezone.utc)
        await _write_audit(
            session, ds.id, uid, actor, _role_label(authorization), "Gán việc", task.name,
        )
        await session.commit()
        await session.refresh(task)
        return await _task_out_after(session, task)


@router.post("/{dataset_id}/tasks/{task_id}/unassign", response_model=TaskOut)
async def unassign_task(
    dataset_id: str, task_id: str, authorization: Optional[str] = Header(None)
) -> TaskOut:
    """Owner/admin clear any task's assignee; a member may release their own task."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds, role = await _load_dataset_any(session, dataset_id, uid, _is_admin(authorization))
        task = await _load_task(session, ds.id, task_id)
        if not _can_work_task(role, task.assignee_user_id, uid):
            raise HTTPException(status_code=403, detail="Bạn không được phân công công việc này.")
        actor = await _actor_name(session, uid)
        task.assignee_user_id = None
        task.assignment_mode = "auto"
        task.updated_at = datetime.now(tz=timezone.utc)
        await _write_audit(
            session, ds.id, uid, actor, _role_label(authorization), "Bỏ gán việc", task.name,
        )
        await session.commit()
        await session.refresh(task)
        return await _task_out_after(session, task)


@router.post("/{dataset_id}/tasks/{task_id}/finalize", response_model=TaskOut)
async def finalize_task(
    dataset_id: str, task_id: str, authorization: Optional[str] = Header(None)
) -> TaskOut:
    """TP1: finalize a Label task — advance it to the next stage (or Ground Truth)."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds, role = await _load_dataset_any(session, dataset_id, uid, _is_admin(authorization))
        task = await _load_task(session, ds.id, task_id)
        if not _can_work_task(role, task.assignee_user_id, uid):
            raise HTTPException(status_code=403, detail="Bạn không được phân công công việc này.")
        stages = await _stage_infos(session, ds.id)
        actor = await _actor_name(session, uid)
        try:
            t = compute_finalize(
                task.pipeline_state,
                str(task.current_stage_id) if task.current_stage_id else None,
                stages,
            )
        except TaskPipelineError as exc:
            raise HTTPException(status_code=exc.status, detail=str(exc))
        task.pipeline_state = t.pipeline_state
        task.current_stage_id = uuid.UUID(t.current_stage_id) if t.current_stage_id else None
        task.queue_status = t.queue_status
        task.updated_at = datetime.now(tz=timezone.utc)
        await _write_audit(
            session, ds.id, uid, actor, _role_label(authorization), t.action, task.name,
        )
        await session.commit()
        await session.refresh(task)
        return await _task_out_after(session, task)


@router.post("/{dataset_id}/tasks/{task_id}/review", response_model=TaskOut)
async def review_task(
    dataset_id: str,
    task_id: str,
    payload: ReviewIn = Body(...),
    authorization: Optional[str] = Header(None),
) -> TaskOut:
    """TP2/TP3: accept (optionally with corrections) advances; reject returns to the first stage."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds, role = await _load_dataset_any(session, dataset_id, uid, _is_admin(authorization))
        task = await _load_task(session, ds.id, task_id)
        if not _can_work_task(role, task.assignee_user_id, uid):
            raise HTTPException(status_code=403, detail="Bạn không được phân công công việc này.")
        stages = await _stage_infos(session, ds.id)
        actor = await _actor_name(session, uid)
        try:
            t = compute_review(
                task.pipeline_state,
                str(task.current_stage_id) if task.current_stage_id else None,
                stages,
                payload.decision,
            )
        except TaskPipelineError as exc:
            raise HTTPException(status_code=exc.status, detail=str(exc))
        review_stage_id = task.current_stage_id  # the Review stage where the decision is made
        task.pipeline_state = t.pipeline_state
        task.current_stage_id = uuid.UUID(t.current_stage_id) if t.current_stage_id else None
        task.queue_status = t.queue_status
        task.updated_at = datetime.now(tz=timezone.utc)
        # Record the structured verdict (queryable history + per-reviewer counters + reject reason).
        session.add(
            ImagehubTaskReviewEvent(
                dataset_id=ds.id,
                task_id=task.id,
                stage_id=review_stage_id,
                reviewer_user_id=uid,
                decision=payload.decision,
                note=(payload.note or "").strip(),
            )
        )
        await _write_audit(
            session, ds.id, uid, actor, _role_label(authorization), t.action, task.name,
        )
        await session.commit()
        await session.refresh(task)
        return await _task_out_after(session, task)


@router.get("/{dataset_id}/review-stats", response_model=ReviewStatsOut)
async def review_stats(
    dataset_id: str,
    userId: Optional[str] = None,
    days: int = 30,
    authorization: Optional[str] = Header(None),
) -> ReviewStatsOut:
    """Accept/reject/corrections tallies over the last ``days`` (the productivity panel). Optional
    ``userId`` scopes to one reviewer; omitted = dataset-wide."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds, _role = await _load_dataset_any(session, dataset_id, uid, _is_admin(authorization))
        cutoff = datetime.now(tz=timezone.utc) - timedelta(days=max(1, min(days, 365)))
        conds = [
            ImagehubTaskReviewEvent.dataset_id == ds.id,
            ImagehubTaskReviewEvent.created_at >= cutoff,
        ]
        if userId:
            try:
                conds.append(ImagehubTaskReviewEvent.reviewer_user_id == uuid.UUID(userId))
            except (ValueError, TypeError):
                raise HTTPException(status_code=422, detail="Người dùng không hợp lệ.")
        rows = (
            await session.execute(
                select(ImagehubTaskReviewEvent.decision, func.count())
                .where(*conds)
                .group_by(ImagehubTaskReviewEvent.decision)
            )
        ).all()
        counts = {str(d): int(c) for d, c in rows}
        return ReviewStatsOut(
            accepted=counts.get("accept", 0),
            acceptWithCorrections=counts.get("acceptWithCorrections", 0),
            rejected=counts.get("reject", 0),
        )


@router.post("/{dataset_id}/tasks/{task_id}/skip", response_model=TaskOut)
async def skip_task(
    dataset_id: str, task_id: str, authorization: Optional[str] = Header(None)
) -> TaskOut:
    """Q2: send a task to the end of the queue (still owned)."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds, role = await _load_dataset_any(session, dataset_id, uid, _is_admin(authorization))
        task = await _load_task(session, ds.id, task_id)
        if not _can_work_task(role, task.assignee_user_id, uid):
            raise HTTPException(status_code=403, detail="Bạn không được phân công công việc này.")
        # Read the current max skip-seq BEFORE mutating this row (autoflush guard).
        max_seq = (
            await session.execute(
                select(func.max(ImagehubTask.skipped_seq)).where(ImagehubTask.dataset_id == ds.id)
            )
        ).scalar()
        actor = await _actor_name(session, uid)
        task.queue_status = "skipped"
        task.skipped_seq = (int(max_seq) + 1) if max_seq is not None else 0
        task.updated_at = datetime.now(tz=timezone.utc)
        await _write_audit(
            session, ds.id, uid, actor, _role_label(authorization), "Bỏ qua công việc", task.name,
        )
        await session.commit()
        await session.refresh(task)
        return await _task_out_after(session, task)


@router.post("/{dataset_id}/tasks/{task_id}/priority", response_model=TaskOut)
async def set_task_priority(
    dataset_id: str,
    task_id: str,
    payload: PriorityIn = Body(...),
    authorization: Optional[str] = Header(None),
) -> TaskOut:
    """PR4: set a task's priority (float 0..1; higher floats to the top of queues)."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds, _role = await _load_dataset_admin(session, dataset_id, uid, _is_admin(authorization))
        task = await _load_task(session, ds.id, task_id)
        actor = await _actor_name(session, uid)
        task.priority = float(payload.priority)
        task.updated_at = datetime.now(tz=timezone.utc)
        await _write_audit(
            session, ds.id, uid, actor, _role_label(authorization),
            "Đặt độ ưu tiên", f"{task.name} = {task.priority:.2f}",
        )
        await session.commit()
        await session.refresh(task)
        return await _task_out_after(session, task)


@router.post("/{dataset_id}/tasks/{task_id}/reference", response_model=TaskOut)
async def set_task_reference(
    dataset_id: str,
    task_id: str,
    payload: SetReferenceIn = Body(...),
    authorization: Optional[str] = Header(None),
) -> TaskOut:
    """RS1/RS2: set or unset a Ground Truth task as a project reference standard."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds, _role = await _load_dataset_admin(session, dataset_id, uid, _is_admin(authorization))
        task = await _load_task(session, ds.id, task_id)
        actor = await _actor_name(session, uid)
        try:
            validate_set_reference(task.pipeline_state, payload.isReferenceStandard)
        except TaskPipelineError as exc:
            raise HTTPException(status_code=exc.status, detail=str(exc))
        task.is_reference_standard = bool(payload.isReferenceStandard)
        task.updated_at = datetime.now(tz=timezone.utc)
        await _write_audit(
            session, ds.id, uid, actor, _role_label(authorization),
            "Đặt chuẩn tham chiếu" if task.is_reference_standard else "Bỏ chuẩn tham chiếu",
            task.name,
        )
        await session.commit()
        await session.refresh(task)
        return await _task_out_after(session, task)


@router.get("/{dataset_id}/tasks/{task_id}", response_model=TaskDetailOut)
async def get_task(
    dataset_id: str, task_id: str, authorization: Optional[str] = Header(None)
) -> TaskDetailOut:
    """Owner, member or admin: one task with its saved annotations (for the AnnotationTool)."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds = await _load_dataset_read(session, dataset_id, uid, _is_admin(authorization))
        task = await _load_task(session, ds.id, task_id)
        return await _task_detail_after(session, task)


@router.post("/{dataset_id}/tasks/{task_id}/save", response_model=TaskDetailOut)
async def save_task(
    dataset_id: str,
    task_id: str,
    payload: SaveIn = Body(...),
    authorization: Optional[str] = Header(None),
) -> TaskDetailOut:
    """Persist the labeler's annotations. Q3 save -> 'saved'; Q1 submit-draft -> 'pendingFinalization'
    (a draft stays in the queue until Finalize advances the stage)."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds, role = await _load_dataset_any(session, dataset_id, uid, _is_admin(authorization))
        task = await _load_task(session, ds.id, task_id)
        if not _can_work_task(role, task.assignee_user_id, uid):
            raise HTTPException(status_code=403, detail="Bạn không được phân công công việc này.")
        if task.pipeline_state == "groundTruth":
            raise HTTPException(status_code=409, detail="Tác vụ đã hoàn tất, không thể chỉnh sửa.")
        actor = await _actor_name(session, uid)
        task.annotations = payload.annotations
        task.queue_status = "pendingFinalization" if payload.submit else "saved"
        task.updated_at = datetime.now(tz=timezone.utc)
        await _write_audit(
            session, ds.id, uid, actor, _role_label(authorization),
            "Nộp bản nháp" if payload.submit else "Lưu chú thích", task.name,
        )
        await session.commit()
        await session.refresh(task)
        return await _task_detail_after(session, task)


# --------------------------------------------------------------------------- #
# Endpoints — dataset membership (multi-labeler). Owner + platform-admin manage; a member gets read
# access + may work tasks assigned to them (see _load_dataset_any / _can_work_task).
# --------------------------------------------------------------------------- #
_MEMBER_ROLES = ("member", "project_admin")


class MemberOut(BaseModel):
    userId: str
    email: str = ""
    fullName: str = ""
    role: str = "member"
    createdAt: Optional[datetime] = None


class MemberAddIn(BaseModel):
    email: Optional[str] = None
    userId: Optional[str] = None
    role: str = "member"


@router.get("/{dataset_id}/members", response_model=list[MemberOut])
async def list_members(
    dataset_id: str, authorization: Optional[str] = Header(None)
) -> list[MemberOut]:
    """Owner, admin or project-admin: the dataset's members."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds, _role = await _load_dataset_admin(session, dataset_id, uid, _is_admin(authorization))
        rows = (
            await session.execute(
                select(ImagehubDatasetMember, User.email, User.full_name)
                .join(User, User.id == ImagehubDatasetMember.user_id)
                .where(ImagehubDatasetMember.dataset_id == ds.id)
                .order_by(ImagehubDatasetMember.created_at)
            )
        ).all()
        return [
            MemberOut(
                userId=str(m.user_id),
                email=email or "",
                fullName=full_name or "",
                role=m.role,
                createdAt=m.created_at,
            )
            for m, email, full_name in rows
        ]


@router.post("/{dataset_id}/members", response_model=MemberOut)
async def add_member(
    dataset_id: str, payload: MemberAddIn = Body(...), authorization: Optional[str] = Header(None)
) -> MemberOut:
    """Owner / admin / project-admin: add a member by email or userId so they can be assigned tasks.
    Only the owner + platform-admin may grant the project_admin role (anti-escalation)."""
    _require_db()
    uid = _require_authed_uid(authorization)
    new_role = payload.role if payload.role in _MEMBER_ROLES else "member"
    async with get_session() as session:
        ds, caller_role = await _load_dataset_admin(session, dataset_id, uid, _is_admin(authorization))
        if new_role == "project_admin" and caller_role not in ("owner", "admin"):
            raise HTTPException(
                status_code=403, detail="Chỉ chủ sở hữu mới cấp quyền quản trị dự án."
            )
        # Resolve the target user + dup-check BEFORE add (autoflush-before-commit guard).
        target: Optional[User] = None
        if payload.userId:
            try:
                target = await session.get(User, uuid.UUID(payload.userId))
            except (ValueError, TypeError):
                target = None
        elif payload.email:
            target = (
                await session.execute(
                    select(User).where(func.lower(User.email) == payload.email.strip().lower())
                )
            ).scalar_one_or_none()
        if target is None:
            raise HTTPException(status_code=404, detail="Không tìm thấy người dùng.")
        if target.id == ds.owner_user_id:
            raise HTTPException(status_code=409, detail="Người dùng đã là chủ sở hữu bộ dữ liệu.")
        actor = await _actor_name(session, uid)
        existing = (
            await session.execute(
                select(ImagehubDatasetMember).where(
                    ImagehubDatasetMember.dataset_id == ds.id,
                    ImagehubDatasetMember.user_id == target.id,
                )
            )
        ).scalar_one_or_none()
        if existing is not None:
            raise HTTPException(status_code=409, detail="Người dùng đã là thành viên.")
        member = ImagehubDatasetMember(dataset_id=ds.id, user_id=target.id, role=new_role, added_by=uid)
        session.add(member)
        await _write_audit(
            session, ds.id, uid, actor, _role_label(authorization),
            "Thêm thành viên", target.email or str(target.id),
        )
        try:
            await session.commit()
        except IntegrityError:
            await session.rollback()
            raise HTTPException(status_code=409, detail="Người dùng đã là thành viên.")
        await session.refresh(member)
        return MemberOut(
            userId=str(target.id),
            email=target.email or "",
            fullName=target.full_name or "",
            role=new_role,
            createdAt=member.created_at,
        )


@router.delete("/{dataset_id}/members/{user_id}")
async def remove_member(
    dataset_id: str, user_id: str, authorization: Optional[str] = Header(None)
) -> dict[str, Any]:
    """Owner / admin / project-admin: remove a member. A project-admin may not remove another
    project-admin (only the owner + platform-admin can)."""
    _require_db()
    uid = _require_authed_uid(authorization)
    async with get_session() as session:
        ds, caller_role = await _load_dataset_admin(session, dataset_id, uid, _is_admin(authorization))
        try:
            tid = uuid.UUID(user_id)
        except (ValueError, TypeError):
            raise HTTPException(status_code=404, detail="Không tìm thấy thành viên.")
        member = (
            await session.execute(
                select(ImagehubDatasetMember).where(
                    ImagehubDatasetMember.dataset_id == ds.id,
                    ImagehubDatasetMember.user_id == tid,
                )
            )
        ).scalar_one_or_none()
        if member is None:
            raise HTTPException(status_code=404, detail="Không tìm thấy thành viên.")
        if member.role == "project_admin" and caller_role not in ("owner", "admin"):
            raise HTTPException(
                status_code=403, detail="Chỉ chủ sở hữu mới gỡ quyền quản trị dự án."
            )
        actor = await _actor_name(session, uid)
        await session.delete(member)
        await _write_audit(
            session, ds.id, uid, actor, _role_label(authorization), "Xóa thành viên", user_id,
        )
        await session.commit()
    return {"ok": True}
