"""ImageHub — organ-segmentation linking (Phase D).

A cohesive domain/service module for linking organ-mask files to the image they
segment. The HTTP route (``imagehub_routes.upload_segmentations``) stays a thin
transport layer; the domain rules live here, in one unit-testable place:

  * a mask may only attach to an existing *image* file in the *same* dataset;
  * masks are namespaced under their parent (``<parent-stem>.seg/<name>``) so they
    never collide with — or replace — a sibling image row under the dataset's
    UNIQUE (dataset_id, logical_path);
  * the link is recorded explicitly (``file_kind='segmentation'`` + ``parent_file_id``
    + ``organ_label``), not inferred.

Infrastructure (blob storage, imaging-metadata sniffing, filename safety) is
*injected*, not imported — so this module has no FastAPI/HTTP coupling, no circular
dependency on the router, and can be unit-tested with fakes. This is the project's
pragmatic take on Clean Architecture: a cohesive service over the flat routers, not
the (unwired) domain/application/infrastructure tree.
"""
from __future__ import annotations

import uuid
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any, Awaitable, Callable, Optional

from sqlalchemy import select

from src.initiative_db.models import ImagehubBlob, ImagehubDataset, ImagehubDatasetFile

# Injected infrastructure contracts (kept HTTP-free + testable).
PutBlob = Callable[[bytes, Optional[str]], Awaitable[dict[str, Any]]]
SniffMeta = Callable[[Optional[str], bytes, str], dict[str, Any]]
SafeName = Callable[[Optional[str]], str]

_IMAGING_EXTS = (".nii.gz", ".nii", ".dcm", ".dicom")


class SegmentationError(Exception):
    """A domain rule was violated; ``status`` is the HTTP code the caller should map to."""

    def __init__(self, message: str, status: int = 422) -> None:
        super().__init__(message)
        self.status = status


@dataclass(frozen=True)
class MaskUpload:
    """One organ-mask file to be linked to a parent image."""

    filename: str
    data: bytes
    media_type: str
    organ_label: str


class SegmentationService:
    """Links organ-mask files to a parent image within a single dataset."""

    def __init__(
        self, session, *, put_blob: PutBlob, sniff_meta: SniffMeta, safe_name: SafeName
    ) -> None:
        self._session = session
        self._put_blob = put_blob
        self._sniff_meta = sniff_meta
        self._safe_name = safe_name

    def _mask_logical_path(self, parent_logical_path: str, mask_filename: str) -> str:
        """``<parent-stem>.seg/<safe-mask-basename>`` — groups masks under their parent
        and keeps them out of the parent's logical-path namespace."""
        stem = parent_logical_path
        low = stem.lower()
        for ext in _IMAGING_EXTS:
            if low.endswith(ext):
                stem = stem[: -len(ext)]
                break
        return f"{stem}.seg/{self._safe_name(mask_filename)}"

    async def _resolve_parent(
        self, dataset: ImagehubDataset, parent_file_id: str
    ) -> ImagehubDatasetFile:
        try:
            pid = uuid.UUID(parent_file_id)
        except (ValueError, TypeError):
            raise SegmentationError("Không tìm thấy tệp ảnh gốc.", 404)
        parent = (
            await self._session.execute(
                select(ImagehubDatasetFile).where(
                    ImagehubDatasetFile.id == pid,
                    ImagehubDatasetFile.dataset_id == dataset.id,
                )
            )
        ).scalar_one_or_none()
        if parent is None:
            raise SegmentationError("Không tìm thấy tệp ảnh gốc.", 404)
        if parent.file_kind != "image":
            raise SegmentationError("Chỉ có thể gắn mặt nạ vào tệp ảnh gốc.", 422)
        return parent

    async def _ensure_blob(self, data: bytes, media_type: str) -> dict[str, Any]:
        blob = await self._put_blob(data, media_type)
        if await self._session.get(ImagehubBlob, blob["sha256"]) is None:
            self._session.add(
                ImagehubBlob(
                    sha256=blob["sha256"],
                    size_bytes=blob["size"],
                    media_type=media_type,
                    storage_bucket=blob["bucket"],
                    storage_key=blob["key"],
                )
            )
            await self._session.flush()
        return blob

    async def link_masks(
        self,
        dataset: ImagehubDataset,
        parent_file_id: str,
        masks: list[MaskUpload],
        actor_uid: uuid.UUID,
    ) -> list[ImagehubDatasetFile]:
        """Store + link each mask to the parent image. Returns the created/updated rows.

        Re-linking a mask whose namespaced path already exists replaces it in place
        (a corrected mask for the same organ). Raises ``SegmentationError`` on a bad
        parent or an empty payload; propagates the injected ``put_blob`` errors.
        """
        parent = await self._resolve_parent(dataset, parent_file_id)
        if not masks:
            raise SegmentationError("Chưa chọn tệp mặt nạ nào.", 422)

        linked: list[ImagehubDatasetFile] = []
        for m in masks:
            if not m.data:
                continue
            blob = await self._ensure_blob(m.data, m.media_type)
            logical_path = self._mask_logical_path(parent.logical_path, m.filename)
            meta = self._sniff_meta(m.filename, m.data, m.media_type)
            organ = (m.organ_label or "").strip() or self._safe_name(m.filename)
            row = (
                await self._session.execute(
                    select(ImagehubDatasetFile).where(
                        ImagehubDatasetFile.dataset_id == dataset.id,
                        ImagehubDatasetFile.logical_path == logical_path,
                    )
                )
            ).scalar_one_or_none()
            if row is None:
                row = ImagehubDatasetFile(
                    id=uuid.uuid4(), dataset_id=dataset.id, logical_path=logical_path
                )
                self._session.add(row)
            row.blob_sha256 = blob["sha256"]
            row.size_bytes = blob["size"]
            row.media_type = m.media_type
            row.imaging_meta = meta
            row.file_kind = "segmentation"
            row.parent_file_id = parent.id
            row.organ_label = organ
            row.uploaded_by = actor_uid
            row.updated_at = datetime.now(tz=timezone.utc)
            linked.append(row)

        if not linked:
            raise SegmentationError("Tệp mặt nạ rỗng.", 422)
        return linked
