import time as tt

import atexit as ae

import numpy as np
import pandas as pd

from .base import utils as ut
from .base import validation_parsing  as vp
from .base import core_split as cs

from typing import Tuple, Union


def split(
    input_data:str|pd.DataFrame,
    split_length_limit:int,
    minimum_melting_temperature:float,
    minimum_hamming_distance:int,
    minimum_overlap_length:int,
    maximum_overlap_length:int,
    output_file:str|None=None,
    separate_outputs:bool=False,
    random_seed:int|None=None,
    verbose:bool=True) -> Tuple[Union[pd.DataFrame, list], dict]:
    '''
    Split long oligos into overlapping fragments for downstream assembly, choosing overlaps that meet
    Tm and distance constraints. Returns a DataFrame containing `Split1`, `Split2`, ... columns (or a
    list of separate DataFrames when `separate_outputs` is enabled).

    Required Parameters:
        - `input_data` (`str` / `pd.DataFrame`): Path to a CSV file or DataFrame with annotated oligo pool variants.
        - `split_length_limit` (`int`): Maximum allowed length for split oligos (>= 4).
        - `minimum_melting_temperature` (`float`): Minimum overlap region Tm (>= 4 °C).
        - `minimum_hamming_distance` (`int`): Minimum overlap region pairwise Hamming distance (>= 1).
        - `minimum_overlap_length` (`int`): Minimum overlap region length (>= 15).
        - `maximum_overlap_length` (`int`): Maximum overlap region length (>= 15).

    Optional Parameters:
        - `output_file` (`str` / `None`): Filename for output DataFrame; required in CLI usage,
            optional in library usage (default: `None`).
        - `separate_outputs` (`bool`): If enabled, return per-split outputs (and optionally write per-split files)
            instead of a single combined DataFrame (default: `False`). See Notes.
        - `random_seed` (`int` / `None`): Seed for local RNG (default: `None`).
        - `verbose` (`bool`): If `True`, logs progress to stdout (default: `True`).

    Returns:
        - When `separate_outputs` is disabled: A single pandas DataFrame with `Split1`, `Split2`, ... columns.
        - When `separate_outputs` is enabled: A list of DataFrames `[df_Split1, df_Split2, ...]`, each containing
          the corresponding `SplitN` column.
        - A dictionary of stats from the last step in pipeline.
        - Saves to `output_file` if specified (separate files when `separate_outputs` is enabled).

    Notes:
        - `input_data` must contain a unique 'ID' column; all other columns must be non-empty DNA strings.
        - `minimum_overlap_length` should always be larger than `minimum_hamming_distance`.
        - Fragment count is auto-determined and can vary per oligo.
        - Each `SplitN` column is a separate oligo pool to synthesize; fragments are later assembled
            (e.g., overlap-extension PCR or Gibson) to reconstruct the full-length oligo.
        - Split fragments are returned in PCR assembly order; even-numbered split columns (`Split2`, `Split4`, ...)
            are reverse-complemented by design. Use `revcomp` to visualize overlaps and orientation.
        - Returned DataFrame contains split fragments only (sequence annotations from `input_data` are not preserved).
        - Typical workflow: run `pad` once per split column (e.g., `Split1`, `Split2`, ...) then `final` on each.
        - Use `separate_outputs` for custom per-fragment workflows (e.g., manual padding via `primer`/`motif`/`spacer`).
    '''

    # Preserve return style when the caller intentionally used ID as index.
    id_from_index = ut.get_id_index_intent(input_data)

    # Argument Aliasing
    indata      = input_data
    splitlimit  = split_length_limit
    mintmelt    = minimum_melting_temperature
    minhdist    = minimum_hamming_distance
    minoverlap  = minimum_overlap_length
    maxoverlap  = maximum_overlap_length
    outfile     = output_file
    separate    = separate_outputs
    random_seed = random_seed
    verbose     = verbose

    # Start Liner
    liner = ut.liner_engine(verbose)

    # Splitting Verbiage Print
    liner.send('\n[Oligopool Calculator: Assembly Mode - Split]\n')

    # Required Argument Parsing
    liner.send('\n Required Arguments\n')

    # First Pass indata Parsing and Validation
    (indf,
    indata_valid) = vp.get_parsed_indata_info(
        indata=indata,
        indata_field='    Input Data       ',
        required_fields=('ID',),
        precheck=False,
        liner=liner)
    input_rows = len(indf.index) if isinstance(indf, pd.DataFrame) else 0

    # Full splitlimit Validation
    splitlimit_valid = vp.get_numeric_validity(
        numeric=splitlimit,
        numeric_field='    Split Limit      ',
        numeric_pre_desc=' Split Fragments at most ',
        numeric_post_desc=' Base Pair(s) Each',
        minval=4,
        maxval=float('inf'),
        precheck=False,
        liner=liner)

    # Full mintmelt Validation
    tmelt_valid = vp.get_numeric_validity(
        numeric=mintmelt,
        numeric_field='  Melting Temperature',
        numeric_pre_desc=' At least ',
        numeric_post_desc=' °C b/w On-Target Overlaps',
        minval=4,
        maxval=float('inf'),
        precheck=False,
        liner=liner)

    # Full minhdist Validation
    minhdist_valid = vp.get_numeric_validity(
        numeric=minhdist,
        numeric_field='  Hamming Distance   ',
        numeric_pre_desc=' At least ',
        numeric_post_desc=' Mismatch(es) per Off-Target Overlap Pair ',
        minval=1,
        maxval=splitlimit if splitlimit_valid else float('inf'),
        precheck=False,
        liner=liner)

    # Full minoverlap and maxoverlap Validation
    (minoverlap,
    maxoverlap,
    overlap_valid) = vp.get_parsed_range_info(
        minval=minoverlap,
        maxval=maxoverlap,
        range_field='  Overlap Length     ',
        range_unit='Base Pair(s) Fragment Overlap(s)',
        range_min=15 if not minhdist_valid else max(15, minhdist),
        range_max=splitlimit if splitlimit_valid else float('inf'),
        liner=liner)

    # Full outfile Validation
    outfile_valid = vp.get_outdf_validity(
        outdf=outfile,
        outdf_suffix='.oligopool.split.csv',
        outdf_field='   Output File       ',
        liner=liner)

    # Adjust outfile Suffix
    if not outfile is None:
        outfile = ut.get_adjusted_path(
            path=outfile,
            suffix='.oligopool.split.csv')

    # Optional Argument Parsing
    liner.send('\n Optional Arguments\n')

    # separate_outputs validation + display
    (separate_on,
    separate_valid) = vp.get_parsed_flag_info(
        flag=separate,
        flag_field=' Separate Outputs    ',
        flag_desc_off='Disabled (Single Combined Output)',
        flag_desc_on='Enabled (Separate Per-Split Outputs)',
        liner=liner)
    separate = separate_on

    # Validate random_seed (do not auto-generate)
    (random_seed,
    seed_valid) = vp.get_parsed_random_seed_info(
        random_seed=random_seed,
        random_seed_field='   Random Seed       ',
        liner=liner)

    # First Pass Validation
    if not all([
        indata_valid,
        splitlimit_valid,
        tmelt_valid,
        minhdist_valid,
        overlap_valid,
        outfile_valid,
        separate_valid,
        seed_valid]):
        liner.send('\n')
        raise RuntimeError(
            'Invalid Argument Input(s).')

    # Local RNG
    rng = np.random.default_rng(random_seed)

    # Start Timer
    t0 = tt.time()

    # Adjust Numeric Paramters
    splitlimit = round(splitlimit)
    minhdist   = round(minhdist)
    minoverlap = round(minoverlap)
    maxoverlap = round(maxoverlap)

    # Primer Design Book-keeping
    outdf = None
    stats = None
    warns = {}

    # Parse Oligopool Split Limit Feasibility
    liner.send('\n[Step 1: Parsing Split Limit]\n')

    # Parse splitlimit
    (parsestatus,
    seqlist,
    oligounderflow,
    unevensplit,
    minoligolen,
    maxoligolen,
    minsplitcount,
    maxsplitcount) = cs.get_parsed_splitlimit(
        indf=indf,
        splitlimit=splitlimit,
        liner=liner)

    # splitlimit infeasible
    if not parsestatus:

        # Prepare stats
        stats = {
            'status'  : False,
            'basis'   : 'infeasible',
            'step'    : 1,
            'step_name': 'parsing-split-limit',
	            'vars'    : {
                    'split_limit'     : splitlimit,
                    'oligo_underflow' : oligounderflow,
                    'uneven_split'    : unevensplit,
                    'min_oligo_length': minoligolen,
                    'max_oligo_length': maxoligolen,
                    'min_split_count' : minsplitcount,
                    'max_split_count' : maxsplitcount,},
	            'warns'   : warns}
        stats['random_seed'] = random_seed
        stats['separate_outputs'] = separate
        stats = ut.stamp_stats(
            stats=stats,
            module='split',
            input_rows=input_rows,
            output_rows=0)

        # Return results
        liner.close()
        return (outdf, stats)

    # Define spanlen
    # Note: spanlen is the minimum span
    #       for a split region
    spanlen = max(minhdist, minoverlap)

    # Compute Sequence Matrix
    liner.send('\n[Step 2: Computing Sequence Matrix]\n')

    # Compute padvec and seqmat
    (padvec,
    seqmat) = cs.get_seqmat_padvec(
        seqlist=seqlist,
        maxoligolen=maxoligolen,
        liner=liner,
        rng=rng)

    # Compute Sequence Matrix
    liner.send('\n[Step 3: Computing Entropy Vector]\n')

    # Compute padvec and seqmat
    entvec = cs.get_entvec(
        seqmat=seqmat,
        maxoligolen=maxoligolen,
        liner=liner)

    # Parse Oligopool Limit Feasibility
    liner.send('\n[Step 4: Parsing Variable Contigs]\n')

    # Parse splitlimit
    (parsestatus,
    varcont,
    varcontcount,
    mergedcontcount,
    filtercontcount) = cs.get_varcont(
        entvec=entvec,
        minhdist=minhdist,
        spanlen=spanlen,
        liner=liner)

    # splitlimit infeasible
    if not parsestatus:

        # Prepare stats
        stats = {
            'status'  : False,
            'basis'   : 'infeasible',
            'step'    : 4,
            'step_name': 'parsing-variable-contig',
            'vars'    : {
                'variable_contig_count': varcontcount,
                'merged_contig_count'  : mergedcontcount,
                'filtered_contig_count': filtercontcount},
            'warns'   : warns}
        stats['random_seed'] = random_seed
        stats['separate_outputs'] = separate
        stats = ut.stamp_stats(
            stats=stats,
            module='split',
            input_rows=input_rows,
            output_rows=0)

        # Return results
        liner.close()
        return (outdf, stats)

    # Launching Split Design
    liner.send('\n[Step 5: Computing Split]\n')

    # Define Split Design Stats
    stats = {
        'status'  : False,
        'basis'   : 'unsolved',
        'step'    : 8,
        'step_name': 'computing-split',
        'vars'    : {
            'split_count'    : 0,   # Total Number of Splits
            'split_lengths'  : [],  # Split Oligo Lengths
            'overlap_lengths': [],  # Split Overlap Lengths
            'mean_melting_temperature_distribution': [],  # Mean Tm of Each Split
            'mean_distance_distribution': [],     # Mean HDist of Each Split
            'has_infeasible_contigs'    : False,  # Infeasible Contigs Flag
            'has_uneven_splits'         : False}, # Uneven Splits Flag
        'warns'   : warns}
    stats['random_seed'] = random_seed
    stats['separate_outputs'] = separate

    # Schedule outfile deletion
    ofdeletion = ae.register(
        ut.remove_file,
        outfile)

    # Design Split
    (split,
    overlap,
    stats) = cs.split_engine(
        seqlist=seqlist,
        splitlimit=splitlimit,
        mintmelt=mintmelt,
        minhdist=minhdist,
        maxoverlap=maxoverlap,
        minoligolen=minoligolen,
        maxoligolen=maxoligolen,
        spanlen=spanlen,
        seqmat=seqmat,
        varcont=varcont,
        stats=stats,
        liner=liner)

    # Success Relevant Stats
    if stats['status']:

        # Launching Stats Aggregation
        liner.send('\n[Step 6: Aggregating Stats]\n')

        # Compute Tm and HDist Distribution
        # and Finalize Split Sequences
        (splitstore,
        stats) = cs.aggregate_stats(
            seqlist=seqlist,
            seqmat=seqmat,
            split=split,
            overlap=overlap,
            stats=stats,
            liner=liner,
            rng=rng)

    # Split Status
    if stats['status']:
        splitstatus = 'Successful'
    else:
        splitstatus = 'Failed'

    # Insert split into outdf
    if stats['status']:

        # Prepare outdf
        outdf = pd.DataFrame(
            index=indf.index)

        # Insert Splits into Columns
        for sidx in range(len(split)):
            outdf['Split{}'.format(sidx+1)] = splitstore[sidx]

        # Write outdf to file(s)
        if outfile is not None:
            cs.write_split_df_csv(
                df=outdf,
                outfile=outfile,
                separate_outputs=separate,
                sep=',')

    # Split Design Statistics
    liner.send('\n[Split Design Statistics]\n')

    liner.send(
        '     Design Status  : {}\n'.format(
            splitstatus))

    # Success Relevant Stats
    if stats['status']:

        maxval = max(max(min(stats['vars'][field]) for field in (
            'split_lengths',
            'overlap_lengths',
            'mean_melting_temperature_distribution',
            'mean_distance_distribution')),
            stats['vars']['split_count'])

        sntn, plen = ut.get_notelen(
            printlen=ut.get_printlen(
                value=maxval))

        liner.send(
            '     No. of Splits  : {:{},d} Fragments per Variant\n'.format(
                stats['vars']['split_count'],
                plen))
        liner.send(
            '      Split Length  : {:{},d} {}Base Pair(s)\n'.format(
                min(stats['vars']['split_lengths']),
                plen,
                ['', 'to {:,d} '.format(max(stats['vars']['split_lengths']))][
                    max(stats['vars']['split_lengths']) != min(stats['vars']['split_lengths'])
                ]))
        liner.send(
            '    Overlap Length  : {:{},d} {}Base Pair(s)\n'.format(
                min(stats['vars']['overlap_lengths']),
                plen,
                ['', 'to {:,d} '.format(max(stats['vars']['overlap_lengths']))][
                    max(stats['vars']['overlap_lengths']) != min(stats['vars']['overlap_lengths'])
                ]))
        liner.send(
            '    Overlap Tm      : {:{},d} {}°C\n'.format(
                min(stats['vars']['mean_melting_temperature_distribution']),
                plen,
                ['', 'to {:,d} '.format(max(stats['vars']['mean_melting_temperature_distribution']))][
                    max(stats['vars']['mean_melting_temperature_distribution']) != min(stats['vars']['mean_melting_temperature_distribution'])
                ]))
        liner.send(
            '    Overlap Distance: {:{},d} {}Mismatch(es)\n'.format(
               min(stats['vars']['mean_distance_distribution']),
                plen,
                ['', 'to {:,d} '.format(max(stats['vars']['mean_distance_distribution']))][
                    max(stats['vars']['mean_distance_distribution']) != min(stats['vars']['mean_distance_distribution'])
                ]))

    # Failure Relavant Stats
    else:
        liner.send(
            ' Infeasible Contigs : {}\n'.format(
                ['No', 'Yes'][stats['vars']['has_infeasible_contigs']]))
        liner.send(
            '      Unven Split   : {}\n'.format(
                ['No', 'Yes'][stats['vars']['has_uneven_splits']]))

    # Show Time Elapsed
    liner.send(
        ' Time Elapsed: {:.2f} sec\n'.format(
            tt.time()-t0))

    # Unschedule outfile deletion
    if splitstatus == 'Successful':
        ae.unregister(ofdeletion)

    # Close Liner
    liner.close()

    # Return Solution and Statistics
    stats = ut.stamp_stats(
        stats=stats,
        module='split',
        input_rows=input_rows,
        output_rows=len(outdf.index) if outdf is not None else 0)

    # Prepare return value based on separate_outputs flag
    if separate and (outdf is not None):
        # Return list of DataFrames, one per split column
        outdf_list = []
        for col in outdf.columns:
            split_df = outdf[[col]].copy()
            if not id_from_index:
                split_df = ut.get_df_with_id_column(split_df)
            outdf_list.append(split_df)
        return (outdf_list, stats)
    else:
        # Return single combined DataFrame
        outdf_return = outdf
        if (outdf is not None) and (not id_from_index):
            outdf_return = ut.get_df_with_id_column(outdf)
        return (outdf_return, stats)
