from fastapi import FastAPI
import src.paraphrase.Paraphrase as Paraphrase
import src.translate.Translate as Translate
from typing import Optional

app = FastAPI(docs_url="/")
MTMODELS = {'enro': 'BlackKakapo/opus-mt-en-ro',
            'roen': 'BlackKakapo/opus-mt-ro-en',
            'gemma': 'Gargaz/gemma-2b-romanian-better'}

@app.get("/")
def index():
    return {'endpoints': ['/paraphrase', '/translate'], 'mtmodels': MTMODELS}

@app.get("/paraphrase")
def paraphrase(text: str, model: str):
    resultValue, exception = Paraphrase.paraphraseParaphraseMethod(text, model)
    return {"input": text, "translation": resultValue, "exception": exception}

@app.get("/translate")
def translate(text: str, model: Optional[str] = MTMODELS['gemma']):
    # resultValue, exception = Translate.paraphraseTranslateMethod(text, model)
    resultValue = Translate.gemma_direct(text, model)
    return {"input": text, "result": resultValue, "model": model}