penracourses/atlas/api.py

from collections import defaultdict
from typing import List, Tuple
from django.shortcuts import get_object_or_404
from django.urls import reverse
from ninja import ModelSchema, Router, Schema, Field

from django.db import transaction, IntegrityError

# from .decorators import user_is_author_or_rapid_checker
from django.core.exceptions import PermissionDenied
import pydicom
from pydicom.errors import InvalidDicomError

from generic.decorators import check_user_in_group
from generic.constants import Group

from ninja.security import django_auth

from ninja import NinjaAPI, File
from ninja.files import UploadedFile
import hashlib
import secrets
from django.utils import timezone

from generic.models import Examination, Modality

from .models import Case, DuplicateDicom, Series, SeriesImage, UncategorisedDicom, APIToken
from atlas.helpers import get_cases_available_to_user

from loguru import logger


router = Router()


class TokenAuth:
    """Simple bearer token auth for Ninja endpoints.

    Integrates with `APIToken` model; use as `auth=TokenAuth()` on router
    endpoints that should accept tokens in `Authorization: Bearer <token>`.
    """
    def __call__(self, request):
        # Extract header
        auth = request.META.get("HTTP_AUTHORIZATION")
        if not auth:
            return None
        parts = auth.split()
        if len(parts) != 2 or parts[0].lower() != "bearer":
            return None
        token = parts[1]
        token_hash = hashlib.sha256(token.encode()).hexdigest()
        try:
            api_token = APIToken.objects.get(token_hash=token_hash, revoked=False)
        except APIToken.DoesNotExist:
            return None

        if api_token.expires and api_token.expires < timezone.now():
            return None

        # mark last used (best-effort)
        try:
            api_token.mark_used()
        except Exception:
            pass

        # Attach token object to request for handlers
        request.api_token = api_token
        return api_token.user


class SeriesSchema(ModelSchema):
    case_id: List[int] = []

    class Meta:
        model = Series
        fields = [
            "id",
            "modality",
            "examination",
            "series_instance_uid",
            "author",
        ]

    @staticmethod
    def resolve_case_id(obj):
        return [i.id for i in obj.case.all()]


class CaseSchema(ModelSchema):
    class Meta:
        model = Case
        fields = ["id", "title"]


@router.post("/upload_dicom", auth=TokenAuth())
def upload_dicom(request, files: List[UploadedFile] = File(...)):
    uploaded = []
    duplicate = []
    failed = []
    duplicate_series = set()
    for file in files:
        # data = file.read()
        try:
            ud = UncategorisedDicom(image=file, user=request.user)

            ud.save()

            uploaded.append((file.name, ud.image_blake3_hash))
        except DuplicateDicom as e:
            duplicate.append((file.name, ud.image_blake3_hash))


            if type(e.duplicate) == SeriesImage:
                duplicate_series.add(e.duplicate.series.get_absolute_url())

            pass
        except InvalidDicomError:
            failed.append(file.name)

    return {
        "uploaded": uploaded,
        "duplicates": duplicate,
        "failed": failed,
        "duplicate_series": list(duplicate_series),
    }


@router.post("/generate_image_hash", auth=TokenAuth())
def generate_image_hash(request, id: int):
    s = SeriesImage.objects.get(pk=id)
    s.generate_hashes()

    return {
        "md5": s.image_md5_hash,
        "blake3": s.image_blake3_hash,
        "is_dicom": s.is_dicom,
    }


@router.post("/clear_dicoms", auth=TokenAuth())
def clear_dicoms(request):
    if "selection" in request.POST:
        dicoms = UncategorisedDicom.objects.filter(
            series_instance_uid__in=request.POST.getlist("selection")
        )

    else:
        dicoms = UncategorisedDicom.objects.filter(user=request.user)

    dicoms.delete()

    return True


@router.get("/uncategorised_dicoms", auth=TokenAuth())
def uncategorised_dicoms(request):
    dicoms = UncategorisedDicom.objects.filter(user=request.user)

    data = defaultdict(list)
    for d in dicoms:
        tags = d.get_basic_dicom_tags()

        data[tags["SeriesInstanceUID"]].append(tags)
        data["image_hash"].append(d.image_blake3_hash)

    return data


