from fastapi import FastAPI, HTTPException
import huggingface_hub
import torch
from fairseq.models.bart import BARTModel
import os
import re

EXPECTED_TOKEN = os.environ.get("EXPECTED_TOKEN")
REPO_NAME = os.environ.get("REPO_NAME")
REPO_NAME_HUGG = os.environ.get("REPO_NAME_HUGG")

app = FastAPI()

path_model = huggingface_hub.hf_hub_download(repo_id=f'{REPO_NAME}/{REPO_NAME_HUGG}' , filename='model/checkpoint_best.pt', token=True)

files = ['dict.src.txt', 'dict.tgt.txt', 'preprocess.log',
         'train.src-tgt.src.bin', 'train.src-tgt.src.idx', 'train.src-tgt.tgt.bin', 'train.src-tgt.tgt.idx',
         'valid.src-tgt.src.bin', 'valid.src-tgt.src.idx', 'valid.src-tgt.tgt.bin', 'valid.src-tgt.tgt.idx']

for file in files:
    path_data = huggingface_hub.hf_hub_download(repo_id=f'{REPO_NAME}/{REPO_NAME_HUGG}' , filename=file, subfolder='gec_data-bin_ptbr', token=True)

bart = BARTModel.from_pretrained(
    '/'.join(path_model.split('/')[:-1]),
    checkpoint_file='checkpoint_best.pt',
    data_name_or_path='/'.join(path_data.split('/')[:-1])
)

bart.eval()


def posprocessing(frase: str):
    
    nova_frase = re.sub(r'\s+([.,!?;:])', r'\1', frase)
    
    return nova_frase


@app.post('/inference')
async def predict(frase: str, token: str):

    if token != EXPECTED_TOKEN:
        raise HTTPException(status_code=401, detail="Token inválido")

    with torch.no_grad():

        result = bart.sample([frase], beam=1)
        
    return posprocessing(result[0])