#
# Copyright BrainPad Inc. All Rights Reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
import bz2
import csv
import gzip
import os
import re
import shutil
import tarfile
import tempfile
import zipfile
from typing import Literal

import pandas

from cliboa.adapter.file import File
from cliboa.scenario.file import FileRead, FileWrite
from cliboa.util.base import _warn_deprecated_args
from cliboa.util.date import DateUtil
from cliboa.util.exception import CliboaException, InvalidParameter


class FileBaseTransform(FileRead, FileWrite):
    """
    Base class of file transform classes

    Basically transform class of Cliboa is that find files,
    do something, and output transformed files.
    When output files are created, name of the files will be the same name with the input files
    (original input files will be changed to the transformed files).
    If you would not like to remove the original files,
    give a path to the "dest_dir" for output directory.
    (If a non-existent directory path is specified, the directory is automatically created.)

    Note:
    Output files are not always be the same name with the input file names.
    See documentation for individual classes for details.
    """

    class Arguments(FileRead.Arguments, FileWrite.Arguments):
        force_continue: bool = False

        def resolve_dest_dir(self) -> str:
            if self.dest_dir:
                os.makedirs(self.dest_dir, exist_ok=True)
                return self.dest_dir
            else:
                return self.src_dir

    @property
    @_warn_deprecated_args("3.0", "4.0")
    def _force_continue(self):
        return self.args.force_continue

    def execute(self):
        pass

    def check_output_path(self, input_path, ext):
        root, name = os.path.split(input_path)

        if ext:
            if ext.startswith("."):
                output_name = os.path.splitext(name)[0] + ext
            else:
                output_name = os.path.splitext(name)[0] + "." + ext
        else:
            output_name = name

        if self.args.dest_dir:
            os.makedirs(self.args.dest_dir, exist_ok=True)
            output_dir = self.args.dest_dir
        else:
            output_dir = root

        output_path = os.path.join(output_dir, output_name)

        fd, temp_file = tempfile.mkstemp()
        os.close(fd)
        return output_path, temp_file

    def overwrite_output_path(self, input_path, output_path, temp_file):
        if input_path == output_path:
            os.remove(input_path)
        shutil.move(temp_file, output_path)

    def io_files(self, iterable, ext=None, func=None):
        """
        Iteration which returns input and output path.
        If the parameter "dest_dir" was given, the output file will be created under
        the given directory and returns input and output path.
        If not, the output file will be created to the same directory to the input file.

        Arguments:
            iterable (list): Input file list
            ext=None (str): Set an extension for output file,
                            if input and output extension would like to be changed.
                            "." is not necessary.
        """
        for input_path in iterable:
            output_path, temp_file = self.check_output_path(input_path, ext)

            try:
                func(input_path, temp_file)
            except Exception as e:
                if self.args.force_continue is True:
                    self.handle_error(e, input_path)
                else:
                    raise e

            self.overwrite_output_path(input_path, output_path, temp_file)

    def io_writers(self, iterable, mode="t", encoding="utf-8", ext=None):
        """
        Iteration which returns input and output writer.
        If the parameter "dest_dir" was given, the output file will be created under
        the given directory and returns input and output writer.
        If not, the output file will be created to the same directory to the input file.

        Arguments:
            iterable (list): Input file list
            mode="t" (str): Mode to open input files, opening file with text mode by default.
            encoding="utf-8" (str): File encoding. It will be ignored when mode is "b"
            ext=None (str): Set an extension for output file,
                            if input and output extension would like to be changed.
                            "." is not necessary.
        """
        if not "t" == mode and not "b" == mode:
            raise InvalidParameter("Unknown mode. One of the following is allowed [t, b]")

        if "b" in mode:
            encoding = None

        for input_path in iterable:
            output_path, temp_file = self.check_output_path(input_path, ext)

            if mode == "t":
                with (
                    open(input_path, mode="r", encoding=encoding, newline="") as i,
                    open(temp_file, mode="w", encoding=encoding, newline="") as o,
                ):
                    yield i, o
            elif mode == "b":
                with open(input_path, mode="rb") as i, open(temp_file, mode="wb") as o:
                    yield i, o

            self.overwrite_output_path(input_path, output_path, temp_file)

    def handle_error(self, e: Exception, input_path: str):
        # Please implement in a subclass if you would like to do something.
        self.logger.warning(
            (
                e,
                input_path,
            )
        )


