import os
import signal
import sys
import logging
from segflow4d.propagation.propagation_pipeline import PropagationPipeline
from segflow4d.common.types.propagation_input import PropagationInputFactory
from segflow4d.utility.file_writer import async_writer

os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

class StreamToLogger:
    """Redirect stdout/stderr to logger"""
    def __init__(self, logger, log_level=logging.INFO):
        self.logger = logger
        self.log_level = log_level
        self.linebuf = ''

    def write(self, buf):
        for line in buf.rstrip().splitlines():
            self.logger.log(self.log_level, line.rstrip())

    def flush(self):
        pass

def configure_logging(log_level='INFO', log_dir=''):
    log_level_int = getattr(logging, log_level.upper(), logging.INFO)
    handlers: list[logging.Handler] = [logging.StreamHandler()]
    if log_dir:
        os.makedirs(log_dir, exist_ok=True)
        # add a timestamped log file
        from datetime import datetime
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        log_file = os.path.join(log_dir, f'segflow4d_{timestamp}.log')
        handlers.append(logging.FileHandler(log_file))
    logging.basicConfig(
        level=log_level_int,
        format='[%(asctime)s - %(name)s - %(levelname)s] - %(message)s',
        handlers=handlers,
        force=True
    )
    
    # Redirect stdout and stderr to logging
    if log_level=='DEBUG':
        sys.stdout = StreamToLogger(logging.getLogger('stdout'), logging.INFO)
        sys.stderr = StreamToLogger(logging.getLogger('stderr'), logging.ERROR)

logger = logging.getLogger(__name__)

