#!/usr/bin/env python3
"""Command-line interface for raccoon toolkit."""

import sys
import os
import argparse
import logging
import shlex
from raccoon import __version__, _program
from raccoon.commands import alignment as alignment_cmd, phylo as phylo_cmd, combine as combine_cmd, mask as mask_cmd
from raccoon.utils import constants as rc

def build_parser():
    parser = argparse.ArgumentParser(prog=_program, description="raccoon toolkit")
    parser.add_argument('-v','--version', action='version', version=f'raccoon {__version__}', help='Show version and exit')
    parser.add_argument('-V','--verbose', action='count', default=0, help='Increase verbosity (can specify multiple times)')
    sub = parser.add_subparsers(dest='command')
    sub.required = True

    c = sub.add_parser('seq-qc', help='sequence QC (combine and harmonise FASTA files)')
    c.add_argument('-f','--fasta', nargs='+', help='Input FASTA files (one or more)')
    c.add_argument('-m','--metadata', nargs='+', action='extend', default=None, help='Input metadata CSV(s) for FASTA header population (one or more; optional)')
    c.add_argument('--metadata-delimiter', default=rc.DEFAULT_METADATA_DELIMITER, help=f'Metadata delimiter (default: {rc.DEFAULT_METADATA_DELIMITER}; auto-detects .csv and .tsv files)')
    c.add_argument('--metadata-id-field', default=rc.DEFAULT_ID_FIELD, help=f'Metadata id column used for matching sequence ID (default: {rc.DEFAULT_ID_FIELD})')
    c.add_argument('--metadata-location-field', default=rc.DEFAULT_LOCATION_FIELD, help=f'Metadata location column (default: {rc.DEFAULT_LOCATION_FIELD})')
    c.add_argument('--metadata-date-field', default=rc.DEFAULT_DATE_FIELD, help=f'Metadata date column (default: {rc.DEFAULT_DATE_FIELD})')
    c.add_argument('--seq-id-delimiter', default=rc.DEFAULT_ID_DELIMITER, help=f'Delimiter for parsing IDs from input headers (default: {rc.DEFAULT_ID_DELIMITER})')
    c.add_argument('--seq-id-field-index', type=int, default=rc.DEFAULT_ID_FIELD_INDEX, help=f'0-based field index for parsing IDs from input headers (default: {rc.DEFAULT_ID_FIELD_INDEX})')
    c.add_argument('--min-length', type=int, default=None, help='Minimum sequence length to keep')
    c.add_argument('--max-n-content', type=float, default=None, help='Maximum N content proportion to keep (e.g. 0.1)')
    c.add_argument('--header-fields', default=None, help=f'Template for constructing harmonised sequence header (e.g. "{rc.DEFAULT_HEADER_FIELDS}")')
    c.add_argument('--header-separator', default=rc.DEFAULT_HEADER_SEPARATOR, help=f'Header separator (default: {rc.DEFAULT_HEADER_SEPARATOR})')
    c.add_argument('-o', '--outfile', default=rc.DEFAULT_OUTPUT_FILE, help='Output FASTA file (use - for stdout)')
    c.set_defaults(func=combine_cmd.main)

    a = sub.add_parser('aln-qc', help='alignment QC')
    a.add_argument('alignment', help='Input alignment fasta file')
    a.add_argument('-t','--sequence-type', choices=['nt','aa'], default='nt', dest='sequence_type', help='Sequence type (default: nt)')
    a.add_argument('-d','--outdir', default='.', dest='outdir', help='Output directory')
    a.add_argument('--genbank', help='GenBank file for frame-breaking indel checks', dest='genbank', default=None)
    a.add_argument('--reference-id', help='Reference sequence ID in alignment (for GenBank features)', dest='reference_id', default=None)
    a.add_argument('--max-n-content', type=float, default=rc.DEFAULT_MAX_N_CONTENT, help=f'N content threshold for flagging (default: {rc.DEFAULT_MAX_N_CONTENT})')
    a.add_argument('--cluster-window', type=int, default=rc.DEFAULT_CLUSTER_WINDOW, dest='cluster_window', help=f'Window size for clustered SNP detection (default: {rc.DEFAULT_CLUSTER_WINDOW}bp)')
    a.add_argument('--cluster-count', type=int, default=rc.DEFAULT_CLUSTER_COUNT, dest='cluster_count', help=f'Min SNPs in window to flag as clustered (default: {rc.DEFAULT_CLUSTER_COUNT})')
    a.add_argument('--no-flag-clustered', action='store_true', dest='no_flag_clustered', help='Skip flagging clustered SNPs ')
    a.add_argument('--no-flag-n-adjacent', action='store_true', dest='no_flag_n_adjacent', help='Skip flagging SNPs adjacent to Ns ')
    a.add_argument('--no-flag-gap-adjacent', action='store_true',  dest='no_flag_gap_adjacent', help='Skip flagging SNPs adjacent to gaps ')
    a.add_argument('--no-flag-frame-break', action='store_true', dest='no_flag_frame_break', help='Skip flagging frame-breaking indels (GenBank file required)')
    a.add_argument('--flag-removal-threshold', type=int, default=rc.DEFAULT_FLAG_REMOVAL_THRESHOLD, dest='flag_removal_threshold', help=f'Flag sequences for removal based on criteria (default: {rc.DEFAULT_FLAG_REMOVAL_THRESHOLD} flagged sites)')
    a.set_defaults(func=alignment_cmd.main)

    m = sub.add_parser('mask', help='apply a mask file to an alignment')
    m.add_argument('alignment', help='Input alignment fasta file')
    m.add_argument('--mask-file', required=True, dest='mask_file', help='Mask CSV file from aln-qc')
    m.add_argument('--mask-character', default=rc.DEFAULT_MASK_CHARACTER, dest='mask_character', help=f'Character to use for masking (default: {rc.DEFAULT_MASK_CHARACTER})')
    m.add_argument('-o', '--outfile', dest='outfile', default=None, help='Output masked alignment file')
    m.add_argument('-d', '--outdir', dest='outdir', default='.', help='Output directory')
    m.add_argument('-t', '--sequence-type', choices=['nt','aa'], default='nt', dest='sequence_type', help='Sequence type (default: nt)')
    m.set_defaults(func=mask_cmd.main)

    p = sub.add_parser('tree-qc', help='phylogenetic QC')
    p.add_argument('-d', '--outdir', help='Output directory', dest='outdir', default='.')
    p.add_argument('-t','--tree', help='Phylogeny tree file or base name', dest='phylogeny', required=True)
    p.add_argument('--tree-format', choices=['auto','newick','nexus'], default='auto', dest='tree_format', help='Tree format (default: auto)')
    p.add_argument('--outgroup-ids', help='Comma-separated outgroup ids', dest='outgroup_ids', default=None)
    p.add_argument('--alignment', help='Alignment fasta used with ASR state file', dest='alignment', default=None)
    p.add_argument('--asr-state', help='Ancestral state reconstruction file in the format output by IQTREE', dest='asr_state', default=None)
    p.add_argument('--mask-file', help='Mask output file', dest='mask_file', default=None)
    p.add_argument('--assembly-refs', help='Assembly references fasta', dest='assembly_refs', default=None)
    p.add_argument('--long-branch-sd', type=float, default=3.0, dest='long_branch_sd', help='Std dev threshold for long branch flagging (default: 3.0)')
    p.add_argument('--adar-window', type=int, default=300, dest='adar_window', help='Max distance (bp) for ADAR cluster window (default: 300)')
    p.add_argument('--adar-min-count', type=int, default=3, dest='adar_min_count', help='Min ADAR sites in window to flag a branch (default: 3)')
    p.add_argument('--run-apobec', action='store_true', dest='run_apobec', help='Run APOBEC3 phylo checks')
    p.add_argument('--run-adar', action='store_true', dest='run_adar', help='Run ADAR phylo checks')
    p.add_argument('--tip-fields', default=rc.DEFAULT_HEADER_FIELDS, dest=rc.KEY_TIP_FIELDS, help=f'Template for parsing tip label fields for colouring (default: {rc.DEFAULT_HEADER_FIELDS})')
    p.add_argument('--tip-field-delimiter', default=rc.DEFAULT_HEADER_SEPARATOR, dest='tip_field_delimiter', help=f'Delimiter for parsing tip label fields (default: {rc.DEFAULT_HEADER_SEPARATOR})')
    p.add_argument('--tip-date-field', default=rc.DEFAULT_DATE_FIELD, dest='tip_date_field', help=f'Date field for parsing tip label fields (default: {rc.DEFAULT_DATE_FIELD})')
    p.add_argument('--midpoint-root', action='store_true', dest='midpoint_root', help='Midpoint-root tree for report visualization (applied only when --asr-state is not provided)')
    p.add_argument('--height', help='Figure height', dest='height', default=None)
    p.set_defaults(func=phylo_cmd.main)

    return parser


def main(sysargs=None):
    if sysargs is None:
        sysargs = sys.argv[1:]
    parser = build_parser()
    args = parser.parse_args(sysargs)
    args.input_cmd_line = f"{_program} {shlex.join(sysargs)}" if sysargs else _program

    # configure logging
    level = logging.WARNING
    if getattr(args,'verbose',0) >=2:
        level = logging.DEBUG
    elif getattr(args,'verbose',0) == 1:
        level = logging.INFO
    logging.basicConfig(level=level, format='%(levelname)s: %(message)s')

    try:
        return args.func(args)
    except Exception as e:
        logging.exception("Error running raccoon")
        return 2


if __name__ == '__main__':
    sys.exit(main())