from time import sleep
from django.core.mail import send_mail
from django.http import HttpResponse
from django.shortcuts import get_object_or_404
from enum import Enum

try:
    from django_tasks import task
    HAS_DJANGO_TASKS = True
except ImportError:
    from celery import shared_task
    from celery.result import AsyncResult

    HAS_DJANGO_TASKS = False

    class _CompatTaskResultStatus(Enum):
        READY = "READY"
        RUNNING = "RUNNING"
        FAILED = "FAILED"
        SUCCESSFUL = "SUCCESSFUL"

    class _CompatError:
        def __init__(self, traceback):
            self.traceback = traceback

    class _CompatTaskResult:
        def __init__(self, async_result):
            self._async_result = async_result

        @property
        def id(self):
            return self._async_result.id

        @property
        def status(self):
            state = self._async_result.state
            if state in ("PENDING",):
                return _CompatTaskResultStatus.READY
            if state in ("STARTED", "RETRY"):
                return _CompatTaskResultStatus.RUNNING
            if state == "SUCCESS":
                return _CompatTaskResultStatus.SUCCESSFUL
            if state == "FAILURE":
                return _CompatTaskResultStatus.FAILED
            return _CompatTaskResultStatus.RUNNING

        def refresh(self):
            return self

        @property
        def return_value(self):
            return self._async_result.result

        @property
        def errors(self):
            if self._async_result.state == "FAILURE" and self._async_result.traceback:
                return [_CompatError(self._async_result.traceback)]
            return []

    def task(func=None, *, takes_context=False, **kwargs):
        def decorator(inner_func):
            if takes_context:
                @shared_task(bind=True)
                def wrapped(self, *args, **inner_kwargs):
                    class _Ctx:
                        class _TaskResultRef:
                            id = self.request.id

                        task_result = _TaskResultRef()
                        attempt = 1

                    return inner_func(_Ctx(), *args, **inner_kwargs)
            else:
                wrapped = shared_task(inner_func)

            wrapped.enqueue = wrapped.delay
            wrapped.get_result = lambda result_id: _CompatTaskResult(AsyncResult(result_id))
            return wrapped

        if func is not None:
            return decorator(func)
        return decorator

from atlas.models import Case, Series, SeriesImage
from generic.models import CimarCase
from rad.settings import REMOTE_URL, CIMAR_USERNAME, CIMAR_PASSWORD
from helpers.cimar import CimarAPI, NotFoundError
from pydicom.uid import generate_uid
from django.contrib.auth.models import User
from django.core.files.base import ContentFile
from django.core.cache import cache
import copy
import io