def import_dicoms_helper(request, case_id: int | None = None):
    if "selection" in request.POST:
        dicoms = UncategorisedDicom.objects.filter(
            series_instance_uid__in=request.POST.getlist("selection")
        )
    else:
        dicoms = UncategorisedDicom.objects.filter(user=request.user)

    if "order-series" in request.POST:
        order = request.POST["order-series"]
    else:
        order = None

    # Group dicoms by SeriesInstanceUID, but don't keep all tags in memory
    series_map = {}
    for d in dicoms.iterator():  # Use iterator() to avoid loading all at once
        tags = d.basic_dicom_tags
        series_uid = tags["SeriesInstanceUID"]

        if series_uid not in series_map:
            # Get or create Series and related objects as needed
            try:
                modality = Modality.objects.get(short_code=tags["Modality"])
            except Modality.DoesNotExist:
                modality = Modality.objects.create(short_code=tags["Modality"], modality=tags["Modality"])

            series_tags = tags  # Save the first tag for this series
            try:
                with transaction.atomic():
                    series, created = Series.objects.get_or_create(
                        series_instance_uid=series_uid,
                        defaults={
                            "modality": modality,
                            "description": tags.get("SeriesDescription", ""),
                        }
                    )
                    if created and tags.get("StudyDescription"):
                        examination, created_exam = Examination.objects.get_or_create(
                            examination=tags["StudyDescription"]
                        )
                        if created_exam:
                            examination.modality = modality
                            examination.save()
                        series.examination = examination
                        series.save()
            except IntegrityError:
                series = Series.objects.get(series_instance_uid=series_uid)

            # Add author and case if needed
            series.author.add(request.user)
            if case_id is not None:
                case_object = get_object_or_404(Case, pk=case_id)
                try:
                    series.case.add(case_object)
                except IntegrityError:
                    pass

            series_map[series_uid] = series
        else:
            series = series_map[series_uid]

        # Create and save SeriesImage immediately
        series_image = SeriesImage(
            image=d.image,
            series=series,
            is_dicom=True,
            image_blake3_hash=d.image_blake3_hash,
        )
        series_image.save()
        d.delete()  # Remove UncategorisedDicom immediately

    # Optionally, re-order series if needed
    for series in series_map.values():
        match order:
            case "order-series-instance-number":
                series.order_by_dicom("InstanceNumber")
            case "order-series-slice-location":
                series.order_by_dicom("SliceLocation")
            case _:
                pass

    return [(series, series.get_link()) for series in series_map.values()]


@router.post(
    "/import_dicoms", auth=TokenAuth(), response=List[Tuple[SeriesSchema, str]]
)
def import_dicoms(request):
    return import_dicoms_helper(request)


@router.post(
    "/import_dicoms/{case_id}",
    auth=TokenAuth(),
    response=List[Tuple[SeriesSchema, str]],
)
def import_dicoms_case(request, case_id: int):
    return import_dicoms_helper(request, case_id=case_id)


@router.post("/upload_dicom_case/{case_id}", auth=TokenAuth(), response=List[Tuple[SeriesSchema, str]])
def upload_dicom_case(request, case_id: int, files: List[UploadedFile] = File(...)):
    """Upload DICOM files and immediately import them into the given case.

    This saves each uploaded file as an `UncategorisedDicom` owned by the user,
    then calls the existing import helper which will create Series/SeriesImage
    objects and attach them to the specified case.
    """
    for file in files:
        try:
            ud = UncategorisedDicom(image=file, user=request.user)
            ud.save()
        except DuplicateDicom:
            # skip duplicates
            continue

    # Now reuse import helper which will consume UncategorisedDicom for this user
    return import_dicoms_helper(request, case_id=case_id)


@router.get("/orphan_series", auth=TokenAuth(), response=List[SeriesSchema])
def orphan_series(request):
    return request.user.series.filter(case=None)


@router.get("/series_remove_duplicate_images", auth=TokenAuth())
def series_remove_duplicate_images(request, series_id: int):
    series = get_object_or_404(Series, pk=series_id)

    img_ids = set()
    dupes = set()
    for series_image in series.images.all():
        if series_image.image_blake3_hash in img_ids:
            dupes.add(series_image.id)
        img_ids.add(series_image.image_blake3_hash)

    if dupes:
        SeriesImage.objects.filter(id__in=dupes).delete()

    return len(dupes)


@router.get("/get_cases_user", auth=TokenAuth(), response=List[CaseSchema])
def get_cases_user(request):
    return Case.objects.filter(author=request.user)


@router.get("/get_cases_available", auth=TokenAuth(), response=List[CaseSchema])
def get_cases_available(request):
    """Return cases available to the authenticated user (via get_cases_available_to_user)."""
    qs = get_cases_available_to_user(request.user)
    return qs.order_by('-created_date')[:200]


class APITokenOut(Schema):
    id: int
    name: str
    scopes: str = ""
    created: str
    expires: str | None = None
    revoked: bool = False
    last_used: str | None = None


@router.get("/api_tokens", auth=TokenAuth(), response=List[APITokenOut])
def list_api_tokens(request):
    tokens = APIToken.objects.filter(user=request.user).order_by("-created")
    out = []
    for t in tokens:
        out.append(
            {
                "id": t.id,
                "name": t.name,
                "scopes": t.scopes,
                "created": t.created.isoformat(),
                "expires": t.expires.isoformat() if t.expires else None,
                "revoked": t.revoked,
                "last_used": t.last_used.isoformat() if t.last_used else None,
            }
        )
    return out


