from collections import defaultdict
from typing import List, Tuple
from django.shortcuts import get_object_or_404
from django.urls import reverse
import json
from django.http import JsonResponse, HttpResponse
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, HttpBearer
from ninja.errors import HttpError

from ninja import NinjaAPI, File
from ninja.files import UploadedFile
import hashlib
import secrets
from django.utils import timezone
from django.contrib.auth import authenticate
from django.views.decorators.csrf import csrf_exempt
from django.middleware.csrf import CsrfViewMiddleware

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 BearerAuth(HttpBearer):
    """Ninja HttpBearer auth that validates APIToken entries.

    Returns the token owner user on success or raises HttpError(401) on failure.
    """
    def authenticate(self, request, token: str):
        if not token:
            raise HttpError(status_code=401)

        token_hash = hashlib.sha256(token.encode()).hexdigest()
        try:
            api_token = APIToken.objects.get(token_hash=token_hash)
        except APIToken.DoesNotExist:
            raise HttpError(status_code=401)

        if api_token.revoked:
            raise HttpError(status_code=401)

        if api_token.expires and api_token.expires < timezone.now():
            raise HttpError(status_code=401)

        try:
            api_token.mark_used()
        except Exception:
            pass

        # Attach for downstream handlers
        request.api_token = api_token
        # Also set `request.user` so code that expects Django's user works
        try:
            request.user = api_token.user
            # Some middleware caches the user on `_cached_user`
            setattr(request, "_cached_user", api_token.user)
        except Exception:
            pass

        return api_token.user


class SessionOrBearer:
    """Composite auth: prefer Django session (request.user), fall back to bearer token.

    Use as `auth=SessionOrBearer()` on Ninja router endpoints.
    """
    def __call__(self, request):
        # Session-backed auth already populated by Django middleware
        if getattr(request, "user", None) and getattr(request.user, "is_authenticated", False):
            return request.user

        # Try bearer token from header
        auth = request.META.get("HTTP_AUTHORIZATION")
        if not auth:
            raise HttpError(status_code=401)
        parts = auth.split()
        if len(parts) != 2 or parts[0].lower() != "bearer":
            raise HttpError(status_code=401)
        token = parts[1]
        # reuse BearerAuth logic
        bearer = BearerAuth()
        return bearer.authenticate(request, token)


# Backwards-compatible alias used in existing decorators
TokenAuth = SessionOrBearer


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, dicoms=None):
    # Allow callers to pass an explicit queryset of UncategorisedDicom objects
    if dicoms is 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
            # Attach author to the series
            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())
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.
    """
    uploaded = []
    duplicate = []
    failed = []
    duplicate_series = set()

    created_ids = []

    for file in files:
        try:
            ud = UncategorisedDicom(image=file, user=request.user)
            ud.save()

            uploaded.append((file.name, ud.image_blake3_hash))
            created_ids.append(ud.pk)
        except DuplicateDicom as e:
            duplicate.append((file.name, getattr(ud, "image_blake3_hash", None)))

            dup_obj = getattr(e, "duplicate", None)
            if isinstance(dup_obj, SeriesImage):
                try:
                    if dup_obj.series is not None:
                        duplicate_series.add(dup_obj.series.get_absolute_url())
                except Exception:
                    pass

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

    # Import into the target case (consume only the uploaded UncategorisedDicom objects)
    if created_ids:
        try:
            dicom_qs = UncategorisedDicom.objects.filter(pk__in=created_ids)
            import_dicoms_helper(request, case_id=case_id, dicoms=dicom_qs)
        except Exception:
            # don't let import errors prevent returning upload status
            logger.error("Error importing dicoms after upload", exc_info=True)

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


@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)."""
    logger.error(f"Getting cases available to user {request.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")


class TokenAuthIn(Schema):
    username: str
    password: str


@csrf_exempt
@router.post("/create_api_token")
def token_auth(request, payload: TokenAuthIn):
    """Authenticate username/password and return a new API token.

    Returns JSON: {"token": "<raw-token>", "id": <token_id>}
    """
    user = authenticate(request, username=payload.username, password=payload.password)
    if user is None:
        return HttpResponse(status=401, content=json.dumps({"detail": "Invalid credentials"}), content_type="application/json")

    # Create a new token for this login (no scopes by default)
    token, obj = APIToken.create_token(user, name=f"login-{timezone.now().isoformat()}")

    return {"token": token, "id": obj.id, "username": user.username}