class FileDecompress(FileBaseTransform):
    """
    Decompress the specified file
    """

    class Arguments(FileBaseTransform.Arguments):
        chunk_size: int | None = None
        password: str | None = None

    def execute(self, *args):
        files = self.get_src_files()
        self.check_file_existence(files)

        for f in files:
            _, ext = os.path.splitext(f)
            if ext == ".zip":
                self.logger.info("Decompress zip file %s" % f)
                if self.args.password is not None:
                    pwd = self.args.password.encode(self.args.encoding)
                else:
                    pwd = self.args.password
                with zipfile.ZipFile(f) as zp:
                    zp.extractall(self.args.resolve_dest_dir(), pwd=pwd),  # nosec
            elif ext == ".tar":
                self.logger.info("Decompress tar file %s" % f)
                with tarfile.open(f, "r:*") as tf:
                    tf.extractall(self.args.resolve_dest_dir())  # nosec
            elif ext == ".bz2":
                self.logger.info("Decompress bz2 file %s" % f)
                dcom_name = os.path.splitext(os.path.basename(f))[0]
                decom_path = os.path.join(self.args.resolve_dest_dir(), dcom_name)
                with bz2.open(f, mode="rb") as i, open(decom_path, mode="wb") as o:
                    while True:
                        buf = i.read(self.args.chunk_size)
                        if buf == b"":
                            break
                        o.write(buf)
            elif ext == ".gz":
                self.logger.info("Decompress gz file %s" % f)
                dcom_name = os.path.splitext(os.path.basename(f))[0]
                decom_path = os.path.join(self.args.resolve_dest_dir(), dcom_name)
                with gzip.open(f, "rb") as i, open(decom_path, "wb") as o:
                    while True:
                        buf = i.read(self.args.chunk_size)
                        if buf == b"":
                            break
                        o.write(buf)
            else:
                raise CliboaException("Unmatched any available decompress type %s" % f)


class FileCompress(FileBaseTransform):
    """
    Compress files
    """

    class Arguments(FileBaseTransform.Arguments):
        format: Literal["zip", "gz", "gzip", "bz2", "bzip2"]
        chunk_size: int = 1048576

    def execute(self, *args):
        files = self.get_src_files()
        self.check_file_existence(files)

        dest_dir = self.args.resolve_dest_dir()
        for f in files:
            if self.args.format == "zip":
                self.logger.info("Compress file %s to zip." % f)
                with zipfile.ZipFile(
                    os.path.join(dest_dir, (os.path.basename(f) + ".zip")),
                    "w",
                    zipfile.ZIP_DEFLATED,
                ) as o:
                    o.write(f, arcname=os.path.basename(f))
            elif self.args.format in ("gz", "gzip"):
                self.logger.info("Compress file %s to gzip." % f)
                com_path = os.path.join(dest_dir, (os.path.basename(f) + ".gz"))
                with open(f, "rb") as i, gzip.open(com_path, "wb") as o:
                    while True:
                        buf = i.read(self.args.chunk_size)
                        if buf == b"":
                            break
                        o.write(buf)
            elif self.args.format in ("bz2", "bzip2"):
                self.logger.info("Compress file %s to bzip2." % f)
                com_path = os.path.join(dest_dir, (os.path.basename(f) + ".bz2"))
                with open(f, "rb") as i, bz2.open(com_path, "wb") as o:
                    while True:
                        buf = i.read(self.args.chunk_size)
                        if buf == b"":
                            break
                        o.write(buf)


