import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from .protector import protect_words, restore_words
from .rules import apply_custom_rules
from .config import MODEL_NAME

class SmartTranslator:
    def __init__(self, src_lang, tgt_lang):
        self.src_lang = src_lang
        self.tgt_lang = tgt_lang

        # Detect device
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print("Using device:", self.device)

        # Load tokenizer and model
        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
        self.model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)

        # Move model to the device
        self.model.to(self.device)
        self.model.eval()

    def translate(self, text, keep_words=None):
        keep_words = keep_words or []

        protected_text, mapping = protect_words(text, keep_words)

        self.tokenizer.src_lang = self.src_lang
        inputs = self.tokenizer(protected_text, return_tensors="pt")

        # Move input tensors to same device as model
        inputs = {k: v.to(self.device) for k, v in inputs.items()}

        tgt_lang_id = self.tokenizer.convert_tokens_to_ids(self.tgt_lang)

        outputs = self.model.generate(
            **inputs,
            forced_bos_token_id=tgt_lang_id,
            max_length=256
        )

        translated = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        translated = restore_words(translated, mapping)
        translated = apply_custom_rules(translated)

        return translated
