# Copyright © Michal Čihař <michal@weblate.org>
#
# SPDX-License-Identifier: GPL-3.0-or-later

from __future__ import annotations

from celery import current_task
from django.conf import settings
from django.core.exceptions import PermissionDenied
from django.db import transaction
from django.db.models.functions import MD5

from weblate.machinery.models import MACHINERY
from weblate.trans.models import Change, Component, Suggestion, Unit
from weblate.trans.util import split_plural
from weblate.utils.state import STATE_APPROVED, STATE_FUZZY, STATE_TRANSLATED


class AutoTranslate:
    def __init__(
        self,
        user,
        translation,
        filter_type: str,
        mode: str,
        component_wide: bool = False,
    ):
        self.user = user
        self.translation = translation
        translation.component.batch_checks = True
        self.filter_type = filter_type
        self.mode = mode
        self.updated = 0
        self.progress_steps = 0
        self.target_state = STATE_TRANSLATED
        if mode == "fuzzy":
            self.target_state = STATE_FUZZY
        elif mode == "approved":
            self.target_state = STATE_APPROVED
        self.component_wide = component_wide

    def get_units(self, filter_mode=True):
        units = self.translation.unit_set.all()
        if self.mode == "suggest" and filter_mode:
            units = units.filter(suggestion__isnull=True)
        return units.filter_type(self.filter_type)

    def set_progress(self, current):
        if current_task and current_task.request.id and self.progress_steps:
            current_task.update_state(
                state="PROGRESS",
                meta={
                    "progress": 100 * current // self.progress_steps,
                    "translation": self.translation.pk,
                },
            )

    def update(self, unit, state, target):
        if isinstance(target, str):
            target = [target]
        if self.mode == "suggest" or any(
            len(item) > unit.get_max_length() for item in target
        ):
            Suggestion.objects.add(unit, target, None, False)
        else:
            unit.is_batch_update = True
            unit.translate(
                self.user, target, state, Change.ACTION_AUTO, propagate=False
            )
        self.updated += 1

    def post_process(self):
        if self.updated > 0:
            if not self.component_wide:
                self.translation.component.update_source_checks()
                self.translation.component.run_batched_checks()
            self.translation.invalidate_cache()
            if self.user:
                self.user.profile.increase_count("translated", self.updated)

    @transaction.atomic
    def process_others(self, source: int | None):
        """Perform automatic translation based on other components."""
        kwargs = {
            "translation__plural": self.translation.plural,
            "state__gte": STATE_TRANSLATED,
        }
        source_language = self.translation.component.source_language
        exclude = {}
        if source:
            component = Component.objects.get(id=source)

            if (
                not component.project.contribute_shared_tm
                and component.project != self.translation.component.project
            ):
                raise PermissionDenied(
                    "Project has disabled contribution to shared translation memory."
                )
            if component.source_language != source_language:
                raise PermissionDenied("Component have different source languages.")
            kwargs["translation__component"] = component
        else:
            project = self.translation.component.project
            kwargs["translation__component__project"] = project
            kwargs["translation__component__source_language"] = source_language
            exclude["translation"] = self.translation
        sources = Unit.objects.filter(**kwargs)

        # Use memory_db for the query in case it exists. This is supposed
        # to be a read-only replica for offloading expensive translation
        # queries.
        if "memory_db" in settings.DATABASES:
            sources = sources.using("memory_db")

        if exclude:
            sources = sources.exclude(**exclude)

        # Fetch translations
        translations = {
            source: split_plural(target)
            for source, state, target in sources.filter(
                source__md5__in=self.get_units()
                .annotate(source__md5=MD5("source"))
                .values("source__md5")
            ).values_list("source", "state", "target")
        }

        # We need to skip mode (suggestions) filtering here as SELECT FOR UPDATE
        # cannot be used with JOIN
        units = (
            self.get_units(False)
            .filter(source__in=translations.keys())
            .prefetch_bulk()
            .select_for_update()
        )
        self.progress_steps = len(units)

        for pos, unit in enumerate(units):
            # Get update
            try:
                target = translations[unit.source]
            except KeyError:
                # Happens on MySQL due to case-insensitive lookup
                continue

            self.set_progress(pos)

            # No save if translation is same or unit does not exist
            if unit.state == self.target_state and unit.target == target:
                continue
            # Copy translation
            self.update(unit, self.target_state, target)

        self.post_process()

    def fetch_mt(self, engines, threshold):
        """Get the translations."""
        units = self.get_units()
        num_units = len(units)

        machinery_settings = self.translation.component.project.get_machinery_settings()

        engines = sorted(
            (
                MACHINERY[engine](setting)
                for engine, setting in machinery_settings.items()
                if engine in MACHINERY and engine in engines
            ),
            key=lambda engine: engine.get_rank(),
            reverse=True,
        )

        self.progress_steps = 2 * (len(engines) + num_units)

        for pos, translation_service in enumerate(engines):
            batch_size = translation_service.batch_size

            for batch_start in range(0, num_units, batch_size):
                translation_service.batch_translate(
                    units[batch_start : batch_start + batch_size],
                    self.user,
                    threshold=threshold,
                )
                self.set_progress(pos * num_units + batch_start)

        return {
            unit.id: unit.machinery["translation"]
            for unit in units
            if unit.machinery and any(unit.machinery["quality"])
        }

    def process_mt(self, engines: list[str], threshold: int):
        """Perform automatic translation based on machine translation."""
        translations = self.fetch_mt(engines, int(threshold))

        # Adjust total number to show correct progress
        offset = self.progress_steps / 2
        self.progress_steps = offset + len(translations)

        with transaction.atomic():
            # Perform the translation
            for pos, unit in enumerate(
                self.translation.unit_set.filter(id__in=translations.keys())
                .prefetch_bulk()
                .select_for_update()
            ):
                # Copy translation
                self.update(unit, self.target_state, translations[unit.pk])
                self.set_progress(offset + pos)

            self.post_process()
