"""DICOM download utilities (DICOMweb) with optional NIfTI conversion."""

from __future__ import annotations

import hashlib
import logging
import os
import pathlib
import shutil
import tempfile
import argparse
import json
import multiprocessing as mp
from dataclasses import dataclass
from typing import Dict, Iterable, List, Optional, Tuple, Union, Sequence

import pandas as pd
import pydicom
import requests
import SimpleITK as sitk
from dicomweb_client.api import DICOMwebClient
from requests.auth import HTTPBasicAuth
from tqdm import tqdm

from ..fhir.auth import Auth

logger = logging.getLogger(__name__)


class DicomDownloadError(RuntimeError):
    """Raised when a DICOM download or conversion fails."""


@dataclass
class DicomResult:
    study_uid: str
    series_uid: Optional[str]
    download_id: str
    path: str
    format: str
    error: Optional[str] = None


class DICOMDownloader:
    """Download DICOM studies/series via DICOMweb and optionally convert to NIfTI.

    Parameters
    ----------
    auth:
        Optional Auth or requests.Session to reuse.
    dicom_web_url:
        Base DICOMweb endpoint (e.g., https://pacs.example.org/dicomweb).
        If omitted, uses DICOM_WEB_URL env.
    dicom_user / dicom_password:
        Optional basic auth credentials for DICOMweb (overrides Auth/session).
    output_format:
        "dicom" or "nifti" (default: nifti).
    use_compression:
        Whether to compress NIfTI output.
    """

    def __init__(
        self,
        auth: Optional[Union[Auth, requests.Session]] = None,
        dicom_web_url: Optional[str] = None,
        dicom_user: Optional[str] = None,
        dicom_password: Optional[str] = None,
        output_format: str = "nifti",
        use_compression: bool = False,
        num_processes: int = 1,
        log_downloads: bool = False,
    ) -> None:
        """Initialize the DICOM downloader."""
        self.dicom_web_url = dicom_web_url or os.environ.get("DICOM_WEB_URL")
        if not self.dicom_web_url:
            raise ValueError("dicom_web_url is required (or set DICOM_WEB_URL).")

        self.use_compression = use_compression
        self.output_format = self._normalize_format(output_format)
        self.num_processes = max(1, num_processes)
        self.log_downloads = log_downloads

        env_user = os.environ.get("DICOM_USER")
        env_password = os.environ.get("DICOM_PASSWORD")
        dicom_user = dicom_user or env_user
        dicom_password = dicom_password or env_password
        self.dicom_user = dicom_user
        self.dicom_password = dicom_password

        if dicom_user or dicom_password:
            session = requests.Session()
            session.auth = HTTPBasicAuth(dicom_user or "", dicom_password or "")
            self.session = session
            self._close_session_on_exit = True
        elif isinstance(auth, Auth):
            self.session = auth.session
            self._close_session_on_exit = False
        elif isinstance(auth, requests.Session):
            self.session = auth
            self._close_session_on_exit = False
        else:
            self.session = requests.Session()
            self._close_session_on_exit = True

        self.client = DICOMwebClient(url=self.dicom_web_url, session=self.session)
        if self.log_downloads:
            logger.info(
                "Initialized DICOMDownloader url=%s format=%s processes=%s",
                self.dicom_web_url,
                self.output_format,
                self.num_processes,
            )

    def __enter__(self) -> "DICOMDownloader":
        """Return self for context manager usage."""
        return self

    def close(self) -> None:
        """Close the underlying session if owned."""
        if self._close_session_on_exit:
            self.session.close()

    def __exit__(self, exc_type, exc, exc_tb) -> None:
        """Close the session on context manager exit."""
        self.close()

    def _run_task(
        self,
        task: Tuple[str, Optional[str], bool],
        output_dir: Union[str, pathlib.Path],
        skip_existing: bool,
        write_manifest: bool,
    ) -> List[DicomResult]:
        """Execute a single download task (series or study)."""
        study_uid, series_uid, download_full_study = task
        try:
            if download_full_study or series_uid is None:
                return self.download_study(
                    study_uid=study_uid,
                    output_dir=output_dir,
                    skip_existing=skip_existing,
                    write_manifest=write_manifest,
                )
            return [
                self.download_series(
                    study_uid=study_uid,
                    series_uid=series_uid,
                    output_dir=output_dir,
                    skip_existing=skip_existing,
                )
            ]
        except Exception as exc:
            logger.exception("Download failed study=%s series=%s", study_uid, series_uid)
            return [
                DicomResult(
                    study_uid=study_uid,
                    series_uid=series_uid,
                    download_id=self.download_id(study_uid, series_uid),
                    path=str(output_dir),
                    format=self.output_format,
                    error=str(exc),
                )
            ]

    @staticmethod
    def download_id(study_uid: str, series_uid: Optional[str] = None) -> str:
        """Compute a deterministic download ID."""
        key = study_uid if series_uid is None else f"{study_uid}_{series_uid}"
        return hashlib.sha256(key.encode()).hexdigest()

    def download_series(
        self,
        study_uid: str,
        series_uid: str,
        output_dir: Union[str, pathlib.Path],
        skip_existing: bool = True,
    ) -> DicomResult:
        """Download a single series to its own folder."""
        download_id = self.download_id(study_uid, series_uid)
        target_dir = pathlib.Path(output_dir) / download_id
        return self._download_series_to_dir(
            study_uid=study_uid,
            series_uid=series_uid,
            target_dir=target_dir,
            skip_existing=skip_existing,
        )

    def download_study(
        self,
        study_uid: str,
        output_dir: Union[str, pathlib.Path],
        skip_existing: bool = True,
        write_manifest: bool = True,
    ) -> List[DicomResult]:
        """Download all series in a study into a study folder with per-series subfolders."""
        study_id = self.download_id(study_uid)
        study_dir = pathlib.Path(output_dir) / study_id
        results: List[DicomResult] = []
        try:
            series_list = self.client.search_for_series(study_uid=study_uid)
        except Exception as exc:
            logger.exception("Failed to list series for study %s", study_uid)
            raise DicomDownloadError("Failed to list series for study") from exc
        for series in series_list:
            series_uid = series.get("0020000E", {}).get("Value", [None])[0]
            if not series_uid:
                continue
            series_dir = study_dir / series_uid
            results.append(
                self._download_series_to_dir(
                    study_uid=study_uid,
                    series_uid=series_uid,
                    target_dir=series_dir,
                    skip_existing=skip_existing,
                )
            )
        if write_manifest:
            _write_study_manifest(study_dir, results)
        return results

    def download_from_dataframe(
        self,
        df: pd.DataFrame,
        output_dir: Union[str, pathlib.Path],
        study_uid_col: str = "study_instance_uid",
        series_uid_col: str = "series_instance_uid",
        download_full_study: bool = False,
        skip_existing: bool = True,
        write_manifest: bool = True,
        results_csv: Optional[Union[str, pathlib.Path]] = None,
    ) -> pd.DataFrame:
        """Download studies/series defined by a DataFrame.

        Returns a DataFrame containing download paths and errors (if any).
        """
        if self.log_downloads:
            logger.info("Starting DICOM download for %s rows", len(df))
        output_dir = pathlib.Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        if study_uid_col not in df.columns:
            raise ValueError(f"Missing study column '{study_uid_col}'.")
        if not download_full_study and series_uid_col not in df.columns:
            raise ValueError(f"Missing series column '{series_uid_col}'.")

        tasks = _build_tasks(
            df=df,
            study_uid_col=study_uid_col,
            series_uid_col=series_uid_col,
            download_full_study=download_full_study,
        )
        results: List[DicomResult] = []
        if self.num_processes > 1 and not (self.dicom_user or self.dicom_password):
            if self.log_downloads:
                logger.info(
                    "Multiprocessing disabled: no DICOM_USER/DICOM_PASSWORD provided."
                )
            self.num_processes = 1

        if self.num_processes > 1:
            with mp.Pool(processes=self.num_processes) as pool:
                for result in tqdm(
                    pool.imap_unordered(
                        _worker_download_task,
                        [
                            (task, self.dicom_web_url, self.output_format, self.use_compression,
                             skip_existing, write_manifest, output_dir,
                             self.dicom_user, self.dicom_password)
                            for task in tasks
                        ],
                    ),
                    total=len(tasks),
                    desc="Downloading DICOM",
                ):
                    results.extend(result)
        else:
            for task in tqdm(tasks, total=len(tasks), desc="Downloading DICOM"):
                results.extend(
                    self._run_task(
                        task=task,
                        output_dir=output_dir,
                        skip_existing=skip_existing,
                        write_manifest=write_manifest,
                    )
                )
        df_out = pd.DataFrame([r.__dict__ for r in results])
        if results_csv:
            pathlib.Path(results_csv).parent.mkdir(parents=True, exist_ok=True)
            df_out.to_csv(results_csv, index=False)
        return df_out

    def download_from_df(
        self,
        df: pd.DataFrame,
        output_dir: Union[str, pathlib.Path],
        study_uid_col: str = "study_instance_uid",
        series_uid_col: str = "series_instance_uid",
        download_full_study: bool = False,
        skip_existing: bool = True,
        write_manifest: bool = True,
        results_csv: Optional[Union[str, pathlib.Path]] = None,
    ) -> pd.DataFrame:
        """Short alias for download_from_dataframe."""
        return self.download_from_dataframe(
            df=df,
            output_dir=output_dir,
            study_uid_col=study_uid_col,
            series_uid_col=series_uid_col,
            download_full_study=download_full_study,
            skip_existing=skip_existing,
            write_manifest=write_manifest,
            results_csv=results_csv,
        )

    def _download_series_to_dir(
        self,
        study_uid: str,
        series_uid: str,
        target_dir: pathlib.Path,
        skip_existing: bool = True,
    ) -> DicomResult:
        """Download a series into a target directory."""
        target_dir.mkdir(parents=True, exist_ok=True)
        download_id = self.download_id(study_uid, series_uid)
        if skip_existing and _series_exists(target_dir, self.output_format):
            if self.log_downloads:
                logger.info("Skipping existing series %s", series_uid)
            return DicomResult(
                study_uid=study_uid,
                series_uid=series_uid,
                download_id=download_id,
                path=str(target_dir),
                format=self.output_format,
                error=None,
            )
        try:
            with tempfile.TemporaryDirectory() as tmpdir:
                tmp_path = pathlib.Path(tmpdir)
                for instance in self.client.iter_series(study_uid=study_uid, series_uid=series_uid):
                    instance.save_as(tmp_path / f"{instance.SOPInstanceUID}.dcm")
                if self.output_format == "dicom":
                    for dcm in tmp_path.glob("*.dcm"):
                        shutil.copy2(dcm, target_dir / dcm.name)
                else:
                    series_reader = sitk.ImageSeriesReader()
                    files = series_reader.GetGDCMSeriesFileNames(str(tmp_path))
                    if not files:
                        raise DicomDownloadError("No DICOM files found for series.")
                    series_reader.SetFileNames(files)
                    image = series_reader.Execute()
                    out_path = target_dir / f"{series_uid}.nii.gz"
                    sitk.WriteImage(image, str(out_path), useCompression=self.use_compression)
        except Exception as exc:
            logger.exception("Failed to download series %s", series_uid)
            return DicomResult(
                study_uid=study_uid,
                series_uid=series_uid,
                download_id=download_id,
                path=str(target_dir),
                format=self.output_format,
                error=str(exc),
            )

        return DicomResult(
            study_uid=study_uid,
            series_uid=series_uid,
            download_id=download_id,
            path=str(target_dir),
            format=self.output_format,
            error=None,
        )

    @staticmethod
    def _normalize_format(output_format: str) -> str:
        """Normalize output format to 'dicom' or 'nifti'."""
        fmt = output_format.strip().lower()
        if fmt in {"dicom", "dcm", ".dcm"}:
            return "dicom"
        if fmt in {"nifti", "nii", "nii.gz", ".nii", ".nii.gz"}:
            return "nifti"
        raise ValueError("output_format must be 'dicom' or 'nifti'.")


