from model.request_models import InputElement
from model.response_models import EntityOIEResponseModel, ResponseOIEModel
import nltk

from typing import List, Dict

from pathlib import Path

from src.Multi2OIE.OIE import OIE
from src.helpers import DownloadModel


class ResponsesOIE:
    responses: Dict[int, ResponseOIEModel] = {}

    @staticmethod
    def check_id(request_id: int) -> bool:
        if ResponsesOIE.responses.get(request_id):
            return True
        return False

    @staticmethod
    def get_response(request_id: int) -> ResponseOIEModel:
        if ResponsesOIE.responses.get(request_id):
            return ResponsesOIE.responses[request_id]


class OIEResponse:
    def __init__(self, name="./data/mbertEntail-2GBEn.bin", BERT='persiannlp/mbert-base-parsinlu-entailment',
                 batch_size: int = 1, device: str = "gpu"):

        nltk.download('stopwords')
        # load the model you trained
        if name == "./data/mbertEntail-2GBEn.bin":
            model_path = Path(name)
            if not model_path.is_file():
                DownloadModel().download_oie_model(name)
        self.model = OIE(name, BERT, batch_size)

    def predict_json(self, input_json: List[InputElement], request_id: int) -> ResponseOIEModel:
        ResponsesOIE.responses[request_id] = ResponseOIEModel(progression="in going")
        sentences = []
        for element in input_json:
            text = element.text
            lang = element.lang
            sentences.append(text)

        result = self.model.predict(sentences)
        ResponsesOIE.responses[request_id] = ResponseOIEModel(progression="done", result=result)
        print(request_id, " done ")
        return ResponsesOIE.responses[request_id]
