import gradio as gr
from transformers import pipeline
import json
import re

# Load the phrase mapping from the JSON file
with open('phrase_mappings.json', 'r') as f:
    phrase_mapping = json.load(f)

# Define the UnifiedTranslator class
class UnifiedTranslator:
    def __init__(self, model_name, phrase_mapping):
        self.model = pipeline("translation", model=model_name)
        self.phrase_mapping = phrase_mapping

    def translate(self, text):
        # Normalize text to lowercase and strip extra spaces
        text_lower = text.lower().strip()
        
        # Debugging output
        print(f"Input text: {text_lower}")

        # Check if the text matches any pattern in the phrase_mapping
        for pattern, translation in self.phrase_mapping.items():
            try:
                # Use regex to match the pattern with placeholders
                # Make sure pattern is properly escaped and case insensitive
                pattern_regex = re.compile(
                    re.escape(pattern).replace(r"\{name\}", r"([\w'-]+)").strip(), 
                    re.IGNORECASE
                )
                match = pattern_regex.fullmatch(text_lower)
                if match:
                    # Replace the placeholder with the actual value if needed
                    if '{name}' in pattern:
                        return translation.format(name=match.group(1))
                    else:
                        return translation
            except re.error as e:
                print(f"Regex error with pattern {pattern}: {e}")
        
        # Fallback to model translation if no pattern matches
        try:
            translation = self.model(text)[0]
            return translation['translation_text']
        except Exception as e:
            print(f"Model translation error: {e}")
            return "Translation error occurred"

# Initialize the UnifiedTranslator with your model and custom phrases
translator = UnifiedTranslator("Bildad/Swahili-English_Translation", phrase_mapping)

# Define the Gradio interface
def translate_text(text):
    return translator.translate(text)

iface = gr.Interface(
    fn=translate_text,
    inputs="text",
    outputs="text",
    title="Swahili-English Translation",
    description="Translate Swahili to English with custom phrase mappings."
)

# Launch the interface
iface.launch(share=True)