import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import uvicorn

# Define request model
class TranslateRequest(BaseModel):
    text: str
    src: str
    tgt: str

# Check for available device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Initialize FastAPI app
app = FastAPI()

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Define model and tokenizer dictionary
model_dict = {
    "zh2en": "Helsinki-NLP/opus-mt-zh-en",
    "en2zh": "Helsinki-NLP/opus-mt-en-zh",
}

# Load models and tokenizers
models = {}
for key, value in model_dict.items():
    tokenizer = AutoTokenizer.from_pretrained(value)
    model = AutoModelForSeq2SeqLM.from_pretrained(value)
    model.to(device)
    models[key] = (model, tokenizer)

# Translation function
def translate(model, tokenizer, text):
    inputs = tokenizer(text, return_tensors="pt").input_ids.to(device)
    outputs = model.generate(inputs)
    translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return translated_text

# FastAPI route
@app.post("/translate")
async def translate_route(translate_request: TranslateRequest):
    text = translate_request.text
    src = translate_request.src
    tgt = translate_request.tgt

    if src == "zh" or tgt == "en":
        model_key = "zh2en"
    elif src == "en" or tgt == "zh":
        model_key = "en2zh"
    else:
        raise HTTPException(status_code=400, detail="Unsupported language pair")

    model, tokenizer = models[model_key]
    translated_text = translate(model, tokenizer, text)
    return {"text": translated_text}


# Run server
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=4200)