import re
import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

app = FastAPI()

model_id = "sagawa/ReactionT5v2-retrosynthesis-USPTO_50k"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSeq2SeqLM.from_pretrained(model_id)

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

class CompletionRequest(BaseModel):
    model: str
    prompt: str
    max_tokens: int = 128
    temperature: float = 0.0

# Офіційний токенізатор авторів ReactionT5
def smi_tokenizer(smi):
    pattern = re.compile(r"(\[[^\]]+\]|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\(|\)|\.|=|#|-|\+|\\|\/|:|~|@|\?|>|\*|\$|\%[0-9]{2}|[0-9])")
    tokens = [token for token in pattern.split(smi) if token]
    return " ".join(tokens)

@app.post("/chat/completions")
async def text_completions(request: CompletionRequest):
    try:
        raw_prompt = request.prompt.strip()
        
        pure_smiles = raw_prompt.split()[-1] 
        
        templated_smiles = smi_tokenizer(pure_smiles)
        
        formatted_prompt = f"Predict reactants: {templated_smiles}"

        inputs = tokenizer(formatted_prompt, return_tensors="pt").to(device)

        with torch.inference_mode():
            outputs = model.generate(
                **inputs,
                max_new_tokens=request.max_tokens,
                num_beams=5,
                num_return_sequences=5,
                do_sample=False,
                early_stopping=True,
            )

        results = []
        for i in range(5):
            decoded = tokenizer.decode(outputs[i], skip_special_tokens=True)
            cleaned = decoded.replace(" ", "").rstrip(".")
            results.append(cleaned)

        return {
            "id": "cmpl-raw-123",
            "object": "text_completion",
            "model": request.model,
            "choices": [{
                "text": results[0],
                "all_predictions": results,
                "index": 0,
                "logprobs": None,
                "finish_reason": "stop",
            }],
        }
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc))

@app.get("/")
def health():
    return {"status": "healthy"}