@csrf_exempt
@router.post("/token_check")
def token_check(request):
    """Check whether a token is valid.

    Accepts either `Authorization: Bearer <token>` header or POST JSON
    body `{ "token": "<token>" }`.

    Returns JSON with keys: `valid` (bool), and when valid: `id`, `username`,
    `scopes`, `expires`, `revoked`.
    """
    auth = request.META.get("HTTP_AUTHORIZATION")
    token_obj = None

    # If header present, validate bearer token explicitly
    if auth:
        parts = auth.split()
        if len(parts) != 2 or parts[0].lower() != "bearer":
            return JsonResponse({"valid": False}, status=401)
        token = parts[1]
        try:
            bearer = BearerAuth()
            user = bearer.authenticate(request, token)
            token_obj = getattr(request, "api_token", None)
        except HttpError:
            return JsonResponse({"valid": False}, status=401)
    else:
        # try JSON body
        try:
            data = json.loads(request.body.decode("utf-8")) if request.body else {}
        except Exception:
            data = {}
        token = data.get("token")
        if token:
            token_hash = hashlib.sha256(token.encode()).hexdigest()
            try:
                t = APIToken.objects.get(token_hash=token_hash)
                token_obj = t
            except APIToken.DoesNotExist:
                token_obj = None

    if not token_obj:
        return JsonResponse({"valid": False}, status=401)

    # Validate token state
    valid = True
    if token_obj.revoked:
        valid = False
    if token_obj.expires and token_obj.expires < timezone.now():
        valid = False

    if not valid:
        return JsonResponse({"valid": False, "revoked": token_obj.revoked}, status=401)

    return JsonResponse(
        {
            "valid": True,
            "id": token_obj.id,
            "username": token_obj.user.username,
            "scopes": token_obj.scopes,
            "expires": token_obj.expires.isoformat() if token_obj.expires else None,
            "revoked": token_obj.revoked,
        },
        status=200,
    )


#@csrf_exempt
#@router.post("/create_api_token")
#def create_api_token(request, payload: CreateTokenIn):
#    """Create an API token.
#
#    This endpoint accepts either normal session auth (`request.user`) or
#    a bearer token. If no authenticated user is found a 403 is returned.
#    """
#    # If caller is not using Authorization header, enforce CSRF via middleware
#    if not request.META.get("HTTP_AUTHORIZATION"):
#        mw = CsrfViewMiddleware()
#        resp = mw.process_view(request, create_api_token, (), {})
#        if resp is not None:
#            return resp
#    # Determine authenticated user: prefer session auth, fall back to token auth
#    user = None
#    if getattr(request, "user", None) and getattr(request.user, "is_authenticated", False):
#        user = request.user
#    else:
#        user = TokenAuth()(request)
#
#    if not user:
#        raise PermissionDenied()
#
#    expires = None
#    if payload.expires_days:
#        expires = timezone.now() + timezone.timedelta(days=payload.expires_days)
#
#    token, obj = APIToken.create_token(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"}, ...}
    """

    try:
        hash_status = {}

        for hash in hashes:
            # Prefer filter().first() to avoid MultipleObjectsReturned
            series_image = SeriesImage.objects.filter(image_blake3_hash=hash).select_related("series").first()
            if series_image:
                try:
                    url = series_image.series.get_absolute_url()
                except Exception:
                    url = None
                data = {"id": series_image.pk, "url": url, "type": "series"}
                hash_status[hash] = data
                continue

            uncategorised = UncategorisedDicom.objects.filter(image_blake3_hash=hash).first()
            if uncategorised:
                try:
                    uploads_url = reverse("atlas:user_uploads")
                except Exception:
                    uploads_url = None
                data = {"id": uncategorised.pk, "url": uploads_url, "type": "uncategorised"}
                hash_status[hash] = data
                continue

            hash_status[hash] = {"id": False, "url": False}

        return hash_status
    except Exception as e:
        logger.exception("check_images_hashes failed")
        return JsonResponse({"detail": "Internal server error"}, status=500)


# @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}"]