def _series_exists(target_dir: pathlib.Path, output_format: str) -> bool:
    """Return True if a series output already exists."""
    if output_format == "dicom":
        return any(target_dir.glob("*.dcm"))
    return any(target_dir.glob("*.nii")) or any(target_dir.glob("*.nii.gz"))


def _write_study_manifest(study_dir: pathlib.Path, results: Sequence[DicomResult]) -> None:
    """Write a study manifest JSON file with series results."""
    study_dir.mkdir(parents=True, exist_ok=True)
    manifest = {
        "series": [
            {
                "study_uid": r.study_uid,
                "series_uid": r.series_uid,
                "download_id": r.download_id,
                "path": r.path,
                "format": r.format,
                "error": r.error,
            }
            for r in results
        ]
    }
    with (study_dir / "study_manifest.json").open("w") as f:
        json.dump(manifest, f, indent=2)


def _build_tasks(
    df: pd.DataFrame,
    study_uid_col: str,
    series_uid_col: str,
    download_full_study: bool,
) -> List[Tuple[str, Optional[str], bool]]:
    """Build download tasks from a DataFrame."""
    tasks: List[Tuple[str, Optional[str], bool]] = []
    for row in df.itertuples(index=False):
        study_uid = getattr(row, study_uid_col)
        series_uid = getattr(row, series_uid_col, None) if not download_full_study else None
        tasks.append((study_uid, series_uid, download_full_study))
    return tasks


