# pyright: strictDictionaryInference=true
#
# Copyright 2003-2009 Zuza Software Foundation
#
# This file is part of the Translate Toolkit.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <http://www.gnu.org/licenses/>.

"""Count strings and words for supported localization files.

These include: XLIFF, TMX, Gettex PO and MO, Qt .ts and .qm, Wordfast TM, etc

See: http://docs.translatehouse.org/projects/translate-toolkit/en/latest/commands/pocount.html
for examples and usage instructions.
"""
from __future__ import annotations

import csv
import logging
import os
import re
import sys
from argparse import ArgumentParser
from collections import defaultdict
from dataclasses import dataclass
from operator import itemgetter

from translate.lang.common import Common
from translate.misc.multistring import multistring
from translate.storage import factory
from translate.storage.workflow import StateEnum

try:
    from functools import cached_property
except ImportError:
    # Fix for python 3.7
    cached_property = property

extended_state_strings = {
    StateEnum.EMPTY: "empty",
    StateEnum.NEEDS_WORK: "needs-work",
    StateEnum.REJECTED: "rejected",
    StateEnum.NEEDS_REVIEW: "needs-review",
    StateEnum.UNREVIEWED: "unreviewed",
    StateEnum.FINAL: "final",
}

UNTRANSLATED = StateEnum.EMPTY
FUZZY = StateEnum.NEEDS_WORK
TRANSLATED = StateEnum.UNREVIEWED

state_strings = {
    UNTRANSLATED: "untranslated",
    FUZZY: "fuzzy",
    TRANSLATED: "translated",
}

logger = logging.getLogger(__name__)

# kdepluralre = re.compile("^_n: ") #Restore this if you really need support for old kdeplurals
brtagre = re.compile(r"<br\s*?/?>")
# xmltagre is a direct copy of the from placeables/general.py
xmltagre = re.compile(
    r"""
        <                         # start of opening tag
        ([\w.:]+)                 # tag name, possibly namespaced
        (\s([\w.:]+=              # space and attribute name followed by =
            ((".*?")|('.*?'))     # attribute value, single or double quoted
        )?)*/?>                   # end of opening tag, possibly self closing
        |</([\w.]+)>              # or a closing tag
        """,
    re.VERBOSE,
)
numberre = re.compile("\\D\\.\\D")


class ConsoleColor:
    """Class to implement color mode."""

    # print using color? Default to true
    color_mode = True
    COLOR_PURPLE = "\033[95m"
    COLOR_GREEN = "\033[92m"
    COLOR_YELLOW = "\033[93m"
    COLOR_RED = "\033[91m"
    COLOR_DEFAULT = "\033[0m"

    @classmethod
    def HEADER(cls):
        return cls.COLOR_PURPLE if ConsoleColor.color_mode else ""

    @classmethod
    def OKGREEN(cls):
        return cls.COLOR_GREEN if ConsoleColor.color_mode else ""

    @classmethod
    def WARNING(cls):
        return cls.COLOR_YELLOW if ConsoleColor.color_mode else ""

    @classmethod
    def FAIL(cls):
        return cls.COLOR_RED if ConsoleColor.color_mode else ""

    @classmethod
    def ENDC(cls):
        return cls.COLOR_DEFAULT if ConsoleColor.color_mode else ""


def wordcount(string):
    # TODO: po class should understand KDE style plurals ##
    # string = kdepluralre.sub("", string) #Restore this if you really need support for old kdeplurals
    string = brtagre.sub("\n", string)
    string = xmltagre.sub("", string)
    string = numberre.sub(" ", string)
    # TODO: This should still use the correct language to count in the target
    # language
    return len(Common.words(string))


def wordsinunit(unit):
    """Counts the words in the unit's source and target, taking plurals into
    account. The target words are only counted if the unit is translated.
    """
    (sourcewords, targetwords) = (0, 0)
    if isinstance(unit.source, multistring):
        sourcestrings = unit.source.strings
    else:
        sourcestrings = [unit.source or ""]
    for s in sourcestrings:
        sourcewords += wordcount(s)
    if not unit.istranslated():
        return sourcewords, targetwords
    if isinstance(unit.target, multistring):
        targetstrings = unit.target.strings
    else:
        targetstrings = [unit.target or ""]
    for s in targetstrings:
        targetwords += wordcount(s)
    return sourcewords, targetwords