@task(takes_context=True)
def series_downsample_task(context, series_id, user_id, downsample_pct):
    """Downsample a series asynchronously while preserving hash continuity."""
    from loguru import logger
    from atlas import views as atlas_views

    progress_key = f"series_downsample_progress:{context.task_result.id}"
    cache.set(
        progress_key,
        {"current": 0, "total": 0, "message": "Preparing downsample..."},
        timeout=60 * 60,
    )

    series = get_object_or_404(Series, pk=series_id)
    user = get_object_or_404(User, pk=user_id)

    if not series.check_user_can_edit(user):
        raise PermissionError("Permission denied")

    images = list(series.get_images())
    total_images = len(images)
    replacement_series_uid = generate_uid()
    source_series_uid = series.series_instance_uid
    updated_images = []

    cache.set(
        progress_key,
        {"current": 0, "total": total_images, "message": "Downsample started"},
        timeout=60 * 60,
    )

    for idx, image in enumerate(images, start=1):
        ds = atlas_views._read_series_image_dataset(image)
        if ds is None:
            cache.set(
                progress_key,
                {"current": idx, "total": total_images, "message": f"Skipping unreadable image ({idx}/{total_images})"},
                timeout=60 * 60,
            )
            continue

        try:
            original_hash = image.image_blake3_hash
            ds_out = atlas_views._downsample_dicom_dataset(
                ds,
                int(downsample_pct),
                series_instance_uid=replacement_series_uid,
                source_series_instance_uid=source_series_uid,
                series_description=series.description,
            )
            out_io = io.BytesIO()
            ds_out.save_as(out_io, write_like_original=False)
            out_io.seek(0)

            image.image.save(
                f"downsample_{downsample_pct}_{image.pk}.dcm",
                ContentFile(out_io.getvalue()),
                save=False,
            )
            image.image_blake3_hash = original_hash
            image.save(update_fields=["image", "image_blake3_hash"])
            updated_images.append(image.pk)
        except Exception as exc:
            logger.warning("Downsample failed for image {}: {}", image.pk, exc)

        cache.set(
            progress_key,
            {
                "current": idx,
                "total": total_images,
                "message": f"Downsampling image {idx}/{total_images}",
            },
            timeout=60 * 60,
        )

    if updated_images:
        series.modified = Series.SeriesModifiedChocies.RE
        series.source_series_instance_uid = source_series_uid
        series.series_instance_uid = replacement_series_uid
        series.save(update_fields=["modified", "source_series_instance_uid", "series_instance_uid"])
        try:
            series.get_total_image_size(refresh=True)
        except Exception as exc:
            logger.warning("Failed to refresh total_image_size for series {}: {}", series.pk, exc)

    cache.set(
        progress_key,
        {
            "current": total_images,
            "total": total_images,
            "message": "Downsample complete",
        },
        timeout=10 * 60,
    )

    return {
        "series_id": series.pk,
        "downsample_pct": int(downsample_pct),
        "updated_images": len(updated_images),
        "series_instance_uid": replacement_series_uid,
        "source_series_instance_uid": source_series_uid,
    }

@task
def push_case_to_cimar_task(case_id):
    """Sends an email when the feedback form has been submitted."""
    case = get_object_or_404(Case, pk=case_id)

    api = CimarAPI()
    api.login(username=CIMAR_USERNAME, password=CIMAR_PASSWORD)

    # We use the same study_uid for all images in a case
    study_uid = generate_uid()

    for series in case.get_series():
        print(f"Upload series: {series}")
        for image in series.images.filter(removed=False):
            data = api.upload_dicom(image.image.path, study_uid=study_uid)

    
    retries = 5
    delay = 20
    for attempt in range(retries):
        print("attempt 1")
        try:
            cimar_uuid = api.get_study_by_study_uid(data["study_uid"])["uuid"]
            break
        except NotFoundError:
            if attempt < retries - 1:
                sleep(delay)
                delay *= 2
            else:
                cimar_uuid = api.get_study_by_study_uid(data["study_uid"])["uuid"]


    case.cimar_uuid = cimar_uuid

    case.save()

    cimar_case, created = CimarCase.objects.get_or_create(uuid=cimar_uuid)

    cimar_case.refresh_study()

    return 10


