"""Python wrapper for the mmpdblib code that is used in this package."""

import shlex
import subprocess
from collections.abc import Generator
from multiprocessing import cpu_count
from pathlib import Path
from typing import Any, Literal

import pandas as pd

from nonadditivity.utils.errors import FragmentationError, IndexingError
from nonadditivity.utils.log import get_logger

logger = get_logger()


def _run_mmpdb_command(command: str) -> int:
    """Run command in shell and pipe output to logger.

    Wrapper for the mmpdb commandline interface so it would run the
    commands using subprocess while streaming its output to the logger
    in realtime.

    Args:
        command (str): command to execute, will be formatted with shlex
        logger (logging.Logger): logger to write messages to.

    Returns:
        int: return code of the application.
    """

    def run_subprocess(
        command: str,
    ) -> Generator[tuple[str, int | None] | tuple[Literal[""], int], Any, None]:
        process = subprocess.Popen(
            shlex.split(command),
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
        )
        stdout = process.stdout.fileno()  # type: ignore
        with open(stdout, encoding="utf-8") as nice_stdout:
            for line in nice_stdout:
                yield line.replace("\n", " "), process.poll()

            yield "", process.wait()

    for line, returncode in run_subprocess(command=command):
        if returncode is not None:
            return returncode
        if line.isspace():
            continue
        logger.info(line)
    return 0


def run_mmpdb_fragment(
    cache: bool,
    fragments_file: Path,
    structure_file: Path,
) -> None:
    """Run `mmpdb fragment` command.

    Args:
        cache (bool): whether to use old fragmentation
        fragments_file (Path): path to fragmentation file
        structure_file (Path): path to input structure file
    """
    command = (
        f"mmpdb fragment --num-jobs {cpu_count() -2 } --delimiter tab "
        f"--output {fragments_file.as_posix()} {structure_file.as_posix()} "
    )
    if cache:
        command += f"--cache {fragments_file.as_posix()}"
    else:
        command += "-i smi --max-rotatable-bonds 20"

    if _run_mmpdb_command(command=command) != 0:
        raise FragmentationError(
            "Something happened during mmpdb fragmentation. Please check the logs!",
        )


def run_mmpdb_index(
    fragments_file: Path,
    mmp_outputfile: Path,
    max_variable_heavies: int,
) -> None:
    """Run `mmpdb index` command.

    Args:
        fragments_file (Path): fragmentation file
        mmp_outputfile (Path): output file
        max_variable_heavies (int): max heavy atom count to be exchanged in
        a mmp.
    """
    command = (
        f"mmpdb index --out csv --symmetric --output {mmp_outputfile.as_posix()} "
        f"--max-variable-heavies {max_variable_heavies} {fragments_file.as_posix()}"
    )
    if _run_mmpdb_command(command=command) != 0:
        raise IndexingError(
            "Something happened during mmpdb indexing. Please check the logs!",
        )


def read_raw_mmps(path: Path) -> pd.DataFrame:
    """Read in file generated by `mmpdb index`.

    Args:
        path (str): path to the file to be read in.

    Returns:
        pd.DataFrame: file content in a pandas dataframe.
    """
    mmp_dataframe = pd.read_table(
        filepath_or_buffer=path,
        sep="\t",
        header=None,
        names=[
            "SMILES_LHS",
            "SMILES_RHS",
            "ID_LHS",
            "ID_RHS",
            "TRANSFORMATION",
            "CONSTANT",
        ],
    )
    return mmp_dataframe


def run_mmpdlib_code(
    cache: bool,
    fragments_file: Path,
    mmp_outputfile: Path,
    structure_file: Path,
    max_variable_heavies: int,
) -> pd.DataFrame:
    """Run mmpdb fragment and index.

    Runs all the mmpdblib code in one and returns the matched molecular pairs as a list
    of strings. The strings contain SMILES and IDs of the whole molecules as well as
    fragment SMILES

    Args:
        cache (bool): wheter you want to use previously fragmented molecules.
        fragments_file (str): path to where the fragments are stored and
        where previously fragmented molecules would be stored.
        mmp_outputfile (str): path to write the matched molecular pair data.
        structure_file (str): input file containing smiles of the molecules
        you want to find mmps for.
        max_variable_heavies (int): max heavy atoms to be replaced in an mmp.

    Returns:
        DataFrame: Pandas Dataframe containing the mmpdb results.
    """
    if cache and not fragments_file.exists():
        logger.warning(
            "Was not able to locate results from previous fragmentation.",
        )
        logger.warning("Will redo all fragmentation.")
        cache = False

    run_mmpdb_fragment(
        cache=cache,
        fragments_file=fragments_file,
        structure_file=structure_file,
    )
    run_mmpdb_index(
        fragments_file=fragments_file,
        mmp_outputfile=mmp_outputfile,
        max_variable_heavies=max_variable_heavies,
    )
    return read_raw_mmps(path=mmp_outputfile)
