import os
import logging
import argparse

def make_parser():
    """
    Construct the command line parser for the CLI.
    """
    parser = argparse.ArgumentParser()

    parser.add_argument('-v', '--verbose', action='store_true')

    subparsers = parser.add_subparsers(dest='command')

    model_default = 'klue/bert-base'
    classification_model_default = os.path.join('models', 'kotok_model')
    error_classification_model_default = os.path.join('models', 'kotok_error_model')
    spacing_classification_model_default = os.path.join('models', 'kotok_spacing_model')
    cache_default = os.path.join('cache')
    data_default = os.path.join('data', 'labeled.json')
    lemma_data_default = os.path.join('data', 'lemma')

    data_dl = subparsers.add_parser('data_dl')
    data_dl.add_argument('-o', '--out_dir', type=str, default=os.path.join('data', 'txt'), help='Output directory for the downloaded data')

    data = subparsers.add_parser('data')
    data.add_argument('-m', '--model', type=str, default=model_default, help='Pretrained model name or path for tokenization')
    data.add_argument('-c', '--cache', type=str, default=cache_default, help='Cache directory')
    data.add_argument('-i', '--input', type=str, default=os.path.join('data', 'txt'), help='Input data directory containing text files')
    data.add_argument('-n', '--normalize_mode', type=str, default=None, help='Unicode normalization mode')
    data.add_argument('-o', '--output', type=str, default=data_default, help='Output file path')
    data.add_argument('-s', '--split', type=float, default=0.8, help='Train-test split ratio')

    train = subparsers.add_parser('train')
    train.add_argument('-m', '--model', type=str, default=model_default, help='Pretrained model name or path for tokenization')
    train.add_argument('-c', '--cache', type=str, default=cache_default, help='Cache directory')
    train.add_argument('-d', '--data', type=str, default=data_default, help='Data file path, generated by the data command')
    train.add_argument('-o','--output', type=str, default=classification_model_default, help='Output directory for the trained model')
    train.add_argument('-l', '--logs', type=str, default='logs', help='Output directory for the logs')

    inference = subparsers.add_parser('inference')
    inference.add_argument('-cm', '--classification_model', type=str, default=classification_model_default, help='Classification model path, generated by the train command')
    inference.add_argument('-m', '--model', type=str, default=model_default, help='Pretrained model name or path for tokenization')
    inference.add_argument('-c', '--cache', type=str, default=cache_default, help='Cache directory')
    inference.add_argument('-f', '--format', type=str, default='pretty', help='Output format')
    inference.add_argument('-n', '--normalize_mode', type=str, default=None, help='Unicode normalization mode')
    inference.add_argument('-u', '--user_dict', type=str, default=None, help='User dictionary file or directory path')
    inference.add_argument('-ld', '--lemma_data', type=str, default=lemma_data_default, help='Lemmatization data directory')
    inference.add_argument('-nl', '--no_lemma', action='store_true', default=False, help='Disable lemmatization')
    inference.add_argument('-ne', '--no_error_correction', action='store_true', default=False, help='Disable error correction')
    inference.add_argument('-ecm', '--error_classification_model', type=str, default=error_classification_model_default, help='Error classification model path, generated by the train command')
    inference.add_argument('-em', '--error_model', type=str, default=model_default, help='Pretrained model name or path for error correction')
    inference.add_argument('-ns', '--no_spacing_correction', action='store_true', default=False, help='Disable spacing correction')
    inference.add_argument('-scm', '--spacing_classification_model', type=str, default=spacing_classification_model_default, help='Spacing classification model path, generated by the train command')
    inference.add_argument('-sm', '--spacing_model', type=str, default=model_default, help='Pretrained model name or path for spacing correction')

    lemmatize = subparsers.add_parser('lemmatize')
    lemmatize.add_argument('-d', '--data-dir', type=str, default=lemma_data_default, help='Lemmatization data directory')

    return parser

def main():
    logging.basicConfig(level=logging.INFO, format='[%(levelname)s] %(message)s')

    args = make_parser().parse_args()

    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    # Only import sub modules if needed to reduce startup time
    if args.command == 'data_dl':
        from .data_dl import data_dl
        data_dl(**args.__dict__)
    elif args.command == 'data':
        from .data import data
        data(**args.__dict__)
    elif args.command == 'train':
        from .train import train
        train(args)
    elif args.command == 'inference':
        from .inference import inference
        # If no error or spacing correction is needed, set the model to None to override the default model
        if args.no_error_correction:
            args.error_model = None
            args.error_classification_model = None
        if args.no_spacing_correction:
            args.spacing_model = None
            args.spacing_classification_model = None    
        inference(**args.__dict__)
    elif args.command == 'lemmatize':
        from .lemmatize import lemmatize
        lemmatize(**args.__dict__)


if __name__ == '__main__':
    main()
