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

from generic.models import Examination, Modality

from .models import Case, DuplicateDicom, Series, SeriesImage, UncategorisedDicom

from loguru import logger


router = Router()


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

    class Config:
        model = Series
        model_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 Config:
        model = Case
        model_fields = ["id", "title"]


@router.post("/upload_dicom", auth=django_auth)
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=django_auth)
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=django_auth)
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=django_auth)
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=django_auth, response=List[Tuple[SeriesSchema, str]]
)
def import_dicoms(request):
    return import_dicoms_helper(request)


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


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


@router.get("/series_remove_duplicate_images", auth=django_auth)
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("/series_truncate/{series_id}/{start}/{end}/", auth=django_auth)
def series_truncate(request, series_id: int, start: int, end: int):
    series = get_object_or_404(Series, pk=series_id)

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

    images_removed = []
    for n, image in enumerate(series.get_images()):
        if n >= start and n <= end:
            continue
        else:
            image.removed = True
            image.image = None
            image.save()
            images_removed.append(image.pk)
    
    series.get_total_image_size(refresh=True)

    return {"status": "success", "images_removed": images_removed}


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


@router.get("/check_image_hash/{hash}", auth=django_auth)
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=django_auth)
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=django_auth)
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=django_auth)
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=django_auth)
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}"]