From 823c54d7f01bdb14b0b82cd4f1ef86d8e2d5b45a Mon Sep 17 00:00:00 2001
From: Bryon Lewis <Bryon.Lewis@kitware.com>
Date: Mon, 8 Jan 2024 08:52:59 -0500
Subject: [PATCH] complete test annotation endpoints

---
 bats_ai/core/admin/__init__.py    |  2 ++
 bats_ai/core/admin/annotations.py | 21 +++++++++++++++++
 bats_ai/core/views/recording.py   | 38 +++++++++++++++++--------------
 3 files changed, 44 insertions(+), 17 deletions(-)
 create mode 100644 bats_ai/core/admin/annotations.py

diff --git a/bats_ai/core/admin/__init__.py b/bats_ai/core/admin/__init__.py
index e471619..d949e53 100644
--- a/bats_ai/core/admin/__init__.py
+++ b/bats_ai/core/admin/__init__.py
@@ -1,7 +1,9 @@
+from .annotations import AnnotationsAdmin
 from .image import ImageAdmin
 from .recording import RecordingAdmin
 
 __all__ = [
     'ImageAdmin',
     'RecordingAdmin',
+    'AnnotationsAdmin',
 ]
diff --git a/bats_ai/core/admin/annotations.py b/bats_ai/core/admin/annotations.py
new file mode 100644
index 0000000..34f1bfd
--- /dev/null
+++ b/bats_ai/core/admin/annotations.py
@@ -0,0 +1,21 @@
+from django.contrib import admin
+
+from bats_ai.core.models import Annotations
+
+
+@admin.register(Annotations)
+class AnnotationsAdmin(admin.ModelAdmin):
+    list_display = [
+        'pk',
+        'recording',
+        'owner',
+        'start_time',
+        'end_time',
+        'low_freq',
+        'high_freq',
+        'comments',
+    ]
+    list_select_related = True
+    # list_select_related = ['owner']
+    filter_horizontal = ('species',)  # or filter_vertical
+    autocomplete_fields = ['owner']
diff --git a/bats_ai/core/views/recording.py b/bats_ai/core/views/recording.py
index 17f5959..6578fad 100644
--- a/bats_ai/core/views/recording.py
+++ b/bats_ai/core/views/recording.py
@@ -47,20 +47,24 @@ class AnnotationSchema(Schema):
     comments: str
 
 
-def get_owner_id(request: HttpRequest):
-    token = request.headers.get('Authorization').replace('Bearer ', '')
-    token_found = AccessToken.objects.get(token=token)
-    if not token_found:
-        raise HttpError(401, 'Authentication credentials were not provided.')
-
-    return token_found.user.pk
+def get_user(request: HttpRequest):
+    auth_header = request.headers.get('Authorization', None)
+    if auth_header is not None:
+        token = request.headers.get('Authorization').replace('Bearer ', '')
+        token_found = AccessToken.objects.get(token=token)
+        if not token_found:
+            raise HttpError(401, 'Authentication credentials were not provided.')
+        return token_found.user
+    elif request.user:
+        logger.warning(f'User: {request.user}')
+        return request.user
 
 
 @router.post('/')
 def create_recording(
     request: HttpRequest, payload: Form[RecordingUploadSchema], audio_file: File[UploadedFile]
 ):
-    user_id = get_owner_id(request)
+    user_id = get_user(request).pk
     converted_date = datetime.strptime(payload.recorded_date, '%Y-%m-%d')
     recording = Recording(
         name=payload.name,
@@ -78,7 +82,7 @@ def create_recording(
 @router.get('/')
 def get_recordings(request: HttpRequest):
     # Check if the user is authenticated and get userId
-    user_id = get_owner_id(request)
+    user_id = get_user(request)
 
     # Filter recordings based on the owner's id
     recordings = Recording.objects.filter(owner=user_id).values()
@@ -106,7 +110,7 @@ def get_spectrogram(request: HttpRequest, id: int):
 
 @router.get('/{id}/annotations')
 def get_annotations(request: HttpRequest, id: int):
-    user_id = get_owner_id(request)
+    user_id = get_user(request)
 
     try:
         recording = Recording.objects.get(pk=id, owner=user_id)
@@ -128,20 +132,20 @@ def get_annotations(request: HttpRequest, id: int):
 def put_annotation(
     request,
     id: int,
-    annotation: Form[AnnotationSchema],
+    annotation: AnnotationSchema,
     species_ids: list[int],
 ):
-    user_id = get_owner_id(request)
+    user = get_user(request)
 
     try:
-        recording = Recording.objects.get(pk=id, owner=user_id)
+        recording = Recording.objects.get(pk=id, owner=user.pk)
     except Recording.DoesNotExist:
         return {'error': 'Recording not found'}
 
     # Create a new annotation
     new_annotation = Annotations.objects.create(
         recording=recording,
-        owner=request.auth.user,  # Assuming request.auth.user contains the authenticated user
+        owner=user,
         start_time=annotation.start_time,
         end_time=annotation.end_time,
         low_freq=annotation.low_freq,
@@ -166,10 +170,10 @@ def patch_annotation(
     request,
     recording_id: int,
     id: int,
-    annotation: Form[AnnotationSchema],
+    annotation: AnnotationSchema,
     species_ids: list[int],
 ):
-    user_id = get_owner_id(request)
+    user_id = get_user(request)
 
     try:
         recording = Recording.objects.get(pk=recording_id, owner=user_id)
@@ -204,7 +208,7 @@ def patch_annotation(
 
 @router.delete('/{recording_id}/annotations/{id}')
 def delete_annotation(request, recording_id: int, id: int):
-    user_id = get_owner_id(request)
+    user_id = get_user(request)
 
     try:
         recording = Recording.objects.get(pk=recording_id, owner=user_id)