class DateFormatConvert(FileBaseTransform):
    """
    Convert csv (tsv) date field columns to another date field format columns
    """

    class Arguments(FileBaseTransform.Arguments):
        columns: list[str]
        formatter: str

    def execute(self, *args):
        files = self.get_src_files()
        self.check_file_existence(files)

        _, ext = os.path.splitext(files[0])
        if ext == ".csv":
            delimiter = ","
        elif ext == ".tsv":
            delimiter = "\t"

        for ins, ous in self.io_writers(files, encoding=self.args.encoding):
            reader = csv.DictReader(ins, delimiter=delimiter)
            writer = csv.DictWriter(ous, reader.fieldnames)
            writer.writeheader()
            date_util = DateUtil()
            for row in reader:
                for column in self.args.columns:
                    r = row.get(column)
                    if not r:
                        continue
                    row[column] = date_util.convert_date_format(r, self.args.formatter)
                writer.writerow(row)


class ExcelConvert(FileBaseTransform):
    """
    Convert excel to other format
    """

    class Arguments(FileBaseTransform.Arguments):
        pass

    def execute(self, *args):
        files = self.get_src_files()
        self.check_file_existence(files)

        # TODO Currently only excel to csv is supported.
        self.io_files(files, ext="csv", func=self.convert)

    def convert(self, fi, fo):
        self.logger.info("Convert %s to %s" % (fi, fo))
        df = pandas.read_excel(fi)
        df.to_csv(fo, encoding=self.args.encoding)


class FileCopy(FileBaseTransform):
    """
    Copy the file
    """

    class Arguments(FileBaseTransform.Arguments):
        dest_dir: str

    def execute(self, *args):
        files = self.get_src_files()
        self.check_file_existence(files)

        dest_dir = self.args.resolve_dest_dir()
        for file in files:
            shutil.copy(file, dest_dir)


class FileDivide(FileBaseTransform):
    """
    Divide a file to plural files
    """

    class Arguments(FileBaseTransform.Arguments):
        divide_rows: int
        header: bool = False
        suffix_pattern: str = ".%d"

    def divide_rows(self, divide_rows):
        self._divide_rows = divide_rows

    def header(self, header):
        self._header = header

    def suffix_pattern(self, suffix_pattern):
        self._suffix_pattern = suffix_pattern

    def execute(self, *args):
        files = self.get_src_files()
        self.check_file_existence(files)

        for file in files:
            fname = os.path.basename(file)

            px = ""
            if fname.startswith("."):
                fname = fname[1:]
                px = "."

            if "." in fname:
                nameonly, ext = fname.split(".", 1)
                ext = "." + ext
            else:
                nameonly = fname
                ext = ""

            if self.args.header:
                with open(file, encoding=self.args.encoding) as i:
                    self._header_row = i.readline()

            row = self._ifile_reader(file)
            newfilename = px + nameonly + self.args.suffix_pattern + ext

            dest_dir = self.args.resolve_dest_dir()
            has_left = True
            index = 1
            while has_left:
                ofile_path = os.path.join(dest_dir, newfilename % index)
                has_left = self._ofile_generator(ofile_path, row)
                index = index + 1

    def _ifile_reader(self, filepath):
        with open(filepath, encoding=self.args.encoding) as i:
            if self.args.header is True:
                i.readline()
            for line in i:
                yield line

    def _ofile_generator(self, filepath, row):
        left = False
        written = False
        with open(filepath, mode="w", encoding=self.args.encoding) as o:
            if self.args.header is True:
                o.write(self._header_row)
            for i, line in enumerate(row):
                written = True
                o.write(line)
                if i + 1 >= self.args.divide_rows:
                    left = True
                    break
        if written is False:
            os.remove(filepath)
        return left


