import argparse
import logging
import filetype
import progressbar
import json
import os
import signal
import sys

import pynetdicom._config

# Try to setup Redis logging
redis_handler = None
pwlog_redis_handler = None
pynetdicomlog_redis_handler = None
pwlog = logging.getLogger('peewee')
pynetdicomlog = logging.getLogger('pynetdicom')
try:
    from redis_logger import setup_redis_logging, RedisLogHandler
    redis_handler = setup_redis_logging(
        logger_name='sampacs',
        max_history=10000,
        ttl_days=7,
        enable_console_fallback=False  # Don't fallback to console if Redis is available
    )

    if redis_handler and redis_handler._redis_available:
        print("Redis logging initialized successfully")

        # Also setup Redis logging for peewee and pynetdicom loggers
        try:
            # Setup Redis handler for peewee logger
            pwlog_redis_handler = RedisLogHandler(
                host=os.environ.get('REDIS_HOST', 'localhost'),
                port=int(os.environ.get('REDIS_PORT', 6379)),
                max_history=10000,
                ttl_days=7,
                enable_console_fallback=False
            )
            pwlog_redis_handler.setFormatter(logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(module)s - %(message)s'
            ))

            # Setup Redis handler for pynetdicom logger
            pynetdicomlog_redis_handler = RedisLogHandler(
                host=os.environ.get('REDIS_HOST', 'localhost'),
                port=int(os.environ.get('REDIS_PORT', 6379)),
                max_history=10000,
                ttl_days=7,
                enable_console_fallback=False
            )
            pynetdicomlog_redis_handler.setFormatter(logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(module)s - %(message)s'
            ))

            # Add handlers to loggers
            pwlog = logging.getLogger('peewee')
            pynetdicomlog = logging.getLogger('pynetdicom')

            # Remove existing handlers if Redis is available
            if pwlog_redis_handler._redis_available:
                for handler in pwlog.handlers[:]:
                    pwlog.removeHandler(handler)
                pwlog.addHandler(pwlog_redis_handler)
                pwlog.propagate = False
                print("Redis logging enabled for peewee (SQL queries)")

            if pynetdicomlog_redis_handler._redis_available:
                for handler in pynetdicomlog.handlers[:]:
                    pynetdicomlog.removeHandler(handler)
                pynetdicomlog.addHandler(pynetdicomlog_redis_handler)
                pynetdicomlog.propagate = False
                print("Redis logging enabled for pynetdicom (DICOM protocol)")

        except Exception as e:
            print(f"Failed to setup Redis logging for additional loggers: {e}")

    else:
        print("Redis logging not available, using console logging")
        if redis_handler:
            print(f"Redis handler exists but not available. Type: {type(redis_handler)}")
            print(f"Redis connection status: {redis_handler._redis_available}")

    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
except ImportError as e:
    print(f"Redis logger not available: {e}")
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
except Exception as e:
    print(f"Error setting up Redis logging: {e}")
    import traceback
    traceback.print_exc()
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )

LOGGER = logging.getLogger('sampacs')

from dbmodel import *
from handlers import *
from utils import saveFile, encodingTag, thread_send_images, sendImage, getImage
from config import ModalityConfig, Config, config as global_config

from pydicom.filereader import dcmread
from pydicom.uid import JPEGLosslessSV1 #, RLELossless
import pynetdicom.sop_class as sop_class
import ownTranslit
import kometa

from threading import Thread

from pynetdicom import (
    AE, evt,
    StoragePresentationContexts
)

def setup_signal_handler():
    """Setup SIGUSR1 handler to reload config in dicomServer process"""
    def sigusr1_handler(signum, frame):
        try:
            new_config = Config.from_file()
            global_config.__dict__.update(new_config.__dict__)
            LOGGER.info(f"Config reloaded successfully in dicomServer (pid={os.getpid()})")
        except Exception as e:
            LOGGER.error(f"Failed to reload config in dicomServer: {e}")

    signal.signal(signal.SIGUSR1, sigusr1_handler)
    LOGGER.info(f"SIGUSR1 handler registered in dicomServer (pid={os.getpid()})")

pynetdicom._config.VALIDATORS['AE'] = lambda x: pynetdicom._validators.validate_ae(x.strip(chr(0)))
#_config.UNRESTRICTED_STORAGE_SERVICE = True

copyImportType = 0
moveImportType = 1
indexImportType = 2

