import sys
import os
import glob
from pydicom import dcmread

from pynetdicom import AE, StoragePresentationContexts


def get_dicom_files(path):
    """Get list of DICOM files from path (file or directory)"""
    if os.path.isfile(path):
        return [path]
    elif os.path.isdir(path):
        # Find all .dcm files in directory and subdirectories
        dicom_files = []
        for root, dirs, files in os.walk(path):
            for file in files:
                if file.lower().endswith('.dcm'):
                    dicom_files.append(os.path.join(root, file))
        dicom_files.sort()
        return dicom_files
    else:
        return []


def process_dicom_file(file_path, client_aetitle, new_patient_id, new_patient_name, new_accession_number, new_study_date):
    """Process and send a single DICOM file"""
    try:
        ds = dcmread(file_path)

        print(f"Обработка файла: {os.path.basename(file_path)}")

        # Update DICOM fields
        ds.PatientID = new_patient_id
        ds.PatientName = new_patient_name
        ds.AccessionNumber = new_accession_number
        ds.StudyDate = new_study_date
        del ds.StudyDescription

        return ds
    except Exception as e:
        print(f"Ошибка при чтении файла {file_path}: {e}")
        return None


def main():
    # Get file or directory path
    path = None
    if len(sys.argv) > 1:
        path = sys.argv[1]

    if not path:
        path = input("Путь к файлу или директории: ")

    if not os.path.exists(path):
        print("Путь не существует")
        exit(1)

    # Get DICOM files
    dicom_files = get_dicom_files(path)

    if not dicom_files:
        print("DICOM файлы не найдены")
        exit(1)

    print(f"Найдено {len(dicom_files)} DICOM файлов")

    # Get connection parameters
    client_aetitle = input("Aetitle аппарата: ")

    # Get new values for DICOM fields
    new_patient_id = input("Новый ид пациента: ")
    new_patient_name = input("Новое имя пациента: ")
    new_accession_number = input("Новый номер исследования: ")
    new_study_date = input("Новая дата исследования: ")

    # Create AE and establish association
    ae = AE(ae_title=client_aetitle)
    ae.requested_contexts = StoragePresentationContexts

    print("\nУстановка соединения с PACS...")
    assoc = ae.associate("127.0.0.1", 104, ae_title='SAMPACS')

    if not assoc.is_established:
        print('Подключиться не удалось')
        exit(1)

    success_count = 0
    error_count = 0

    # Process each DICOM file
    for file_path in dicom_files:
        ds = process_dicom_file(file_path, client_aetitle, new_patient_id,
                               new_patient_name, new_accession_number, new_study_date)

        if ds:
            try:
                status = assoc.send_c_store(ds)
                if status:
                    success_count += 1
                else:
                    print('Ошибка при отправке файла')
                    error_count += 1
            except Exception as e:
                print(f'Ошибка при отправке файла {file_path}: {e}')
                error_count += 1
        else:
            error_count += 1

    # Release association
    assoc.release()

    print(f"\nИтог:")
    print(f"Успешно отправлено: {success_count}")
    print(f"Ошибок: {error_count}")
    print(f"Всего обработано: {len(dicom_files)}")


if __name__ == "__main__":
    main()