class CreateTokenIn(Schema):
    name: str = Field("", description="Friendly name for token")
    scopes: str = Field("", description="Space-separated scopes")
    expires_days: int | None = Field(None, description="Expire after N days")


@router.post("/api_tokens", auth=TokenAuth())
def create_api_token(request, payload: CreateTokenIn):
    expires = None
    if payload.expires_days:
        expires = timezone.now() + timezone.timedelta(days=payload.expires_days)

    token, obj = APIToken.create_token(request.user, name=payload.name, scopes=payload.scopes, expires=expires)
    # Return raw token once — callers must store it
    return {"token": token, "id": obj.id}


@router.post("/api_tokens/{token_id}/revoke", auth=TokenAuth())
def revoke_api_token(request, token_id: int):
    try:
        t = APIToken.objects.get(pk=token_id, user=request.user)
    except APIToken.DoesNotExist:
        return {"status": "not found"}
    t.revoked = True
    t.save(update_fields=["revoked"])
    return {"status": "revoked"}


@router.get("/check_image_hash/{hash}", auth=TokenAuth())
def check_image_hash(request, hash: str):
    try:
        series_image = SeriesImage.objects.get(image_blake3_hash=hash)
        data = {
            "status": "success",
            "id": series_image.pk,
            "url": series_image.series.get_absolute_url(),
        }
        return data
    except SeriesImage.DoesNotExist:
        data = {"status": "success", "id": False}
        return data


@router.post("/check_image_hashes/", auth=TokenAuth())
def check_images_hashes(request, hashes: List[str]):
    """Checks a list of image hashes and returns the series id / url if found

    Return format
    { "hash_id": {"id": "series_id|false", "url": "series_url|false"}, ...}
    """

    hash_status = {}

    for hash in hashes:
        try:
            # TOOD also check against uncategorised dicoms?
            series_image = SeriesImage.objects.get(image_blake3_hash=hash)
            data = {
                "id": series_image.pk,
                "url": series_image.series.get_absolute_url(),
                "type": "series",
            }
        except SeriesImage.DoesNotExist:
            try:
                uncategorised_dicom = UncategorisedDicom.objects.get(
                    image_blake3_hash=hash
                )
                data = {
                    "id": uncategorised_dicom.pk,
                    "url": reverse("atlas:user_uploads"),
                    "type": "uncategorised",
                }
            except UncategorisedDicom.DoesNotExist:
                data = {"id": False, "url": False}

        hash_status[hash] = data

    return hash_status


# @router.get("/generate_image_hash/{id}", auth=django_auth)
# def generate_image_hash(request, id: int):
#        series_image = SeriesImage.objects.get(pk=id)
#
#        series_image.image_blake3_hash = ""
#
#        series_image.save()
#
#        print(series_image)


@router.get("/view_dicom_tags/{hash}", auth=TokenAuth())
def view_dicom_tags(request, hash: str):
    item = SeriesImage.objects.get(image_blake3_hash=hash)
    return item.get_dicom_json()


@router.get("/series_split_by_dicom_tag/{series_id}/{dicom_tag}", auth=TokenAuth())
def series_split_by_tag(request, series_id: int, dicom_tag: str):
    series = get_object_or_404(Series, pk=series_id)

    if not series.check_user_can_edit(request.user):
        return {"status": "permission denied"}

    if dicom_tag.startswith("(") and dicom_tag.endswith(")"):
        dicom_tag = tuple([hex(int(i.strip(), 16)) for i in dicom_tag[1:-1].split(",")])

    image_map = defaultdict(list)

    # Use iterator to avoid loading all images into memory
    for image in series.images.iterator():
        ds = image.get_dicom_data()
        if dicom_tag in ds:
            val = ds[dicom_tag].value
            if isinstance(val, pydicom.multival.MultiValue):
                key = ",".join(map(str, val))
            else:
                key = str(val)
            image_map[key].append(image)
        else:
            return [f"{dicom_tag} not found {image}"]

    case = list(series.case.all())
    authors = list(series.author.all())
    new_series = []

    # Use a transaction for all DB changes
    with transaction.atomic():
        for n, (option, images) in enumerate(image_map.items()):
            if n == 0:
                # Update the original series
                series.images.set(images)
                new_series.append(series.pk)
            else:
                # Clone the series
                series.pk = None
                series.save()
                series.case.set(case)
                series.author.set(authors)
                series.images.set(images)
                new_series.append(series.pk)

    return new_series


@router.get("/split_order_by_dicom_tag/{series_id}/{dicom_tag}", auth=TokenAuth())
def series_order_by_tag(request, series_id: int, dicom_tag: str):
    series = get_object_or_404(Series, pk=series_id)

    series.order_by_dicom(dicom_tag)

    return [f"{series} ordered by {dicom_tag}"]