def startDicomServer():
    handlers = [(evt.EVT_C_STORE, handle_store_wrapper),
                        (evt.EVT_CONN_OPEN, handle_open),
                        (evt.EVT_C_FIND, handle_find),
                        (evt.EVT_C_GET, handle_get),
                        (evt.EVT_C_MOVE, handle_move),
                        (evt.EVT_N_ACTION, handle_n_action),
                        (evt.EVT_N_CREATE, handle_n_create),
                        (evt.EVT_N_SET, handle_n_set),
                        (evt.EVT_C_ECHO, handle_c_echo)
                    ]

    ae = AE()
    ae.supported_contexts = StoragePresentationContexts
    ae.requested_contexts = StoragePresentationContexts

    ae.add_supported_context(sop_class.Verification)
    ae.add_supported_context(sop_class.PatientRootQueryRetrieveInformationModelFind)
    ae.add_supported_context(sop_class.StudyRootQueryRetrieveInformationModelFind)
    ae.add_supported_context(sop_class.ModalityWorklistInformationFind)
    ae.add_supported_context(sop_class.PatientRootQueryRetrieveInformationModelGet)
    ae.add_supported_context(sop_class.StudyRootQueryRetrieveInformationModelGet)
    ae.add_supported_context(sop_class.PatientRootQueryRetrieveInformationModelMove)
    ae.add_supported_context(sop_class.StudyRootQueryRetrieveInformationModelMove)
    ae.add_supported_context(sop_class.HangingProtocolInformationModelGet)
    ae.add_supported_context(sop_class.HangingProtocolInformationModelMove)
    ae.add_supported_context(sop_class.StorageCommitmentPushModel)
    ae.add_supported_context(sop_class.UltrasoundImageStorage, JPEGLosslessSV1)
    #ae.add_supported_context(sop_class.UltrasoundImageStorage, RLELossless)
    ae.add_supported_context(sop_class.ModalityPerformedProcedureStep)
    ae.add_supported_context(sop_class.TwelveLeadECGWaveformStorage)
    ae.add_supported_context("1.2.392.200036.9125.1.1.2")
    for cx in ae.supported_contexts:
        cx.scp_role = True
        cx.scu_role = True

    commitment_thread = Thread(target=thread_storage_commitment, args=(commitment_queue,))
    commitment_thread.daemon = True
    commitment_thread.start()

    send_thread = Thread(target=thread_send_images)
    send_thread.daemon = True
    send_thread.start()

    if config.kometa and config.kometa.host:
        kometa_thread = Thread(target=kometa.thread_kometa_send)
        kometa_thread.daemon = True
        kometa_thread.start()

    ae.require_called_aet = config.require_called_aet
    ae.require_calling_aet = [modality.aetitle for modality in config.modalities] if config.require_calling_aet else []
    LOGGER.debug('require_calling_aet list [%s]'%','.join([modality.aetitle for modality in config.modalities]))


    LOGGER.info('Starting DICOM server on port %s with AE Title: %s', config.pacs_port, config.pacs_aetitle)
    # Check if Redis logging is enabled by looking for RedisLogHandler in handlers
    redis_enabled = False
    redis_handler_instance = None
    if redis_handler:
        # Check if any handler is an instance of the same class as redis_handler
        for h in LOGGER.handlers:
            if isinstance(h, type(redis_handler)):
                redis_enabled = True
                redis_handler_instance = h
                break

    LOGGER.info('Redis logging: %s', 'ENABLED' if redis_enabled else 'DISABLED')
    if redis_enabled and redis_handler_instance:
        LOGGER.debug('Redis handler found: %s', type(redis_handler_instance).__name__)
        if hasattr(redis_handler_instance, '_redis_available'):
            LOGGER.debug('Redis connection available: %s', redis_handler_instance._redis_available)
    elif not redis_enabled:
        LOGGER.debug('No Redis handler found. Total handlers: %d', len(LOGGER.handlers))
        for i, h in enumerate(LOGGER.handlers):
            LOGGER.debug('  Handler %d: %s', i, h.__class__.__name__)


    setup_signal_handler()
    ae.start_server(('', config.pacs_port), evt_handlers=handlers, ae_title=config.pacs_aetitle)


