from fastapi import FastAPI
from pydantic import BaseModel
from transformers import MarianMTModel, MarianTokenizer, pipeline

app = FastAPI()

# SK → EN
model_name_sk_en = "Helsinki-NLP/opus-mt-sk-en"
tokenizer_sk_en = MarianTokenizer.from_pretrained(model_name_sk_en, cache_dir="/app/cache")
model_sk_en = MarianMTModel.from_pretrained(model_name_sk_en, cache_dir="/app/cache")
sk_to_en = pipeline("translation", model=model_sk_en, tokenizer=tokenizer_sk_en)

# EN → SK
model_name_en_sk = "Helsinki-NLP/opus-mt-en-sk"
tokenizer_en_sk = MarianTokenizer.from_pretrained(model_name_en_sk, cache_dir="/app/cache")
model_en_sk = MarianMTModel.from_pretrained(model_name_en_sk, cache_dir="/app/cache")
en_to_sk = pipeline("translation", model=model_en_sk, tokenizer=tokenizer_en_sk)

class TranslationRequest(BaseModel):
    data: list

@app.post("/predict/")
async def predict(req: TranslationRequest):
    try:
        text = req.data[0]
        direction = req.data[1]
        if direction == "sk-en":
            translation = sk_to_en(text)[0]["translation_text"]
        elif direction == "en-sk":
            translation = en_to_sk(text)[0]["translation_text"]
        else:
            return {"error": "Nepodporovaný smer prekladu."}
        return {"data": [translation]}
    except Exception as e:
        return {"error": str(e)}