def calcstats(filename):
    # ignore totally blank or header units
    try:
        store = factory.getobject(filename)
    except ValueError as e:
        logger.warning(e)
        return {}

    units = [unit for unit in store.units if unit.istranslatable()]  # type: ignore
    translated = translatedmessages(units)
    fuzzy = fuzzymessages(units)
    review = [unit for unit in units if unit.isreview()]
    untranslated = untranslatedmessages(units)
    wordcounts = {id(unit): wordsinunit(unit) for unit in units}
    sourcewords = lambda elementlist: sum(
        wordcounts[id(unit)][0] for unit in elementlist
    )
    targetwords = lambda elementlist: sum(
        wordcounts[id(unit)][1] for unit in elementlist
    )
    stats = {"filename": filename}

    # units
    stats["translated"] = len(translated)
    stats["fuzzy"] = len(fuzzy)
    stats["untranslated"] = len(untranslated)
    stats["review"] = len(review)
    stats["total"] = stats["translated"] + stats["fuzzy"] + stats["untranslated"]

    # words
    stats["translatedsourcewords"] = sourcewords(translated)
    stats["translatedtargetwords"] = targetwords(translated)
    stats["fuzzysourcewords"] = sourcewords(fuzzy)
    stats["untranslatedsourcewords"] = sourcewords(untranslated)
    stats["reviewsourcewords"] = sourcewords(review)
    stats["totalsourcewords"] = (
        stats["translatedsourcewords"]
        + stats["fuzzysourcewords"]
        + stats["untranslatedsourcewords"]
    )

    stats["extended"] = file_extended_totals(units, wordcounts)

    return stats


def file_extended_totals(units, wordcounts):
    """
    Provide extended statuses (used by XLIFF)
    """

    stats = {}

    for unit in units:
        state = unit.get_state_n()

        # if state is not standard (xliff)
        # search for the default one to use
        # each unit defines its own states
        if state not in extended_state_strings:
            for k in unit.STATE.keys():
                val = unit.STATE[k]
                if val[0] <= int(state.__str__()) <= val[1]:
                    state = k

        extended_state = extended_state_strings[state]

        state_stats = stats.get(extended_state, defaultdict(int))
        state_stats["units"] += 1
        state_stats["sourcewords"] += wordcounts[id(unit)][0]
        state_stats["targetwords"] += wordcounts[id(unit)][1]

        stats[extended_state] = state_stats

    return stats


@dataclass
class Renderer:
    stats: StatCollector

    def header(self):
        pass

    def entry(self, title: str, stats: dict):
        pass

    def footer(self):
        pass


@dataclass
class CsvRenderer(Renderer):
    def __post_init__(self):
        self._fields = {
            "title": "Filename",
            "translated": "Translated Messages",
            "translatedsourcewords": "Translated Source Words",
            "translatedtargetwords": "Translated Target Words",
            "fuzzy": "Fuzzy Messages",
            "fuzzysourcewords": "Fuzzy Source Words",
            "untranslated": "Untranslated Messages",
            "untranslatedsourcewords": "Untranslated Source Words",
            "total": "Total Message",
            "totalsourcewords": "Total Source Words",
            "review": "Review Messages",
            "reviewsourcewords": "Review Source Words",
        }
        self._writer = csv.DictWriter(sys.stdout, self._fields.values())

    def header(self):
        self._writer.writeheader()

    def entry(self, title, stats):
        data = stats.copy()
        data.update(title=title)
        row = {v: data[k] for k, v in self._fields.items()}
        self._writer.writerow(row)