def parse_arguments():
    import argparse

    parser = argparse.ArgumentParser(description="SegFlow4D Propagation Pipeline")
    parser.add_argument('--image4d', type=str, help='Path to the input 4d image')
    parser.add_argument('--output', type=str, help='Path to the output directory')
    parser.add_argument('--tp-ref', type=int, help='Reference time point')
    parser.add_argument('--tp-targets', type=int, nargs='+', help='Target time points')
    parser.add_argument('--seg-ref', type=str, help='Path to the segmentation reference image')
    parser.add_argument('--log-level', type=str, default='INFO', help='Logging level (DEBUG, INFO, WARNING, ERROR)')
    parser.add_argument('--log-dir', type=str, default='', help='Directory to store log files. Prints to console if not set.')
    parser.add_argument('--add-mesh', action='append', dest='additional_meshes', 
                        help='Add additional mesh: --add-mesh <name>:<path_to_mesh_file>')
    parser.add_argument('--lowres-factor', type=float, default=0.5, help='Low resolution resample factor')
    parser.add_argument('--dilation-radius', type=int, default=2, help='Dilation radius for segmentation')
    parser.add_argument('--registration-backend', type=str, default='fireants', help='Registration backend to use')
    parser.add_argument('--min-vram', type=int, default=10, help='Minimum required VRAM in GB')
    parser.add_argument('--debug', action='store_true', help='Enable debug mode')
    parser.add_argument('--debug-dir', type=str, default='', help='Directory to store debug outputs')
    parser.add_argument('--smooth-grad-sigma-mm', type=float, default=3.0, help='Gaussian smoothing sigma for gradient field in mm (default: 3.0)')
    parser.add_argument('--smooth-warp-sigma-mm', type=float, default=1.5, help='Gaussian smoothing sigma for warp field in mm (default: 1.5)')
    parser.add_argument('--scales', type=float, nargs='+', default=None, help='Multi-resolution scales for registration, coarse to fine (e.g. --scales 4 2 1)')
    parser.add_argument('--affine-iterations', type=int, nargs='+', default=None, help='Iterations per scale for affine stage (e.g. --affine-iterations 200 100 50)')
    parser.add_argument('--deformable-iterations', type=int, nargs='+', default=None, help='Iterations per scale for deformable stage (e.g. --deformable-iterations 200 100 25)')
    parser.add_argument('--loss-type', type=str, default=None, choices=['mse', 'cc'], help='Similarity metric: mse or cc (cross-correlation)')
    parser.add_argument('--cc-kernel-size', type=int, default=None, help='Kernel size for cross-correlation loss (used when --loss-type=cc)')

    # Greedy-specific options (picsl-greedy backend)
    parser.add_argument('--metric', type=str, default=None, choices=['NCC', 'SSD', 'NMI'],
                        help='Similarity metric for greedy backend: NCC, SSD, or NMI')
    parser.add_argument('--metric-radius', type=int, nargs='+', default=None,
                        help='Patch radius for NCC metric, one value per image dimension (e.g. --metric-radius 2 2 2)')
    parser.add_argument('--affine-dof', type=int, default=None,
                        help='Degrees of freedom for greedy affine stage: 6 (rigid), 7 (similarity), 12 (full affine)')
    parser.add_argument('--smooth-sigma-pre-mm', type=float, default=None,
                        help='Pre-smoothing sigma for greedy gradient fields in mm (default: 2.0)')
    parser.add_argument('--smooth-sigma-post-mm', type=float, default=None,
                        help='Post-smoothing sigma for greedy warp fields in mm (default: 0.5)')
    parser.add_argument('--threads', type=int, default=None,
                        help='Number of CPU threads for greedy (default: all cores)')
    parser.add_argument('--greedy-use-float', action=argparse.BooleanOptionalAction, default=True,
                        help='Use single-precision (float32) for greedy registration via -float flag (default: true)')
    parser.add_argument('--label-interpolation', type=str, default=None,
                        help='Smoothing parameter for greedy LABEL reslice interpolation, e.g. "0.2vox" or "0.5mm" (default: 0.2vox)')

    parser.add_argument('--propagation-strategy-combo', type=str, default='sequential_star',
                        choices=['sequential_star', 'sasd_star'],
                        help='Strategy combo for lowres+highres stages (default: sequential_star)')
    parser.add_argument('--config', type=str, default='', help='Path to YAML configuration file')
    
    args = parser.parse_args()

    if args.additional_meshes:
        additional_meshes_dict = {}
        for item in args.additional_meshes:
            name, path = item.split(':', 1)
            additional_meshes_dict[name] = path
        args.additional_meshes = additional_meshes_dict

    return args

def load_config(config_path):
    """Load and parse YAML configuration file"""
    import yaml
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    return config

def _install_signal_handlers():
    """Register SIGTERM/SIGINT handlers to kill all worker subprocesses on exit."""
    from segflow4d.registration.registration_manager import RegistrationManager

    def _handler(signum, frame):
        logger.warning(f"Received signal {signum}, shutting down all workers...")
        try:
            RegistrationManager.get_instance().shutdown(wait=False)
        except Exception:
            pass
        try:
            async_writer.shutdown(wait=False)
        except Exception:
            pass
        os._exit(1)

    signal.signal(signal.SIGTERM, _handler)
    signal.signal(signal.SIGINT, _handler)


