"""
DICOM finder service - Implements findscu functionality using pynetdicom.
"""
import os
from datetime import datetime
from pydicom.dataset import Dataset
from pynetdicom import AE, evt, build_role, debug_logger
from pynetdicom.sop_class import (
    PatientRootQueryRetrieveInformationModelFind,
    StudyRootQueryRetrieveInformationModelFind,
    PatientStudyOnlyQueryRetrieveInformationModelFind
)
from config import settings
from utils.logger import main_logger as logger
from utils.error_handler import DicomQueryError, dicom_retry
from utils.dicom_utils import parse_dicom_date

# Set debug level
# debug_logger()  # Uncomment for detailed debug logs

class DicomFinder:
    """
    Class to perform DICOM C-FIND operations at different levels (STUDY, SERIES, IMAGE).
    """
    
    def __init__(self, pacs_config=None):
        """
        Initialize DicomFinder with PACS settings.
        
        Args:
            pacs_config (dict, optional): PACS configuration dict containing host, port, aet
        """
        self.pacs_config = pacs_config or settings.SOURCE_PACS
        self.ae = AE(ae_title=settings.SOURCE_AET)
        
        # Add the supported presentation contexts (Query/Retrieve SOP classes)
        self.ae.add_requested_context(StudyRootQueryRetrieveInformationModelFind)
        self.ae.add_requested_context(PatientRootQueryRetrieveInformationModelFind)
        self.ae.add_requested_context(PatientStudyOnlyQueryRetrieveInformationModelFind)
        
        logger.info(f"DicomFinder initialized with PACS: {self.pacs_config['aet']}@{self.pacs_config['host']}:{self.pacs_config['port']}")
    
    @dicom_retry(exception_types=(DicomQueryError, ConnectionError))
    def find_studies_by_date_range(self, start_date, end_date, additional_filters=None):
        """
        Find studies within a date range.
        
        Args:
            start_date (str): Start date in YYYYMMDD format
            end_date (str): End date in YYYYMMDD format
            additional_filters (dict, optional): Additional DICOM attributes to filter by
            
        Returns:
            list: List of study datasets
        """
        logger.info(f"Finding studies from {start_date} to {end_date}")
        
        # Create query dataset
        ds = Dataset()
        ds.QueryRetrieveLevel = 'STUDY'
        
        # Required fields (minimal set)
        ds.StudyInstanceUID = ''
        ds.StudyDate = ''
        
        # Additional fields we want to retrieve
        ds.AccessionNumber = ''
        ds.PatientID = ''  # MedrecID
        ds.StudyDescription = ''
        ds.StudyTime = ''
        ds.NumberOfStudyRelatedSeries = ''
        
        # Set date range
        ds.StudyDate = f"{start_date}-{end_date}"
        
        # Add any additional filters
        if additional_filters:
            for key, value in additional_filters.items():
                setattr(ds, key, value)
        
        studies = []
        
        # Create association
        assoc = self.ae.associate(
            self.pacs_config['host'],
            self.pacs_config['port'],
            ae_title=self.pacs_config['aet']
        )
        
        if assoc.is_established:
            try:
                logger.debug("Association established, sending C-FIND request")
                
                # Send C-FIND request
                responses = assoc.send_c_find(
                    ds, 
                    StudyRootQueryRetrieveInformationModelFind
                )
                
                # Process responses
                for (status, dataset) in responses:
                    if status and status.Status == 0xFF00:  # Pending
                        if dataset:
                            studies.append(dataset)
                            logger.debug(f"Found study: {dataset.StudyInstanceUID}")
                    elif status and status.Status != 0x0000:  # Not success
                        logger.error(f"C-FIND error: {status}")
                
                logger.info(f"Found {len(studies)} studies in date range {start_date}-{end_date}")
                
            except Exception as e:
                logger.error(f"Error during C-FIND: {str(e)}")
                raise DicomQueryError(f"C-FIND operation failed: {str(e)}")
            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)
            raise DicomQueryError(error_msg)
        
        return studies
    
    @dicom_retry(exception_types=(DicomQueryError, ConnectionError))
    def find_series_for_study(self, study_instance_uid, accession_number=None):
        """
        Find all series for a given study.
        
        Args:
            study_instance_uid (str): StudyInstanceUID
            
        Returns:
            list: List of series datasets
        """
        logger.info(f"Finding series for {accession_number} Study_IUID: {study_instance_uid}")
        
        # Create query dataset
        ds = Dataset()
        ds.QueryRetrieveLevel = 'SERIES'
        
        # Set the StudyInstanceUID filter
        ds.StudyInstanceUID = study_instance_uid
        
        # Required fields
        ds.SeriesInstanceUID = ''
        
        # Additional fields to retrieve
        ds.SeriesNumber = ''
        ds.SeriesDescription = ''
        ds.Modality = ''
        ds.NumberOfSeriesRelatedInstances = ''
        
        series_list = []
        
        # Create association
        assoc = self.ae.associate(
            self.pacs_config['host'],
            self.pacs_config['port'],
            ae_title=self.pacs_config['aet']
        )
        
        if assoc.is_established:
            try:
                logger.debug(f"Association established, sending SERIES-level C-FIND for study {study_instance_uid}")
                
                # Send C-FIND request
                responses = assoc.send_c_find(
                    ds, 
                    StudyRootQueryRetrieveInformationModelFind
                )
                
                # Process responses
                for (status, dataset) in responses:
                    if status and status.Status == 0xFF00:  # Pending
                        if dataset:
                            series_list.append(dataset)
                            logger.debug(f"Found series: {dataset.SeriesInstanceUID}")
                    elif status and status.Status != 0x0000:  # Not success
                        logger.error(f"C-FIND error: {status}")
                
                logger.info(f"Found {len(series_list)} series for {accession_number} with Study_IUID {study_instance_uid}")
                
            except Exception as e:
                logger.error(f"Error during SERIES C-FIND: {str(e)}")
                raise DicomQueryError(f"SERIES C-FIND operation failed: {str(e)}")
            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)
            raise DicomQueryError(error_msg)
        
        return series_list
    
    @dicom_retry(exception_types=(DicomQueryError, ConnectionError))
    def find_first_instance_for_series(self, study_instance_uid, series_instance_uid):
        """
        Find the first instance (SOP) for a given series.
        
        Args:
            study_instance_uid (str): StudyInstanceUID
            series_instance_uid (str): SeriesInstanceUID
            
        Returns:
            Dataset: Dataset of the first instance found or None
        """
        logger.info(f"Finding first instance for series: {series_instance_uid}")
        
        # Create query dataset
        ds = Dataset()
        ds.QueryRetrieveLevel = 'IMAGE'
        
        # Set the StudyInstanceUID and SeriesInstanceUID filters
        ds.StudyInstanceUID = study_instance_uid
        ds.SeriesInstanceUID = series_instance_uid
        
        # Required fields
        ds.SOPInstanceUID = ''
        ds.SOPClassUID = ''
        ds.InstanceNumber = ''
        
        # Create association
        assoc = self.ae.associate(
            self.pacs_config['host'],
            self.pacs_config['port'],
            ae_title=self.pacs_config['aet']
        )
        
        first_instance = None
        
        if assoc.is_established:
            try:
                logger.debug(f"Association established, sending IMAGE-level C-FIND for series {series_instance_uid}")
                
                # Send C-FIND request
                responses = assoc.send_c_find(
                    ds, 
                    StudyRootQueryRetrieveInformationModelFind
                )
                
                # Process responses - just get the first one
                for (status, dataset) in responses:
                    if status and status.Status == 0xFF00:  # Pending
                        if dataset:
                            # Found one instance, break
                            first_instance = dataset
                            logger.debug(f"Found first instance: {dataset.SOPInstanceUID}")
                            break
                    elif status and status.Status != 0x0000:  # Not success
                        logger.error(f"C-FIND error: {status}")
                
                if first_instance:
                    logger.info(f"Found first instance {first_instance.SOPInstanceUID} for series {series_instance_uid}")
                else:
                    logger.warning(f"No instances found for series {series_instance_uid}")
                
            except Exception as e:
                logger.error(f"Error during IMAGE C-FIND: {str(e)}")
                raise DicomQueryError(f"IMAGE C-FIND operation failed: {str(e)}")
            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)
            raise DicomQueryError(error_msg)
        
        return first_instance
    
    @dicom_retry(exception_types=(DicomQueryError, ConnectionError))
    def find_study_by_uid(self, study_instance_uid):
        """
        Find a specific study by its StudyInstanceUID.
        
        Args:
            study_instance_uid (str): StudyInstanceUID to find
            
        Returns:
            Dataset: Study dataset or None if not found
        """
        logger.info(f"Finding study with UID: {study_instance_uid}")
        
        # Create query dataset
        ds = Dataset()
        ds.QueryRetrieveLevel = 'STUDY'
        
        # Set the StudyInstanceUID filter
        ds.StudyInstanceUID = study_instance_uid
        
        # Required fields (minimal set)
        ds.StudyDate = ''
        
        # Additional fields we want to retrieve
        ds.AccessionNumber = ''
        ds.PatientID = ''  # MedrecID
        ds.StudyDescription = ''
        ds.StudyTime = ''
        ds.NumberOfStudyRelatedSeries = ''
        
        # Create association
        assoc = self.ae.associate(
            self.pacs_config['host'],
            self.pacs_config['port'],
            ae_title=self.pacs_config['aet']
        )
        
        study = None
        
        if assoc.is_established:
            try:
                logger.debug("Association established, sending C-FIND request")
                
                # Send C-FIND request
                responses = assoc.send_c_find(
                    ds, 
                    StudyRootQueryRetrieveInformationModelFind
                )
                
                # Process responses - just get the first one
                for (status, dataset) in responses:
                    if status and status.Status == 0xFF00:  # Pending
                        if dataset:
                            study = dataset
                            logger.debug(f"Found study: {dataset.StudyInstanceUID}")
                            break
                    elif status and status.Status != 0x0000:  # Not success
                        logger.error(f"C-FIND error: {status}")
                
                if study:
                    logger.info(f"Found study {study_instance_uid}")
                else:
                    logger.warning(f"No study found with UID {study_instance_uid}")
                
            except Exception as e:
                logger.error(f"Error during C-FIND: {str(e)}")
                raise DicomQueryError(f"C-FIND operation failed: {str(e)}")
            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)
            raise DicomQueryError(error_msg)
        
        return study