class FileRename(FileBaseTransform):
    """
    Change file names with adding either prefix or suffix.
    """

    class Arguments(FileBaseTransform.Arguments):
        prefix: str = ""
        suffix: str = ""
        regex_pattern: str | None = None
        rep_str: str | None = None
        ext: str = ""
        without_ext: bool = True

    def prefix(self, prefix):
        self._prefix = prefix

    def suffix(self, suffix):
        self._suffix = suffix

    def regex_pattern(self, regex_pattern):
        self._regex_pattern = regex_pattern

    def rep_str(self, rep_str):
        self._rep_str = rep_str

    def ext(self, ext):
        self._ext = ext

    def without_ext(self, without_ext):
        self._without_ext = without_ext

    def execute(self, *args):
        files = self.get_src_files()
        self.check_file_existence(files)
        if self.args.rep_str is None and self.args.regex_pattern is not None:
            raise InvalidParameter("Replace string is not defined in yaml file: rep_str")
        if self.args.rep_str is not None and self.args.regex_pattern is None:
            raise InvalidParameter(
                "The conversion pattern is not defined in yaml file: regex_pattern"
            )

        for file in files:
            dirname = os.path.dirname(file)
            basename = os.path.basename(file)

            if self.args.without_ext is False:
                if self.args.prefix != "" or self.args.suffix != "" or self.args.ext != "":
                    raise InvalidParameter("Cannot be specified in without_ext mode")
                self._replace(
                    dirname,
                    file,
                    os.path.join(
                        dirname, re.sub(self.args.regex_pattern, self.args.rep_str, basename)
                    ),
                )
            else:
                px = ""
                if basename.startswith("."):
                    basename = basename[1:]
                    px = "."

                if "." in basename:
                    nameonly, extension = basename.split(".", 1)
                    extension = "." + extension
                else:
                    nameonly = basename
                    extension = ""

                if self.args.regex_pattern is not None and self.args.rep_str is not None:
                    nameonly = re.sub(self.args.regex_pattern, self.args.rep_str, nameonly)

                if self.args.ext:
                    extension = "." + self.args.ext

                newfilename = self.args.prefix + px + nameonly + self.args.suffix + extension
                self._replace(dirname, file, newfilename)

    def _replace(self, dirname, file, newfilename):
        newfilepath = os.path.join(dirname, newfilename)
        os.rename(file, newfilepath)
        self.logger.info("File name changed %s -> %s" % (file, newfilepath))


class FileConvert(FileBaseTransform):
    """
    Convert file encoding
    """

    class Arguments(FileBaseTransform.Arguments):
        encoding_from: str
        encoding_to: str
        errors: str | None = None

    def execute(self, *args):
        files = self.get_src_files()
        self.check_file_existence(files)

        self.io_files(files, func=self.convert)

    def convert(self, fi, fo):
        File().convert_encoding(
            fi,
            fo,
            self.args.encoding_from,
            self.args.encoding_to,
            self.args.errors,
        )

        self.logger.info("Encoded file %s" % fi)


class FileArchive(FileBaseTransform):
    """
    Create archeve object.a
    """

    class Arguments(FileBaseTransform.Arguments):
        dest_name: str
        format: Literal["zip", "tar"]
        create_dir: bool = False

    def execute(self, *args):
        files = self.get_src_files()
        self.check_file_existence(files)

        dest_dir = self.args.resolve_dest_dir()

        dest_path = os.path.join(dest_dir, (self.args.dest_name + ".%s" % self.args.format))

        if self.args.format == "tar":
            with tarfile.open(dest_path, "w") as tar:
                for file in files:
                    arcname = (
                        os.path.join(self.args.dest_name, os.path.basename(file))
                        if self.args.create_dir
                        else os.path.basename(file)
                    )
                    tar.add(file, arcname=arcname)
        elif self.args.format == "zip":
            with zipfile.ZipFile(dest_path, "w") as zp:
                for file in files:
                    arcname = (
                        os.path.join(self.args.dest_name, os.path.basename(file))
                        if self.args.create_dir
                        else os.path.basename(file)
                    )
                    zp.write(file, arcname=arcname)
        else:
            raise InvalidParameter("'format' must set one of the followings [tar, zip]")
