"""
DICOM sender service - Implements storescu functionality using pynetdicom.
"""
import os
import glob
import pydicom
import shutil
from pydicom.dataset import Dataset
from pynetdicom import AE, StoragePresentationContexts, evt, build_role
from config import settings
from utils.logger import main_logger as logger
from utils.error_handler import DicomStoreError, dicom_retry
from utils.dicom_utils import create_directory_if_not_exists
from pynetdicom.sop_class import (
    StudyRootQueryRetrieveInformationModelGet,
    PatientRootQueryRetrieveInformationModelGet
)

class DicomSender:
    """
    Class to perform DICOM C-STORE operations to send DICOM studies to a destination PACS.
    """
    
    def __init__(self, pacs_config=None):
        """
        Initialize DicomSender with PACS settings.
        
        Args:
            pacs_config (dict, optional): Destination PACS configuration dict containing host, port, aet
        """
        self.pacs_config = pacs_config or settings.DESTINATION_PACS
        self.ae = AE(ae_title=settings.SOURCE_AET)
        
        # Add the Query/Retrieve SOP classes
        self.ae.add_requested_context(StudyRootQueryRetrieveInformationModelGet)
        self.ae.add_requested_context(PatientRootQueryRetrieveInformationModelGet)
        
        # Add storage presentation contexts (needed for receiving the images)
        # for context in StoragePresentationContexts:
        #     self.ae.add_requested_context(context.abstract_syntax)
        
        storage_uids = [
            '1.2.840.10008.5.1.4.1.1.1',      # CR Storage
            '1.2.840.10008.5.1.4.1.1.1.1',    # Digital X-Ray Image Storage
            '1.2.840.10008.5.1.4.1.1.2',      # CT Image Storage
            '1.2.840.10008.5.1.4.1.1.4',      # MR Image Storage
            '1.2.840.10008.5.1.4.1.1.7',      # Secondary Capture Image Storage
            '1.2.840.10008.5.1.4.1.1.6.1',    # Ultrasound Image Storage
            '1.2.840.10008.5.1.4.1.1.128',    # PET Image Storage
            '1.2.840.10008.5.1.4.1.1.20',     # Nuclear Medicine Image Storage
            '1.2.840.10008.5.1.4.1.1.9.1.1',  # 12-lead ECG Waveform Storage
            '1.2.840.10008.5.1.4.1.1.9.1.2',  # General ECG Waveform Storage
        ]
        self.ext_neg = []
        for uid in storage_uids:
            self.ae.add_requested_context(uid)
            role = build_role(uid, scp_role=True)
            self.ext_neg.append(role)  

        
        logger.info(f"DicomSender initialized with destination PACS: {self.pacs_config['aet']}@{self.pacs_config['host']}:{self.pacs_config['port']}")
    
    def _remove_dicom_files(self, directory):
        """
        Remove all DICOM files in the specified directory.
        
        Args:
            directory (str): Path to the directory to remove files from.
        """
        try:
            if os.path.exists(directory):
                shutil.rmtree(directory)
                logger.info(f"Removed DICOM files from directory: {directory}")
        except Exception as e:
            logger.error(f"Failed to remove DICOM files from {directory}: {str(e)}")

    @dicom_retry(exception_types=(DicomStoreError, ConnectionError))
    def send_study(self, study_path):
        """
        Send all DICOM files in a study directory to the destination PACS.
        
        Args:
            study_path (str): Path to study directory containing series subdirectories with DICOM files
            
        Returns:
            dict: Summary of store operations with counts and status
        """
        logger.info(f"Sending all DICOM files from study: {study_path}")
        
        # Check if study directory exists
        if not os.path.exists(study_path) or not os.path.isdir(study_path):
            error_msg = f"Study directory does not exist: {study_path}"
            logger.error(error_msg)
            raise DicomStoreError(error_msg)
        
        # Find all DICOM files in the study directory (recursively)
        dicom_files = []
        for root, _, files in os.walk(study_path):
            for file in files:
                if file.endswith('.dcm'):
                    dicom_files.append(os.path.join(root, file))
        
        if not dicom_files:
            logger.warning(f"No DICOM files found in study directory: {study_path}")
            return {
                'success': False,
                'total_files': 0,
                'successful_sends': 0,
                'failed_sends': 0,
                'error': 'No DICOM files found'
            }
        
        logger.info(f"Found {len(dicom_files)} DICOM files to send")
        
        # Create association
        assoc = self.ae.associate(
            self.pacs_config['host'],
            self.pacs_config['port'],
            ae_title=self.pacs_config['aet']        
            )
        
        result = {
            'success': False,
            'total_files': len(dicom_files),
            'successful_sends': 0,
            'failed_sends': 0,
            'failures': [],
            'study_path': study_path
        }
        
        if assoc.is_established:
            try:
                logger.debug(f"Association established with {self.pacs_config['aet']}, sending DICOM files")
                
                for file_path in dicom_files:
                    try:
                        # Read the DICOM file
                        dataset = pydicom.dcmread(file_path, force=True)
                        
                        # Get identifying information for logging
                        sop_instance_uid = getattr(dataset, 'SOPInstanceUID', 'Unknown')
                        series_instance_uid = getattr(dataset, 'SeriesInstanceUID', 'Unknown')
                        
                        # Send the dataset
                        status = assoc.send_c_store(dataset)
                        
                        if status and status.Status == 0x0000:  # Success
                            logger.debug(f"Successfully sent: {os.path.basename(file_path)}")
                            result['successful_sends'] += 1
                        else:
                            logger.error(f"Failed to send {os.path.basename(file_path)}: {status}")
                            result['failed_sends'] += 1
                            result['failures'].append({
                                'file': file_path,
                                'sop_instance_uid': sop_instance_uid,
                                'series_instance_uid': series_instance_uid,
                                'status': str(status) if status else 'Unknown error'
                            })
                    except Exception as e:
                        logger.error(f"Error processing file {file_path}: {str(e)}")
                        result['failed_sends'] += 1
                        result['failures'].append({
                            'file': file_path,
                            'error': str(e)
                        })
                
                # Update success flag based on results
                if result['successful_sends'] == result['total_files']:
                    result['success'] = True
                elif result['successful_sends'] > 0:
                    # Partial success
                    result['success'] = True
                    logger.warning(f"Partial success: {result['successful_sends']}/{result['total_files']} files sent")
                
                logger.info(f"C-STORE completed: {result['successful_sends']}/{result['total_files']} files sent")
                
                # Remove DICOM files after sending
                self._remove_dicom_files(study_path)

            except Exception as e:
                error_msg = f"Error during C-STORE operations: {str(e)}"
                logger.error(error_msg)
                result['error'] = error_msg
                raise DicomStoreError(error_msg)
            finally:
                # Release the association
                assoc.release()
                logger.debug("Association released")
        else:
            error_msg = f"Association rejected, aborted or never connected to {self.pacs_config['aet']}"
            logger.error(error_msg)
            result['error'] = error_msg
            raise DicomStoreError(error_msg)
        
        return result
    
    @dicom_retry(exception_types=(DicomStoreError, ConnectionError))
    def send_series(self, study_uid, series_uid, base_dir=None):
        """
        Send all DICOM files in a specific series to the destination PACS.
        
        Args:
            study_uid (str): StudyInstanceUID
            series_uid (str): SeriesInstanceUID to send
            base_dir (str, optional): Base directory where DICOM files are stored. Defaults to settings.DICOM_STORE_DIR.
            
        Returns:
            dict: Summary of store operations with counts and status
        """
        base_dir = base_dir or settings.DICOM_STORE_DIR
        series_path = os.path.join(base_dir, study_uid, series_uid)
        
        logger.info(f"Sending DICOM files from series: {series_uid}")
        
        # Check if series directory exists
        if not os.path.exists(series_path) or not os.path.isdir(series_path):
            error_msg = f"Series directory does not exist: {series_path}"
            logger.error(error_msg)
            raise DicomStoreError(error_msg)
        
        # Find all DICOM files in the series directory
        dicom_files = glob.glob(os.path.join(series_path, '*.dcm'))
        
        if not dicom_files:
            logger.warning(f"No DICOM files found in series directory: {series_path}")
            return {
                'success': False,
                'total_files': 0,
                'successful_sends': 0,
                'failed_sends': 0,
                'error': 'No DICOM files found'
            }
        
        logger.info(f"Found {len(dicom_files)} DICOM files to send")
        
        # Create association
        assoc = self.ae.associate(
            self.pacs_config['host'],
            self.pacs_config['port'],
            ae_title=self.pacs_config['aet']
        )
        
        result = {
            'success': False,
            'total_files': len(dicom_files),
            'successful_sends': 0,
            'failed_sends': 0,
            'failures': [],
            'series_path': series_path
        }
        
        if assoc.is_established:
            try:
                logger.debug(f"Association established with {self.pacs_config['aet']}, sending DICOM files")
                
                for file_path in dicom_files:
                    try:
                        # Read the DICOM file
                        dataset = pydicom.dcmread(file_path, force=True)
                        
                        # Send the dataset
                        status = assoc.send_c_store(dataset)
                        
                        if status and status.Status == 0x0000:  # Success
                            logger.debug(f"Successfully sent: {os.path.basename(file_path)}")
                            result['successful_sends'] += 1
                        else:
                            logger.error(f"Failed to send {os.path.basename(file_path)}: {status}")
                            result['failed_sends'] += 1
                            result['failures'].append({
                                'file': file_path,
                                'status': str(status) if status else 'Unknown error'
                            })
                    except Exception as e:
                        logger.error(f"Error processing file {file_path}: {str(e)}")
                        result['failed_sends'] += 1
                        result['failures'].append({
                            'file': file_path,
                            'error': str(e)
                        })
                
                # Update success flag based on results
                if result['successful_sends'] == result['total_files']:
                    result['success'] = True
                elif result['successful_sends'] > 0:
                    # Partial success
                    result['success'] = True
                    logger.warning(f"Partial success: {result['successful_sends']}/{result['total_files']} files sent")
                
                logger.info(f"C-STORE completed: {result['successful_sends']}/{result['total_files']} files sent")
                
                # Remove DICOM files after sending
                self._remove_dicom_files(series_path)

            except Exception as e:
                error_msg = f"Error during C-STORE operations: {str(e)}"
                logger.error(error_msg)
                result['error'] = error_msg
                raise DicomStoreError(error_msg)
            finally:
                # Release the association
                assoc.release()
                logger.debug("Association released")
        else:
            error_msg = f"Association rejected, aborted or never connected to {self.pacs_config['aet']}"
            logger.error(error_msg)
            result['error'] = error_msg
            raise DicomStoreError(error_msg)
        
        return result
    
    @dicom_retry(exception_types=(DicomStoreError, ConnectionError))
    def send_file(self, file_path):
        """
        Send a single DICOM file to the destination PACS.
        
        Args:
            file_path (str): Path to DICOM file to send
            
        Returns:
            dict: Summary of store operation with status
        """
        logger.info(f"Sending DICOM file: {file_path}")
        
        # Check if file exists
        if not os.path.exists(file_path) or not os.path.isfile(file_path):
            error_msg = f"DICOM file does not exist: {file_path}"
            logger.error(error_msg)
            raise DicomStoreError(error_msg)
        
        # Create association
        assoc = self.ae.associate(
            self.pacs_config['host'],
            self.pacs_config['port'],
            ae_title=self.pacs_config['aet']
        )
        
        result = {
            'success': False,
            'file_path': file_path,
            'status': '',
            'error': ''
        }
        
        if assoc.is_established:
            try:
                logger.debug(f"Association established with {self.pacs_config['aet']}, sending file")
                
                # Read the DICOM file
                dataset = pydicom.dcmread(file_path, force=True)
                
                # Get identifying information for logging
                sop_instance_uid = getattr(dataset, 'SOPInstanceUID', 'Unknown')
                series_instance_uid = getattr(dataset, 'SeriesInstanceUID', 'Unknown')
                study_instance_uid = getattr(dataset, 'StudyInstanceUID', 'Unknown')
                
                # Send the dataset
                status = assoc.send_c_store(dataset)
                
                if status and status.Status == 0x0000:  # Success
                    result['success'] = True
                    result['status'] = 'Success'
                    logger.info(f"Successfully sent file: {file_path}")
                    logger.debug(f"File details: Study={study_instance_uid}, Series={series_instance_uid}, SOP={sop_instance_uid}")
                else:
                    result['status'] = f"Failed - Status: {status}" if status else "Failed - Unknown error"
                    logger.error(f"Failed to send file: {file_path}")
                    logger.error(f"Status: {status}")
                
            except Exception as e:
                error_msg = f"Error sending file {file_path}: {str(e)}"
                logger.error(error_msg)
                result['error'] = error_msg
                raise DicomStoreError(error_msg)
            finally:
                # Release the association
                assoc.release()
                logger.debug("Association released")
        else:
            error_msg = f"Association rejected, aborted or never connected to {self.pacs_config['aet']}"
            logger.error(error_msg)
            result['error'] = error_msg
            raise DicomStoreError(error_msg)
        
        return result