"""
pybergamot - (Somewhat) stable interface for the **Bergamot Translation Engine Python Bindings**.
Copyright (C) 2023  Ad5001

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 3 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 <https://www.gnu.org/licenses/>.

Meta repositories for translation engines
"""
from abc import ABC, abstractmethod
from bergamot import Service, VectorString, ResponseOptions, VectorResponse, TranslationModel


class Engine(ABC):
    """
    An interface for multiple types of translation engine.
    """

    @property
    @abstractmethod
    def source_lang(self) -> str:
        """
        Two-char ISO languages for which the engine translates from.
        """
        pass

    @property
    @abstractmethod
    def target_lang(self) -> str:
        """
        Two-char ISO languages for which the engine translates to.
        """
        pass

    @abstractmethod
    def translate(self, text: str, html: bool = False, alignment: bool = False, quality_scores: bool = False) -> str:
        """
        Translates the text from the engine's source lang its target.

        ## Parameters
        ---
        - **text**: Text to translate.    
        - **html**: Set to True if the text contains an HTML structure which needs to
                     be preserved while translated.    
        - **alignment**: Toggle for alignment.    
        - **quality_scores**: Toggle for whether to include the translation's quality scores
                               for each word in HTML format.    

        ## Returns
        ---
        The translated text.
        """
        pass


class DirectBergamotModelEngine(Engine):
    """
    Internal module. Engine using a single bergamot model to translate.
    """
    def __init__(self, source_lang: str, target_lang: str, model: TranslationModel,
                 service: Service):
        self._source_lang = source_lang
        self._target_lang = target_lang
        self.model = model
        self.service = service

    @property
    def source_lang(self) -> str:
        return self._source_lang

    @property
    def target_lang(self) -> str:
        return self._target_lang

    def translate(self, text: str, html: bool = False, alignment: bool = False, quality_scores: bool = False) -> str:
        opts = ResponseOptions(
            alignment=alignment, qualityScores=quality_scores, HTML=html
        )
        resp: VectorResponse = self.service.translate(self.model, VectorString([text]), opts)
        return resp[0].target.text


class ChainBergamotModelsEngine(Engine):
    """
    Internal module. Engine chaining two bergamot model to translate.
    """
    def __init__(self, source_lang: str, target_lang: str, model1: TranslationModel,
                 model2: TranslationModel, service: Service):
        self._source_lang = source_lang
        self._target_lang = target_lang
        self.model1 = model1
        self.model2 = model2
        self.service = service

    @property
    def source_lang(self) -> str:
        return self._source_lang

    @property
    def target_lang(self) -> str:
        return self._target_lang

    def translate(self, text: str, html: bool = False, alignment: bool = False, quality_scores: bool = False) -> str:
        opts = ResponseOptions(
            alignment=alignment, qualityScores=quality_scores, HTML=html
        )
        resp: VectorResponse = self.service.pivot(self.model1, self.model2, VectorString([text]), opts)
        return resp[0].target.text