import gradio as gr
from transformers import pipeline
from utils import preprocess_text, postprocess_translation, TextProcessor
from css import custom_css
from cultural_utils import get_cultural_context
import logging
import sys

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    stream=sys.stdout
)
logger = logging.getLogger(__name__)

class TranslationService:
    """Handles the translation pipeline and text processing."""

    def __init__(self):
        try:
            logger.info("Loading translation pipeline...")
            self.pipeline = pipeline(
                "translation",
                model="persiannlp/mt5-small-parsinlu-translation_en_fa",
                device="cpu"
            )
            logger.info("Translation pipeline loaded successfully")
        except Exception as e:
            logger.error(f"Failed to load translation pipeline: {str(e)}")
            raise

    def translate(self, text: str, source_lang: str, target_lang: str) -> dict:
        """Translate text between English and Farsi with cultural context."""
        try:
            if not text:
                return {"translation": "", "cultural_context": {"idioms": []}}

            # Validate input
            is_valid, error_msg = TextProcessor.validate_input(text)
            if not is_valid:
                return {"translation": error_msg, "cultural_context": {"idioms": []}}

            # Get cultural context before translation
            src_lang = "en" if source_lang == "English" else "fa"
            cultural_context = get_cultural_context(text, src_lang)

            # Preprocess text
            processed_text = preprocess_text(text)
            logger.info(f"Translating text from {source_lang} to {target_lang}")

            # Generate translation
            result = self.pipeline(processed_text)

            # Post-process and return result
            translated_text = postprocess_translation(result[0]['translation_text'])
            return {
                "translation": translated_text,
                "cultural_context": cultural_context
            }

        except Exception as e:
            logger.error(f"Translation error: {str(e)}")
            return {
                "translation": "Translation error occurred. Please try again.",
                "cultural_context": {"idioms": []}
            }

class GradioInterface:
    """Manages the Gradio web interface components."""

    def __init__(self, translation_service):
        self.translation_service = translation_service

    def create(self):
        """Create and configure the Gradio interface."""
        with gr.Blocks(css=custom_css) as interface:
            self._create_header()
            source_lang, target_lang = self._create_language_controls()
            self._create_translation_interface(source_lang, target_lang)
            return interface

    def _create_header(self):
        """Create the interface header."""
        gr.Markdown("""
            # English-Farsi Translation
            Culturally-sensitive translation between English and Farsi
        """)

    def _create_language_controls(self):
        """Create language selection controls."""
        with gr.Row():
            source_lang = gr.Dropdown(
                choices=["English", "Farsi"],
                value="English",
                label="Source Language"
            )
            target_lang = gr.Dropdown(
                choices=["Farsi", "English"],
                value="Farsi",
                label="Target Language"
            )

            swap_btn = gr.Button("🔄 Swap Languages")
            swap_btn.click(
                fn=lambda s, t: (t, s),
                inputs=[source_lang, target_lang],
                outputs=[source_lang, target_lang]
            )
        return source_lang, target_lang

    def _create_translation_interface(self, source_lang, target_lang):
        """Create the translation input/output interface."""
        with gr.Row():
            with gr.Column():
                input_text = gr.Textbox(
                    lines=5,
                    placeholder="Enter text to translate...",
                    label="Input Text"
                )
            with gr.Column():
                output_text = gr.Textbox(
                    lines=5,
                    label="Translation",
                    rtl=True
                )
                cultural_notes = gr.Markdown(
                    label="Cultural Context",
                    value="Cultural context annotations will appear here..."
                )

        def format_cultural_context(result):
            translation = result["translation"]
            context = result["cultural_context"]

            if not context["idioms"]:
                return translation, "No cultural annotations found."

            notes = "### Cultural Context Notes:\n"
            for idiom, literal, explanation in context["idioms"]:
                notes += f"- **{idiom}**\n"
                notes += f"  - Literal: {literal}\n"
                notes += f"  - Context: {explanation}\n\n"

            return translation, notes

        translate_btn = gr.Button("Translate", variant="primary")
        def handle_translation(text, src_lang, tgt_lang):
            result = self.translation_service.translate(text, src_lang, tgt_lang)
            translation = result["translation"]
            context = result["cultural_context"]

            if not context["idioms"]:
                return translation, "No cultural annotations found."

            notes = "### Cultural Context Notes:\n"
            for idiom, literal, explanation in context["idioms"]:
                notes += f"- **{idiom}**\n"
                notes += f"  - Literal: {literal}\n"
                notes += f"  - Context: {explanation}\n\n"

            return translation, notes

        translate_btn.click(
            fn=handle_translation,
            inputs=[input_text, source_lang, target_lang],
            outputs=[output_text, cultural_notes]
        )

        source_lang.change(
            fn=lambda lang: gr.update(rtl=lang == "Farsi"),
            inputs=[source_lang],
            outputs=[input_text]
        )

def main():
    """Initialize and launch the translation application."""
    try:
        logger.info("Initializing translation service...")
        translation_service = TranslationService()

        logger.info("Creating Gradio interface...")
        interface = GradioInterface(translation_service).create()

        interface.launch(
            server_name="0.0.0.0",
            server_port=8000,
            share=True,
            debug=True,
            show_error=True
        )
        logger.info("Gradio interface started successfully")
    except Exception as e:
        logger.error(f"Application startup failed: {str(e)}")
        raise

if __name__ == "__main__":
    main()