def importDirectory(dir: str, importType: int, importAetitle: str = "", decodeFrom: str = None, ignoreImportErrors: bool = False):
    """ Импортируем все снимки из указанной директории и завершаем программу

    Parameters
    ----------
        dir: str
            Директория для поиска файлов
        importType: int
            тип импорта. 0 копируем файлы к себе, 1 перемещаем файлы, 2 создаем только индекс
        importAetitle: str
            указанный aetitle будет установлен для импортированного study
        decodeFrom: str
            кодировка файлов. Умолчание UTF-8
        ignoreImportErrors: bool
            не останавливаться если есть ошибки импорта
    """
    filesList = []
    cached = False
    try:
        with open('filelistcache', 'r') as f:
            filesList = json.load(f)
        cached = True
        print('got %i files from cache'%len(filesList))
    except:
        print('counting files in %s...'%dir)
        for root, subdirs, files in os.walk(dir, followlinks=True):
            for filename in files:
                file_path = os.path.join(root, filename)
                kind = filetype.guess(file_path)
                if kind and kind.mime == 'application/dicom':
                    filesList.append(file_path)

    if len(filesList) == 0:
        print('No dicom files found..')
        return

    if not cached:
        with open('filelistcache', 'w') as f:
            json.dump(filesList, f, indent=2)

    autorouteList = config.autoroute

    readyFilesList = []
    with progressbar.ProgressBar(max_value=len(filesList)) as bar:
        for i, file_path in enumerate(filesList):
            ds = dcmread(file_path)
            if 'SpecificCharacterSet' not in ds and decodeFrom:
                ds.set_original_encoding(ds.is_implicit_VR, ds.is_little_endian, decodeFrom)
                ds.add(encodingTag)

            # Check for group 0002 elements in main dataset (should only be in file_meta)
            group_0002_in_ds = ds.group_dataset(0x0002)
            if len(group_0002_in_ds) > 0:
                LOGGER.warning(f'IMPORT   Found {len(group_0002_in_ds)} group 0002 elements in main dataset from file {os.path.basename(file_path)}')
                for tag, elem in group_0002_in_ds.items():
                    # Handle both DataElement and RawDataElement
                    elem_name = elem.name if hasattr(elem, 'name') else 'Unknown'
                    elem_type = type(elem).__name__
                    LOGGER.warning(f'IMPORT   Removing ({tag.group:04x},{tag.element:04x}) {elem_name} [{elem_type}]: {elem.value} from main dataset')
                    del ds[tag]

            try:
                (ds, error) = prepareDataset(ds, registrateOptions = config.registrateOptions)
                saveFile(ds, importAetitle, importType!=indexImportType)
                if importType == moveImportType:
                    os.remove(file_path)
                readyFilesList.append(file_path)
                for destName in autorouteList:
                    addImageForSend(destName, ds.SOPInstanceUID)
            except Exception as e:
                LOGGER.error('Exception in saveFile. Filename %s; Error: %s'%(file_path, str(e)))
                if not ignoreImportErrors:
                    filesList = list(set(filesList) - set(readyFilesList))
                    with open('filelistcache', 'w') as f:
                        json.dump(filesList, f, indent=2)
                    raise
            bar.update(i)