@task(takes_context=True)
def series_reconstruct_task(
    context,
    series_id,
    user_id,
    recon_planes,
    slice_thickness_val=None,
    slice_spacing_val=None,
    recon_thickness_mode="mean",
):
    """Generate reconstructions asynchronously for a series with progress updates."""
    import numpy as np
    from loguru import logger
    from atlas import views as atlas_views

    progress_key = f"series_reconstruct_progress:{context.task_result.id}"
    cache.set(
        progress_key,
        {"current": 0, "total": 0, "message": "Preparing reconstruction..."},
        timeout=60 * 60,
    )

    series = get_object_or_404(Series, pk=series_id)
    user = get_object_or_404(User, pk=user_id)

    if not series.check_user_can_edit(user):
        raise PermissionError("Permission denied")

    images = list(series.get_images())
    dicom_items = []
    for image in images:
        ds = atlas_views._read_series_image_dataset(image)
        if ds is None:
            continue
        try:
            arr = ds.pixel_array
            if arr.ndim != 2:
                continue
            dicom_items.append((image, ds, arr))
        except Exception:
            continue

    if len(dicom_items) < 2:
        raise ValueError("Need at least 2 valid DICOM images in series for reconstruction")

    base_shape = dicom_items[0][2].shape
    dicom_items = [item for item in dicom_items if item[2].shape == base_shape]
    if len(dicom_items) < 2:
        raise ValueError("Not enough consistently-sized slices for reconstruction")

    geom = atlas_views._extract_recon_geometry(dicom_items)
    sorted_items = geom["sorted_items"]
    volume = np.stack([item[2] for item in sorted_items], axis=0)

    template_ds = sorted_items[0][1]
    origin_ipp = geom["origin_ipp"]
    row_dir = geom["row_dir"]
    col_dir = geom["col_dir"]
    normal_dir = geom["normal_dir"]
    native_row_spacing = geom["row_spacing"]
    native_col_spacing = geom["col_spacing"]
    native_z_spacing = geom["native_z_spacing"]

    requested_slice_spacing = float(slice_spacing_val) if slice_spacing_val is not None else None
    requested_slab_thickness = float(slice_thickness_val) if slice_thickness_val is not None else None

    # Determine source acquisition plane from normal vector.
    dominant_normal_axis = int(np.argmax(np.abs(normal_dir)))
    source_plane = {0: "sagittal", 1: "coronal", 2: "axial"}.get(dominant_normal_axis, "unknown")

    # Reorient the reconstructed volume into patient axes order [z, y, x] so
    # generated plane names are anatomically correct regardless of source plane.
    # NOTE: DICOM stores ImageOrientationPatient as:
    # - first triplet: direction across columns (x-axis in pixel grid)
    # - second triplet: direction across rows (y-axis in pixel grid)
    # In _extract_recon_geometry names are historically row_dir/col_dir, so we
    # map axis spacing/indexing explicitly to avoid directional mixups.
    source_axis_dirs = [normal_dir, col_dir, row_dir]
    source_axis_spacings = [float(native_z_spacing), float(native_row_spacing), float(native_col_spacing)]
    source_axis_sizes = [int(volume.shape[0]), int(volume.shape[1]), int(volume.shape[2])]

    source_for_patient_axis = {}
    sign_for_patient_axis = {}
    patient_axes_claimed = set()

    for src_axis, vec in enumerate(source_axis_dirs):
        patient_axis = int(np.argmax(np.abs(vec)))
        if patient_axis in patient_axes_claimed:
            raise ValueError("Could not infer unique source orientation axes for reconstruction")
        patient_axes_claimed.add(patient_axis)
        source_for_patient_axis[patient_axis] = src_axis
        sign_for_patient_axis[patient_axis] = 1 if float(vec[patient_axis]) >= 0 else -1

    src_x = source_for_patient_axis[0]
    src_y = source_for_patient_axis[1]
    src_z = source_for_patient_axis[2]

    # Build patient-axis basis vectors expressed in source/world space.
    # This preserves source world-frame orientation (including obliquity/sign)
    # for derived series instead of hardcoding cardinal axes.
    patient_x_vec = np.asarray(source_axis_dirs[src_x], dtype=float) * float(sign_for_patient_axis[0])
    patient_y_vec = np.asarray(source_axis_dirs[src_y], dtype=float) * float(sign_for_patient_axis[1])
    patient_z_vec = np.asarray(source_axis_dirs[src_z], dtype=float) * float(sign_for_patient_axis[2])

    def _norm_vec(vec):
        nrm = float(np.linalg.norm(vec))
        if nrm == 0:
            return np.asarray(vec, dtype=float)
        return np.asarray(vec, dtype=float) / nrm

    patient_x_vec = _norm_vec(patient_x_vec)
    patient_y_vec = _norm_vec(patient_y_vec)
    patient_z_vec = _norm_vec(patient_z_vec)

    def _orientation_pair(col_vec, row_vec, desired_normal):
        col = _norm_vec(col_vec)
        row = _norm_vec(row_vec)
        desired = _norm_vec(desired_normal)
        normal = np.cross(col, row)
        if float(np.dot(normal, desired)) < 0:
            row = -row
        return [float(col[0]), float(col[1]), float(col[2]), float(row[0]), float(row[1]), float(row[2])]

    spacing_x = source_axis_spacings[src_x]
    spacing_y = source_axis_spacings[src_y]
    spacing_z = source_axis_spacings[src_z]

    vol_zyx = np.transpose(volume, (src_z, src_y, src_x))

    if sign_for_patient_axis[2] < 0:
        vol_zyx = np.flip(vol_zyx, axis=0)
    if sign_for_patient_axis[1] < 0:
        vol_zyx = np.flip(vol_zyx, axis=1)
    if sign_for_patient_axis[0] < 0:
        vol_zyx = np.flip(vol_zyx, axis=2)

    def source_indices_from_patient_zyx(iz, iy, ix):
        src_indices = [0.0, 0.0, 0.0]

        z_index = float(iz)
        y_index = float(iy)
        x_index = float(ix)

        if sign_for_patient_axis[2] < 0:
            z_index = source_axis_sizes[src_z] - 1 - z_index
        if sign_for_patient_axis[1] < 0:
            y_index = source_axis_sizes[src_y] - 1 - y_index
        if sign_for_patient_axis[0] < 0:
            x_index = source_axis_sizes[src_x] - 1 - x_index

        src_indices[src_z] = z_index
        src_indices[src_y] = y_index
        src_indices[src_x] = x_index
        return src_indices

    def position_from_patient_zyx(iz, iy, ix):
        src_k, src_r, src_c = source_indices_from_patient_zyx(iz, iy, ix)
        pos = (
            np.asarray(origin_ipp, dtype=float)
            + (np.asarray(normal_dir, dtype=float) * (float(src_k) * float(native_z_spacing)))
            + (np.asarray(col_dir, dtype=float) * (float(src_r) * float(native_row_spacing)))
            + (np.asarray(row_dir, dtype=float) * (float(src_c) * float(native_col_spacing)))
        )
        return [float(pos[0]), float(pos[1]), float(pos[2])]

    def _reduce_slab(slab, reduce_axis=0):
        if recon_thickness_mode == "max":
            out = np.max(slab, axis=reduce_axis)
        elif recon_thickness_mode == "min":
            out = np.min(slab, axis=reduce_axis)
        else:
            out = np.mean(slab, axis=reduce_axis)

        if np.issubdtype(volume.dtype, np.integer):
            info = np.iinfo(volume.dtype)
            out = np.clip(np.rint(out), info.min, info.max).astype(volume.dtype)
        else:
            out = out.astype(volume.dtype, copy=False)
        return out

    def _slice_indices_for_axis(axis_length, axis_spacing, output_spacing, slab_thickness):
        if output_spacing <= 0:
            raise ValueError("Slice spacing must be greater than 0")
        if slab_thickness <= 0:
            raise ValueError("Slice thickness must be greater than 0")

        axis_positions_mm = np.arange(axis_length, dtype=float) * float(axis_spacing)
        end_pos = axis_positions_mm[-1] if axis_positions_mm.size else 0.0
        center_positions_mm = np.arange(0.0, end_pos + 1e-6, output_spacing, dtype=float)
        if center_positions_mm.size == 0:
            center_positions_mm = np.array([0.0], dtype=float)

        slabs = []
        centers_idx = []
        half_thickness = slab_thickness / 2.0
        for center_mm in center_positions_mm:
            mask = np.abs(axis_positions_mm - center_mm) <= (half_thickness + 1e-6)
            if not mask.any():
                nearest = int(np.argmin(np.abs(axis_positions_mm - center_mm)))
                idxs = np.array([nearest], dtype=int)
            else:
                idxs = np.where(mask)[0].astype(int)

            slabs.append(idxs)
            centers_idx.append(float(center_mm / axis_spacing) if axis_spacing > 0 else float(idxs[0]))

        return slabs, centers_idx

    plane_builders = {}
    for plane in recon_planes:
        plane_norm = plane.lower()
        if plane_norm == "axial":
            plane_spacing = float(requested_slice_spacing) if requested_slice_spacing is not None else float(spacing_z)
            plane_thickness = float(requested_slab_thickness) if requested_slab_thickness is not None else float(plane_spacing)
            slab_indices, slab_centers_idx = _slice_indices_for_axis(
                int(vol_zyx.shape[0]), float(spacing_z), plane_spacing, plane_thickness
            )
            plane_builders[plane_norm] = {
                "count": int(len(slab_indices)),
                "slice_fn": lambda idx, slabs=slab_indices: _reduce_slab(vol_zyx[slabs[idx], :, :]),
                "pixel_spacing_out": [float(spacing_y), float(spacing_x)],
                "spacing_between_slices_out": float(plane_spacing),
                "slice_thickness_out": float(plane_thickness),
                "image_orientation_out": _orientation_pair(patient_x_vec, patient_y_vec, patient_z_vec),
                "plane_normal_out": [float(patient_z_vec[0]), float(patient_z_vec[1]), float(patient_z_vec[2])],
                "position_fn": lambda idx, centers=slab_centers_idx: position_from_patient_zyx(centers[idx], 0.0, 0.0),
            }
        elif plane_norm == "coronal":
            plane_spacing = float(requested_slice_spacing) if requested_slice_spacing is not None else float(spacing_y)
            plane_thickness = float(requested_slab_thickness) if requested_slab_thickness is not None else float(plane_spacing)
            slab_indices, slab_centers_idx = _slice_indices_for_axis(
                int(vol_zyx.shape[1]), float(spacing_y), plane_spacing, plane_thickness
            )
            plane_builders[plane_norm] = {
                "count": int(len(slab_indices)),
                "slice_fn": lambda idx, slabs=slab_indices: _reduce_slab(vol_zyx[:, slabs[idx], :], reduce_axis=1),
                "pixel_spacing_out": [float(spacing_z), float(spacing_x)],
                "spacing_between_slices_out": float(plane_spacing),
                "slice_thickness_out": float(plane_thickness),
                "image_orientation_out": _orientation_pair(patient_x_vec, patient_z_vec, patient_y_vec),
                "plane_normal_out": [float(patient_y_vec[0]), float(patient_y_vec[1]), float(patient_y_vec[2])],
                "position_fn": lambda idx, centers=slab_centers_idx: position_from_patient_zyx(0.0, centers[idx], 0.0),
            }
        elif plane_norm == "sagittal":
            plane_spacing = float(requested_slice_spacing) if requested_slice_spacing is not None else float(spacing_x)
            plane_thickness = float(requested_slab_thickness) if requested_slab_thickness is not None else float(plane_spacing)
            slab_indices, slab_centers_idx = _slice_indices_for_axis(
                int(vol_zyx.shape[2]), float(spacing_x), plane_spacing, plane_thickness
            )
            plane_builders[plane_norm] = {
                "count": int(len(slab_indices)),
                "slice_fn": lambda idx, slabs=slab_indices: _reduce_slab(vol_zyx[:, :, slabs[idx]], reduce_axis=2),
                "pixel_spacing_out": [float(spacing_z), float(spacing_y)],
                "spacing_between_slices_out": float(plane_spacing),
                "slice_thickness_out": float(plane_thickness),
                "image_orientation_out": _orientation_pair(patient_y_vec, patient_z_vec, patient_x_vec),
                "plane_normal_out": [float(patient_x_vec[0]), float(patient_x_vec[1]), float(patient_x_vec[2])],
                "position_fn": lambda idx, centers=slab_centers_idx: position_from_patient_zyx(0.0, 0.0, centers[idx]),
            }

    if not plane_builders:
        raise ValueError("No valid reconstruction planes selected")

    total_slices = sum(v["count"] for v in plane_builders.values())
    cache.set(
        progress_key,
        {"current": 0, "total": total_slices, "message": "Reconstruction started"},
        timeout=60 * 60,
    )

    processed = 0
    created_series = []

    for plane_norm, cfg in plane_builders.items():
        recon_series = atlas_views._create_series_derivative(series, f"Recon {plane_norm.title()} ({source_plane.title()} src)")

        for idx in range(int(cfg["count"])):
            arr2d = cfg["slice_fn"](idx)
            ds_new = copy.deepcopy(template_ds)
            arr2d = np.asarray(arr2d, dtype=dicom_items[0][2].dtype)

            ds_new.Rows = int(arr2d.shape[0])
            ds_new.Columns = int(arr2d.shape[1])
            ds_new.InstanceNumber = idx + 1
            atlas_views._stamp_derived_dataset(
                ds_new,
                series_instance_uid=recon_series.series_instance_uid,
                source_series_instance_uid=series.series_instance_uid,
                series_description=recon_series.description or f"Recon {plane_norm.title()}",
                derivation_description=(
                f"{plane_norm.title()} reconstruction; mode={recon_thickness_mode}; "
                f"thickness={float(cfg['slice_thickness_out']):.3f}mm; spacing={float(cfg['spacing_between_slices_out']):.3f}mm"
                ),
                image_type=["DERIVED", "SECONDARY", "MPR"],
            )
            ds_new.PixelData = arr2d.tobytes()
            ds_new.PixelSpacing = [float(cfg["pixel_spacing_out"][0]), float(cfg["pixel_spacing_out"][1])]
            ds_new.ImageOrientationPatient = cfg["image_orientation_out"]
            ipp = cfg["position_fn"](idx)
            ds_new.ImagePositionPatient = ipp
            ds_new.SliceThickness = float(cfg["slice_thickness_out"])
            ds_new.SpacingBetweenSlices = float(cfg["spacing_between_slices_out"])
            try:
                normal = np.asarray(cfg["plane_normal_out"], dtype=float)
                ds_new.SliceLocation = float(np.dot(np.asarray(ipp, dtype=float), normal))
            except Exception:
                pass

            out_io = io.BytesIO()
            ds_new.save_as(out_io, write_like_original=False)
            out_io.seek(0)

            recon_image = SeriesImage(
                series=recon_series,
                position=idx + 1,
                upload_filename=f"recon_{plane_norm}_{idx + 1}.dcm",
            )
            recon_image.image.save(
                f"recon_{plane_norm}_{recon_series.pk}_{idx + 1}.dcm",
                ContentFile(out_io.getvalue()),
                save=False,
            )
            recon_image.save()

            processed += 1
            cache.set(
                progress_key,
                {
                    "current": processed,
                    "total": total_slices,
                    "message": f"Generating {plane_norm} reconstruction ({processed}/{total_slices})",
                },
                timeout=60 * 60,
            )

        created_series.append(
            {
                "id": recon_series.pk,
                "url": recon_series.get_absolute_url(),
                "description": recon_series.description or str(recon_series.pk),
            }
        )

    logger.info(
        "Reconstruction task complete for series {} with {} outputs",
        series.pk,
        len(created_series),
    )

    cache.set(
        progress_key,
        {
            "current": total_slices,
            "total": total_slices,
            "message": "Reconstruction complete",
        },
        timeout=10 * 60,
    )

    return {
        "series_id": series.pk,
        "created_series": created_series,
        "target_spacing": float(requested_slice_spacing) if requested_slice_spacing is not None else None,
        "slab_thickness": float(requested_slab_thickness) if requested_slab_thickness is not None else None,
        "mode": recon_thickness_mode,
        "source_plane": source_plane,
    }