import argparse
import os
import sys
from datetime import datetime
from multiprocessing import Pool
from pathlib import Path

import nltk
import orjson
from nltk.stem import WordNetLemmatizer
from tqdm.notebook import tqdm

from ngram_tools.helpers.file_handler import FileHandler

nltk.download('wordnet', quiet=True)


def construct_output_path(input_file, output_dir, compress):
    """
    Construct the output file path based on the input file path and compression.

    Args:
        input_file (str): Path to the input file.
        output_dir (str): Directory where the output will be saved.
        compress (bool): Whether the file should be compressed (lz4).

    Returns:
        str: The constructed output file path.
    """
    input_path = Path(input_file)
    # If the file has a ".lz4" suffix, remove it before constructing the base name.
    base_name = (
        input_path.stem if input_path.suffix == '.lz4' else input_path.name
    )
    return str(Path(output_dir) / (base_name + ('.lz4' if compress else '')))


def set_info(proj_dir, ngram_size, file_range, compress):
    """
    Validate and set up file paths and metadata needed for lemmatizing ngrams.

    Args:
        proj_dir (str): Project base directory path.
        ngram_size (int): The size of the ngrams (1-5).
        file_range (tuple or list of int, optional): Range of files to process.
        compress (bool): Whether output files should be compressed (lz4).

    Returns:
        tuple: Contains:
            - input_dir (str): The input directory path.
            - output_dir (str): The output directory path.
            - file_range (tuple): Updated file range.
            - num_files_available (int): Count of all available input files.
            - input_paths_use (list[str]): The specific input file paths to use.
            - file_range (tuple): Same file range passed in (for printing).
            - num_files_to_use (int): Number of files that will actually be used.
            - first_file (str): The first file in the chosen range.
            - last_file (str): The last file in the chosen range.
            - output_paths (list[str]): Corresponding output file paths.
    """
    input_dir = os.path.join(proj_dir, f'{ngram_size}gram_files/3lowercase')
    output_dir = os.path.join(proj_dir, f'{ngram_size}gram_files/4lemmatize')

    if not os.path.isdir(input_dir):
        raise NotADirectoryError(
            f"Input directory {input_dir} does not exist or isn't a directory."
        )

    input_paths_elig = sorted(
        [entry.path for entry in os.scandir(input_dir) if entry.is_file()]
    )
    num_files_available = len(input_paths_elig)

    if not file_range:
        file_range = (0, len(input_paths_elig) - 1)

    input_paths_use = input_paths_elig[file_range[0]: file_range[1] + 1]
    num_files_to_use = len(input_paths_use)

    first_file = input_paths_use[0]
    last_file = input_paths_use[-1]

    output_paths = sorted([
        construct_output_path(file, output_dir, compress)
        for file in input_paths_use
    ])

    return (
        input_dir, output_dir, file_range, num_files_available,
        input_paths_use, file_range, num_files_to_use, first_file,
        last_file, output_paths
    )


def print_info(input_dir, output_dir, file_range, num_files_available,
               num_files_to_use, first_file, last_file, ngram_size, workers,
               compress, overwrite, start_time, delete_input):
    """
    Print a summary of the lemamtizing process settings for the user.

    Args:
        input_dir (str): The input directory path.
        output_dir (str): The output directory path.
        file_range (tuple): Range of file indices to process.
        num_files_available (int): Count of all available input files.
        num_files_to_use (int): Number of input files to process.
        first_file (str): First file in the selected range.
        last_file (str): Last file in the selected range.
        ngram_size (int): The size of the ngrams (1-5).
        workers (int): Number of worker processes to spawn.
        compress (bool): Whether output files should be compressed (lz4).
        overwrite (bool): Whether to overwrite existing files.
        start_time (datetime): The time the process started.
        delete_input (bool): Whether input files should be deleted.
    """
    print(f'\033[31mStart Time:                {start_time}\n\033[0m')
    print('\033[4mLemmatizing Info\033[0m')
    print(f'Input directory:           {input_dir}')
    print(f'Output directory:          {output_dir}')
    print(f'File index range:          {file_range[0]} to {file_range[1]}')
    print(f'Files available:           {num_files_available}')
    print(f'Files to use:              {num_files_to_use}')
    print(f'First file to get:         {first_file}')
    print(f'Last file to get:          {last_file}')
    print(f'Ngram size:                {ngram_size}')
    print(f'Number of workers:         {workers}')
    print(f'Compress output files:     {compress}')
    print(f'Overwrite existing files:  {overwrite}')
    print(f'Delete input directory:    {delete_input}\n')


lemmatizer = WordNetLemmatizer()
google_to_wordnet = {
    "NOUN": "n",
    "VERB": "v",
    "ADJ": "a",
    "ADV": "r"
}