def _worker_download_task(
    args: Tuple[
        Tuple[str, Optional[str], bool],
        str,
        str,
        bool,
        bool,
        bool,
        Union[str, pathlib.Path],
        Optional[str],
        Optional[str],
    ],
) -> List[DicomResult]:
    """Worker helper for parallel DICOM downloads."""
    task, dicom_web_url, output_format, use_compression, skip_existing, write_manifest, output_dir, user, password = args
    downloader = DICOMDownloader(
        dicom_web_url=dicom_web_url,
        dicom_user=user,
        dicom_password=password,
        output_format=output_format,
        use_compression=use_compression,
    )
    return downloader._run_task(
        task=task,
        output_dir=output_dir,
        skip_existing=skip_existing,
        write_manifest=write_manifest,
    )


def main() -> int:
    """CLI entrypoint for DICOM downloads."""
    parser = argparse.ArgumentParser(description="Fyron DICOM downloader")
    parser.add_argument("--dicom-web-url", default=os.environ.get("DICOM_WEB_URL"))
    parser.add_argument("--dicom-user", default=os.environ.get("DICOM_USER"))
    parser.add_argument("--dicom-password", default=os.environ.get("DICOM_PASSWORD"))
    parser.add_argument("--study-uid")
    parser.add_argument("--series-uid", default=None)
    parser.add_argument("--csv", help="Path to CSV with study/series columns")
    parser.add_argument("--study-col", default="study_instance_uid")
    parser.add_argument("--series-col", default="series_instance_uid")
    parser.add_argument("--full-study", action="store_true")
    parser.add_argument("--output-dir", default="dicom_out")
    parser.add_argument("--output-format", default="nifti")
    parser.add_argument("--compression", action="store_true")
    parser.add_argument("--processes", type=int, default=1)
    parser.add_argument("--skip-existing", action="store_true")
    parser.add_argument("--manifest", action="store_true")
    parser.add_argument("--results-csv", default=None)
    args = parser.parse_args()

    if not args.dicom_web_url:
        raise SystemExit("DICOM_WEB_URL is required (or pass --dicom-web-url).")

    downloader = DICOMDownloader(
        dicom_web_url=args.dicom_web_url,
        dicom_user=args.dicom_user,
        dicom_password=args.dicom_password,
        output_format=args.output_format,
        use_compression=args.compression,
        num_processes=args.processes,
    )

    if args.csv:
        df = pd.read_csv(args.csv)
        results = downloader.download_from_dataframe(
            df=df,
            output_dir=args.output_dir,
            study_uid_col=args.study_col,
            series_uid_col=args.series_col,
            download_full_study=args.full_study,
            skip_existing=args.skip_existing,
            write_manifest=args.manifest,
            results_csv=args.results_csv,
        )
        print(results)
        return 0

    if not args.study_uid:
        raise SystemExit("Provide --study-uid or --csv.")

    if args.series_uid:
        result = downloader.download_series(
            study_uid=args.study_uid,
            series_uid=args.series_uid,
            output_dir=args.output_dir,
            skip_existing=args.skip_existing,
        )
        print(result)
    else:
        results = downloader.download_study(
            study_uid=args.study_uid,
            output_dir=args.output_dir,
            skip_existing=args.skip_existing,
            write_manifest=args.manifest,
        )
        for result in results:
            print(result)
    return 0