def main():
    args = parse_arguments()
    
    
    input_factory = PropagationInputFactory()

    if args.config:
        config = load_config(args.config)

        configure_logging(config.get('log_level', args.log_level), config.get('log_dir', args.log_dir))

        input_factory.set_image_4d_from_disk(config.get('image4d'))
        
        # Set global options
        backend_options = config.get('registration_backend_options', {})
        input_factory.set_options(
            lowres_factor=config.get('lowres_factor', 0.5),
            registration_backend=config.get('registration_backend', 'fireants'),
            dilation_radius=config.get('dilation_radius', 2),
            write_result_to_disk=True,
            output_directory=config.get('output'),
            debug=config.get('debug', False),
            debug_output_directory=config.get('debug_dir', ''),
            minimum_required_vram_gb=config.get('minimum_required_vram_gb', 0),
            propagation_strategy_combo=config.get('propagation_strategy_combo', 'sequential_star'),
            **backend_options
        )
        
        # Parse and add multiple tp_input_groups
        tp_input_groups = config.get('tp_input_groups', [])
        
        if not tp_input_groups:
            logger.error("No tp_input_groups found in configuration")
            return
        
        for group in tp_input_groups:
            tp_ref = group.get('tp_ref')
            tp_targets = group.get('tp_targets', [])
            seg_ref = group.get('seg_ref')
            additional_meshes = group.get('additional_meshes', {})

            logger.info(f"Adding tp_input_group with tp_ref={tp_ref}, tp_targets={tp_targets}")

            input_factory.add_tp_input_group_from_disk(
                tp_ref=tp_ref,
                tp_target=tp_targets,
                seg_ref_path=seg_ref,
                additional_meshes_ref=additional_meshes
            )
    else:
        configure_logging(args.log_level, args.log_dir)

        input_factory.set_image_4d_from_disk(args.image4d)
        # Use command-line arguments
        input_factory.add_tp_input_group_from_disk(
            tp_ref=args.tp_ref,
            tp_target=args.tp_targets,
            seg_ref_path=args.seg_ref,
            additional_meshes_ref=args.additional_meshes
        )
        # Collect optional backend options only when explicitly provided
        cli_backend_options = {}
        if args.smooth_grad_sigma_mm is not None:
            cli_backend_options['smooth_grad_sigma_mm'] = args.smooth_grad_sigma_mm
        if args.smooth_warp_sigma_mm is not None:
            cli_backend_options['smooth_warp_sigma_mm'] = args.smooth_warp_sigma_mm
        if args.scales is not None:
            cli_backend_options['scales'] = args.scales
        if args.affine_iterations is not None:
            cli_backend_options['affine_iterations'] = args.affine_iterations
        if args.deformable_iterations is not None:
            cli_backend_options['deformable_iterations'] = args.deformable_iterations
        if args.loss_type is not None:
            cli_backend_options['loss_type'] = args.loss_type
        if args.cc_kernel_size is not None:
            cli_backend_options['cc_kernel_size'] = args.cc_kernel_size
        # Greedy-specific options
        if args.metric is not None:
            cli_backend_options['metric'] = args.metric
        if args.metric_radius is not None:
            cli_backend_options['metric_radius'] = args.metric_radius
        if args.affine_dof is not None:
            cli_backend_options['affine_dof'] = args.affine_dof
        if args.smooth_sigma_pre_mm is not None:
            cli_backend_options['smooth_sigma_pre_mm'] = args.smooth_sigma_pre_mm
        if args.smooth_sigma_post_mm is not None:
            cli_backend_options['smooth_sigma_post_mm'] = args.smooth_sigma_post_mm
        if args.threads is not None:
            cli_backend_options['threads'] = args.threads
        if args.greedy_use_float is not None:
            cli_backend_options['use_float'] = args.greedy_use_float
        if args.label_interpolation is not None:
            cli_backend_options['label_interpolation'] = args.label_interpolation

        input_factory.set_options(
            lowres_factor=args.lowres_factor,
            registration_backend=args.registration_backend,
            dilation_radius=args.dilation_radius,
            minimum_required_vram_gb=args.min_vram,
            write_result_to_disk=True,
            output_directory=args.output,
            debug=args.debug,
            debug_output_directory=args.debug_dir,
            propagation_strategy_combo=args.propagation_strategy_combo,
            **cli_backend_options
        )

    propagation_input = input_factory.build()
    pipeline = PropagationPipeline(propagation_input)

    _install_signal_handlers()

    logger.info("Starting SegFlow4D Propagation Pipeline")
    pipeline.run()

    async_writer.shutdown(wait=True)
    logger.info("SegFlow4D Propagation Pipeline completed")

if __name__ == "__main__":
    main()