def process_a_line(line):
    """
    Process a single line of JSON by lemmatizing the 'ngram' tokens.

    Args:
        line (bytes or str): A line of JSON (possibly bytes).

    Returns:
        dict: The updated JSON object after lemmatization.
    """
    if isinstance(line, bytes):
        line = line.decode('utf-8')

    try:
        data = orjson.loads(line)
    except orjson.JSONDecodeError as exc:
        print(f"Failed to parse JSON line: {line[:100]}... Error: {exc}")
        raise

    data['ngram'] = {
        token: lemmatizer.lemmatize(
            word, google_to_wordnet.get(data['pos'][token], 'n')
        )
        for token, word in data['ngram'].items()
    }
    del data['pos']

    return data


def process_a_file(args):
    """
    Process a single input file, apply lemmatization to each line, and
    save the results to the output file.

    Args:
        args (tuple): Contains:
            in_handler (FileHandler): File handler for reading the input.
            out_handler (FileHandler): File handler for writing the output.
            overwrite (bool): Whether to overwrite existing output files.
    """
    in_handler, out_handler, overwrite = args

    if not overwrite and os.path.exists(out_handler.path):
        return

    try:
        with in_handler.open() as infile, out_handler.open() as outfile:
            for line_no, line in enumerate(infile, start=1):
                if isinstance(line, bytes):
                    line = line.decode('utf-8')

                json_dict = process_a_line(line)
                jsonl_line = out_handler.serialize(json_dict)
                outfile.write(jsonl_line)
    except Exception as exc:
        print(f"Error processing {in_handler.path}: {exc}")


def process_a_directory(output_dir, input_paths, output_paths, workers,
                        overwrite, compress):
    """
    Process all eligible input files in a directory with multiple workers.

    Args:
        output_dir (str): Path to the output directory.
        input_paths (list[str]): List of input file paths.
        output_paths (list[str]): List of output file paths.
        workers (int): Number of worker processes to use.
        overwrite (bool): Whether to overwrite existing output files.
        compress (bool): Whether output files should be compressed (lz4).
    """
    os.makedirs(output_dir, exist_ok=True)

    handlers = []
    for input_path, output_path in zip(input_paths, output_paths):
        input_ext = Path(input_path).suffix
        input_size = os.path.getsize(input_path)

        # Skip empty files.
        if (
            (input_ext == '.jsonl' and input_size == 0)
            or (input_ext == '.lz4' and input_size == 11)
        ):
            continue

        in_handler = FileHandler(input_path)
        out_handler = FileHandler(output_path, is_output=True, compress=compress)
        handlers.append((in_handler, out_handler))

    args = [
        (in_handler, out_handler, overwrite)
        for in_handler, out_handler in handlers
    ]

    with tqdm(total=len(handlers), desc='Lemmatizing', unit='files') as pbar:
        with Pool(processes=workers) as pool:
            for _ in pool.imap_unordered(process_a_file, args):
                pbar.update()


def clear_a_directory(directory_path):
    """
    Remove all files (and empty subdirectories) from the specified directory.

    Args:
        directory_path (str): Path to the directory to clear.
    """
    for entry in os.scandir(directory_path):
        if entry.is_file():
            os.remove(entry.path)
        elif entry.is_dir():
            os.rmdir(entry.path)


def lemmatize_ngrams(ngram_size, proj_dir, file_range=None, overwrite=False,
                     compress=False, workers=os.cpu_count(), delete_input=False):
    """
    Main function to drive the lemmatizing process. Reads the specified range of
    files from the project directory, applies lemmatization, and writes results.

    Args:
        ngram_size (int): The size of ngrams (1-5).
        proj_dir (str): The base directory path for the project.
        file_range (tuple[int], optional): Range of files to process.
        overwrite (bool, optional): Overwrite existing files.
        compress (bool, optional): Compress output files (lz4).
        workers (int, optional): Number of parallel processes to use.
        delete_input (bool, optional): Delete the input directory after processing.
    """
    start_time = datetime.now()

    (
        input_dir,
        output_dir,
        file_range,
        num_files_available,
        input_paths,
        file_range,
        num_files_to_use,
        first_file,
        last_file,
        output_paths
    ) = set_info(proj_dir, ngram_size, file_range, compress)

    print_info(
        input_dir, output_dir, file_range, num_files_available,
        num_files_to_use, first_file, last_file, ngram_size, workers,
        compress, overwrite, start_time, delete_input
    )

    process_a_directory(
        output_dir, input_paths, output_paths, workers, overwrite, compress
    )

    if delete_input:
        clear_a_directory(input_dir)

    end_time = datetime.now()
    print(f'\033[31m\nEnd Time:                  {end_time}\033[0m')

    total_runtime = end_time - start_time
    print(f'\033[31mTotal runtime:             {total_runtime}\n\033[0m')