def main():
    loggerValues = {'notset': logging.NOTSET, 'debug': logging.DEBUG, 'info': logging.INFO, 'warning': logging.WARNING, 'error': logging.ERROR, 'critical': logging.CRITICAL}
    parser = argparse.ArgumentParser(description='SAMPACS dicom server')
    g = parser.add_mutually_exclusive_group()
    g.add_argument('--copyFrom', help='Directory to import from. Source files are saved while importing')
    g.add_argument('--moveFrom', help='Directory to import from. Source files are deleted while importing')
    g.add_argument('--indexFrom', help='Directory to import from. Source files are saved. Create only index db')
    parser.add_argument('--importAetitle', help='Set aetitle for imported studies')
    parser.add_argument('--echo', help='Echo provided modality name')
    parser.add_argument('--decodeFrom', help='Set encoding for imported files', choices=['cp1251'])
    parser.add_argument('--sqlLevel', help='Set sql log level', choices=loggerValues, default=os.environ.get('sqlLevel'))
    parser.add_argument('--dicomLevel', help='Set dicom log level', choices=loggerValues, default=os.environ.get('dicomLevel'))
    parser.add_argument('--logsLevel', help='Set sampacs log level', choices=loggerValues, default=os.environ.get('logsLevel'))
    parser.add_argument('--ignoreImportErrors', help='Continue importing with errors', default=False, action='store_true')
    sendGroup = parser.add_argument_group('send image', 'send image arguments')
    sendGroup.add_argument('--sendTo', help='Send image to modality')
    sendGroup.add_argument('--instance', help='Send only instance')
    sendGroup.add_argument('--study', help='Send all images from study')
    sendGroup.add_argument('--date', help='Send all studies for date')
    args = parser.parse_args()

    if args.sqlLevel:
        # Don't add StreamHandler if Redis logging is enabled for peewee
        if not pwlog_redis_handler or not pwlog_redis_handler._redis_available:
            pwlog.addHandler(logging.StreamHandler())
        pwlog.setLevel(loggerValues.get(args.sqlLevel))
    if args.dicomLevel:
        # Don't add StreamHandler if Redis logging is enabled for pynetdicom
        if not pynetdicomlog_redis_handler or not pynetdicomlog_redis_handler._redis_available:
            pynetdicomlog.addHandler(logging.StreamHandler())
        pynetdicomlog.setLevel(loggerValues.get(args.dicomLevel))
    if args.logsLevel:
        # Don't add StreamHandler if Redis logging is enabled for sampacs
        if not redis_handler or not redis_handler._redis_available:
            LOGGER.addHandler(logging.StreamHandler())
        LOGGER.setLevel(loggerValues.get(args.logsLevel))



    if args.echo:
        modalityFound = False
        for modality in config.modalities:
            modalityFound = True
            if modality.name == args.echo:
                ae = AE(ae_title=modality.pacsaetitle or config.pacs_aetitle)
                ae.add_requested_context(sop_class.Verification)
                assoc = ae.associate(modality.host, modality.port, ae_title = modality.aetitle)
                if assoc.is_established:
                    status = assoc.send_c_echo()
                    LOGGER.info('ECHO     successfully sended to {} ({})'.format(modality.aetitle, modality.host))
                else:
                    LOGGER.error('ECHO     association with %s failed'%modality.aetitle)
                break
        if not modalityFound:
            LOGGER.error('ECHO     aetitle %s not found in config!'%args.echo)
        return

    try:
        pacs_db.connect()
        pacs_db.create_tables([Patient, Study, Series, Instance, SendLog])
        checkUpdates(pacs_db)
        pacs_db.close()
    except Exception as err:
        LOGGER.error(err)
        exit(-1)

    if args.sendTo:
        dest = None
        for modality in config.modalities:
            if modality.name == args.sendTo:
                dest = modality
                break
        if dest is None:
            LOGGER.warning('SEND    Route destination %s not found' % args.sendTo)
            print('SEND    Route destination %s not found' % args.sendTo)
            return
        if args.instance:
            pacs_db.connect()
            try:
                instance = Instance.select(Instance, Study.nodeCode).join(Series).join(Study).where(Instance.uid == args.instance).get()
            except:
                instance = None
            pacs_db.close()
            if instance:
                try:
                    ds = getImage(instance.uid, instance.study_date.strftime("%Y%m"),
                                  instance.series.study.nodeCode)
                except Exception as e:
                    LOGGER.error('SEND     image for instance uid "%s" not found. %s' % (slog.instance_uid, e))
                    slog.delete_instance()
                else:
                    if sendImage(dest.name, ds):
                        LOGGER.info('SEND     success')
            else:
                LOGGER.error('SEND     instance not found')
        elif args.study:
            pacs_db.connect()
            instanceQuery = Instance.select(Instance.uid, Instance.study_date, Series.uid, Study.uid, Study.nodeCode).join(Series).join(Study)
            instanceQuery = instanceQuery.where(Study.uid == args.study)
            instanceQuery.execute()
            pacs_db.close()
            if len(instanceQuery) == 0:
                LOGGER.error('SEND     no instances found')
            for instance in instanceQuery:
                try:
                    ds = getImage(instance.uid, instance.study_date.strftime("%Y%m"),
                                  instance.series.study.nodeCode)
                except Exception as e:
                    LOGGER.error('SEND     image for instance uid "%s" not found. %s' % (slog.instance_uid, e))
                    slog.delete_instance()
                else:
                    if sendImage(dest.name, ds):
                        LOGGER.info('SEND     success')
        # elif args.date:
        #     pacs_db.connect()
        #     instanceQuery = Instance.select(Instance.uid, Instance.study_date, Series.uid, Study.uid, Study.nodeCode).join(Series).join(Study)
        #     instanceQuery = instanceQuery.where(Study.date == args.date)
        #     instanceQuery.execute()
        #     pacs_db.close()
        #     if len(instanceQuery) == 0:
        #         LOGGER.error('SEND     no instances found')
        #     for instance in instanceQuery:
        #         addImageForSend(dest.name, instance.uid)
        else:
            LOGGER.error('SEND     you must provide one of the parameters: instance, study or date')
        return

    source = args.copyFrom or args.moveFrom or args.indexFrom
    if source:
        if not os.path.isdir(source):
            LOGGER.critical('Директория %s не существует'%source)
            exit(-1)
        type = None
        if args.copyFrom:
            type = copyImportType
        elif args.moveFrom:
            type = moveImportType
        elif args.indexFrom:
            type = indexImportType
        importDirectory(source, type, args.importAetitle, args.decodeFrom, args.ignoreImportErrors)
    else:
        startDicomServer()


if __name__ == "__main__":
    main()