class FullRenderer(Renderer):
    def entry(self, title, stats):
        print(
            "Processing file : " + ConsoleColor.HEADER() + title + ConsoleColor.ENDC()
        )
        print("Type               Strings      Words (source)    Words (translation)")
        print(
            ConsoleColor.OKGREEN()
            + "Translated:   %5d (%3d%%) %10d (%3d%%) %15d"
            % (
                stats["translated"],
                percent(stats["translated"], stats["total"]),
                stats["translatedsourcewords"],
                percent(stats["translatedsourcewords"], stats["totalsourcewords"]),
                stats["translatedtargetwords"],
            )
            + ConsoleColor.ENDC()
        )
        print(
            ConsoleColor.WARNING()
            + "Fuzzy:        %5d (%3d%%) %10d (%3d%%)             n/a"
            % (
                stats["fuzzy"],
                percent(stats["fuzzy"], stats["total"]),
                stats["fuzzysourcewords"],
                percent(stats["fuzzysourcewords"], stats["totalsourcewords"]),
            )
            + ConsoleColor.ENDC()
        )
        print(
            ConsoleColor.FAIL()
            + "Untranslated: %5d (%3d%%) %10d (%3d%%)             n/a"
            % (
                stats["untranslated"],
                percent(stats["untranslated"], stats["total"]),
                stats["untranslatedsourcewords"],
                percent(stats["untranslatedsourcewords"], stats["totalsourcewords"]),
            )
            + ConsoleColor.ENDC()
        )
        print(
            "Total:        %5d %17d %22d"
            % (
                stats["total"],
                stats["totalsourcewords"],
                stats["translatedtargetwords"],
            )
        )
        if "extended" in stats:
            print("")
            for state, e_stats in stats["extended"].items():
                print(
                    "%-11s   %5d (%3d%%) %10d (%3d%%) %15d"
                    % (
                        state.title() + ":",
                        e_stats["units"],
                        percent(e_stats["units"], stats["total"]),
                        e_stats["sourcewords"],
                        percent(e_stats["sourcewords"], stats["totalsourcewords"]),
                        e_stats["targetwords"],
                    )
                )

        if stats["review"] > 0:
            print(
                "review:          %5d %17d                    n/a"
                % (stats["review"], stats["reviewsourcewords"])
            )
        print()

    def footer(self):
        stats = self.stats
        if stats.file_count > 1:
            if stats.incomplete_only:
                self.entry("TOTAL (incomplete only):", stats.totals)
                print("File count (incomplete):   %5d" % len(stats.results))
            else:
                self.entry("TOTAL:", stats.totals)
            print("File count:   %5d" % (stats.file_count))
            print()


@dataclass
class ShortStringsRenderer(Renderer):
    """
    :param indent: indentation of the 2nd column (length of longest filename)
    """

    indent: int = 8

    def entry(self, title, stats):
        spaces = " " * (self.indent - len(title))
        print(
            "%s%s strings: total: %d\t| %st\t%sf\t%su\t| %st\t%sf\t%su"
            % (
                ConsoleColor.HEADER() + title + ConsoleColor.ENDC(),
                spaces,
                stats["total"],
                ConsoleColor.OKGREEN() + str(stats["translated"]) + ConsoleColor.ENDC(),
                ConsoleColor.WARNING() + str(stats["fuzzy"]) + ConsoleColor.ENDC(),
                ConsoleColor.FAIL() + str(stats["untranslated"]) + ConsoleColor.ENDC(),
                ConsoleColor.OKGREEN()
                + str(percent(stats["translated"], stats["total"]))
                + "%"
                + ConsoleColor.ENDC(),
                ConsoleColor.WARNING()
                + str(percent(stats["fuzzy"], stats["total"]))
                + "%"
                + ConsoleColor.ENDC(),
                ConsoleColor.FAIL()
                + str(percent(stats["untranslated"], stats["total"]))
                + "%"
                + ConsoleColor.ENDC(),
            )
        )


@dataclass
class ShortWordsRenderer(Renderer):
    """
    :param indent: indentation of the 2nd column (length of longest filename)
    """

    indent: int = 8

    def entry(self, title, stats):
        spaces = " " * (self.indent - len(title))
        print(
            "%s%s source words: total: %d\t| %st\t%sf\t%su\t| %st\t%sf\t%su"
            % (
                ConsoleColor.HEADER() + title + ConsoleColor.ENDC(),
                spaces,
                stats["totalsourcewords"],
                ConsoleColor.OKGREEN()
                + str(stats["translatedsourcewords"])
                + ConsoleColor.ENDC(),
                ConsoleColor.WARNING()
                + str(stats["fuzzysourcewords"])
                + ConsoleColor.ENDC(),
                ConsoleColor.FAIL()
                + str(stats["untranslatedsourcewords"])
                + ConsoleColor.ENDC(),
                ConsoleColor.OKGREEN()
                + str(
                    percent(stats["translatedsourcewords"], stats["totalsourcewords"])
                )
                + "%"
                + ConsoleColor.ENDC(),
                ConsoleColor.WARNING()
                + str(percent(stats["fuzzysourcewords"], stats["totalsourcewords"]))
                + "%"
                + ConsoleColor.ENDC(),
                ConsoleColor.FAIL()
                + str(
                    percent(stats["untranslatedsourcewords"], stats["totalsourcewords"])
                )
                + "%"
                + ConsoleColor.ENDC(),
            )
        )


