#!/usr/bin/env python
#
# Copyright (c) 2024, the cclib development team
#
# This file is part of cclib (http://cclib.github.io) and is distributed under
# the terms of the BSD 3-Clause License.
"""Script for writing data tables from computational chemistry files."""

import argparse
import os.path
import sys  # noqa: F401
from typing import Iterable, Optional

from cclib.io import ccframe, ccopen
from cclib.parser.utils import find_package

_has_pandas = find_package("pandas")
if _has_pandas:
    import pandas as pd


def process_logfiles(filenames: Iterable[str], output: Optional[str], identifier: str) -> None:
    df = ccframe([ccopen(path) for path in filenames])
    if output is not None:
        outputtype = os.path.splitext(os.path.basename(output))[1][1:]
        if not outputtype:
            raise RuntimeWarning(
                "The output type could not be determined from the given path, "
                "not writing DataFrame to disk"
            )

        if outputtype in {"csv"}:
            df.to_csv(output, mode="w")
        elif outputtype in {"h5", "hdf", "hdf5"}:
            df.to_hdf(output, mode="w", key=identifier)
        elif outputtype in {"json"}:
            df.to_json(output)
        elif outputtype in {"pickle", "pkl"}:
            df.to_pickle(output)
        elif outputtype in {"xlsx"}:
            writer = pd.ExcelWriter(output, mode="w")
            # This overwrites previous sheets
            # (see https://stackoverflow.com/a/42375263/4039050)
            df.to_excel(writer, sheet_name=identifier)
            writer.save()
    else:
        print(df)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-O",
        "--output",
        help=(
            "the output document to write, including an "
            "extension supported by pandas "
            "(csv, h5/hdf/hdf5, json, pickle/pkl, xlsx)"
        ),
    )
    parser.add_argument(
        "compchemlogfiles",
        metavar="compchemlogfile",
        nargs="+",
        help=("one or more computational chemistry output " "files to parse and convert"),
    )
    parser.add_argument(
        "--identifier",
        default="logfiles",
        help=(
            "name of sheet which will contain DataFrame, if "
            "writing to an Excel file, or identifier for "
            "the group in HDFStore, if writing a HDF file"
        ),
    )
    parser.add_argument(
        "-f",
        "--force",
        action="store_true",
        help=("overwrite output file in case it already exists"),
    )
    args = parser.parse_args()
    if args.output is not None and not args.force and os.path.exists(args.output):
        parser.exit(1, f'failure: exiting to avoid overwriting existing file "{args.output}"\n')

    process_logfiles(args.compchemlogfiles, args.output, args.identifier)


if __name__ == "__main__":
    main()