def percent(denominator, devisor):
    if devisor == 0:
        return 0
    else:
        return denominator * 100 / devisor


def fuzzymessages(units):
    return [unit for unit in units if unit.isfuzzy() and unit.target]


def translatedmessages(units):
    return [unit for unit in units if unit.istranslated()]


def untranslatedmessages(units):
    return [
        unit
        for unit in units
        if not (unit.istranslated() or unit.isfuzzy()) and unit.source
    ]


class StatCollector:
    def __init__(self, items: list[str], incomplete_only=False):
        self.incomplete_only = incomplete_only
        self._results: list[dict] = []
        self._handle_items(items)

    def render(self, renderer_class: type[Renderer]):
        if renderer_class in (ShortWordsRenderer, ShortStringsRenderer):
            renderer = renderer_class(self, indent=self.longest_filename)
        else:
            renderer = renderer_class(self)

        renderer.header()
        for entry in self.results:
            renderer.entry(entry["filename"], entry)
        renderer.footer()

    @cached_property
    def results(self):
        if self.incomplete_only:
            return [s for s in self._results if s["total"] != s["translated"]]
        else:
            return self._results

    def _handle_items(self, items: list[str]):
        for item in items:
            if not os.path.exists(item):
                logger.error("cannot process %s: does not exist", item)
                continue
            elif os.path.isdir(item):
                self._handle_dir(item)
            else:
                self._handle_single_file(item)

    def _handle_dir(self, dirname):
        _, name = os.path.split(dirname)
        if name in ["CVS", ".svn", "_darcs", ".git", ".hg", ".bzr"]:
            return
        entries = os.listdir(dirname)
        self._handle_multiple_files(dirname, entries)

    def _handle_multiple_files(self, dirname, filenames):
        for filename in filenames:
            pathname = os.path.join(dirname, filename)
            if os.path.isdir(pathname):
                self._handle_dir(pathname)
            else:
                self._handle_single_file(pathname)

    def _handle_single_file(self, filename):
        try:
            stats = calcstats(filename)
            self._results.append(stats)
        except Exception:  # This happens if we have a broken file.
            logger.error(sys.exc_info()[1])

    @property
    def longest_filename(self):
        filenames = list(map(itemgetter("filename"), self.results)) or [""]
        return max(len(f) for f in filenames)

    @property
    def file_count(self):
        return len(self._results)

    @cached_property
    def totals(self) -> dict:
        """Total stats"""
        totals = defaultdict(int)
        for stats in self._results:
            for key, value in stats.items():
                if key in ["extended", "filename"]:
                    # FIXME: calculate extended totals
                    continue
                totals[key] += value
        return totals


def main(arguments=None):
    parser = ArgumentParser()
    parser.add_argument(
        "--incomplete",
        action="store_true",
        default=False,
        dest="incomplete_only",
        help="skip 100%% translated files.",
    )
    output_group = parser.add_argument_group("Output format")
    megroup = output_group.add_mutually_exclusive_group()
    megroup.add_argument(
        "--full",
        action="store_const",
        const=FullRenderer,
        dest="style",
        default=FullRenderer,
        help="(default) statistics in full, verbose format",
    )
    megroup.add_argument(
        "--csv",
        action="store_const",
        const=CsvRenderer,
        dest="style",
        help="statistics in CSV format",
    )
    megroup.add_argument(
        "--short",
        action="store_const",
        const=ShortStringsRenderer,
        dest="style",
        help="same as --short-strings",
    )
    megroup.add_argument(
        "--short-strings",
        action="store_const",
        const=ShortStringsRenderer,
        dest="style",
        help="statistics of strings in short format - one line per file",
    )
    megroup.add_argument(
        "--short-words",
        action="store_const",
        const=ShortWordsRenderer,
        dest="style",
        help="statistics of words in short format - one line per file",
    )
    output_group.add_argument(
        "--no-color", action="store_true", help="show output without color"
    )

    parser.add_argument("files", nargs="+")

    args = parser.parse_args(arguments)

    logging.basicConfig(format="%(name)s: %(levelname)s: %(message)s")
    ConsoleColor.color_mode = not args.no_color

    StatCollector(args.files, args.incomplete_only).render(args.style)


if __name__ == "__main__":
    main()
