"""
Translation Service for translate.press.zone using Modal.com

This service deploys ML translation models on Modal's infrastructure:
- NLLB-200 Distilled 600M model (T4 GPU) for fast, cost-effective translations
- NLLB-200 3.3B model (A100 GPU) for higher quality translations

NLLB (No Language Left Behind) supports 200+ languages with high quality.

Features:
- HTML tag preservation during translation
- Language code mapping (ISO 639-1 to NLLB codes)
- Token usage tracking
- Production-ready error handling
- Comprehensive logging
"""

import modal
import re
from typing import Dict, List, Tuple, Optional
import logging

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Create Modal app
app = modal.App("translate-nllb")

# Docker image with ML dependencies
image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install(
        "transformers==4.44.2",
        "torch==2.4.1",
        "sentencepiece==0.2.0",
        "protobuf==5.28.2",
        "accelerate==0.34.2",
        "sacremoses==0.1.1",
    )
)

# Language code mapping: ISO 639-1 -> NLLB-200 codes
# NLLB uses script codes (e.g., "eng_Latn" for English with Latin script)
LANGUAGE_CODE_MAP = {
    # Common European languages
    "en": "eng_Latn",  # English
    "es": "spa_Latn",  # Spanish
    "fr": "fra_Latn",  # French
    "de": "deu_Latn",  # German
    "it": "ita_Latn",  # Italian
    "pt": "por_Latn",  # Portuguese
    "nl": "nld_Latn",  # Dutch
    "pl": "pol_Latn",  # Polish
    "ru": "rus_Cyrl",  # Russian
    "uk": "ukr_Cyrl",  # Ukrainian
    "cs": "ces_Latn",  # Czech
    "ro": "ron_Latn",  # Romanian
    "sv": "swe_Latn",  # Swedish
    "da": "dan_Latn",  # Danish
    "fi": "fin_Latn",  # Finnish
    "no": "nob_Latn",  # Norwegian
    "el": "ell_Grek",  # Greek
    "bg": "bul_Cyrl",  # Bulgarian
    "hr": "hrv_Latn",  # Croatian
    "sk": "slk_Latn",  # Slovak
    "sl": "slv_Latn",  # Slovenian
    "et": "est_Latn",  # Estonian
    "lv": "lvs_Latn",  # Latvian
    "lt": "lit_Latn",  # Lithuanian

    # Asian languages
    "zh": "zho_Hans",  # Chinese (Simplified)
    "zh-TW": "zho_Hant",  # Chinese (Traditional)
    "ja": "jpn_Jpan",  # Japanese
    "ko": "kor_Hang",  # Korean
    "vi": "vie_Latn",  # Vietnamese
    "th": "tha_Thai",  # Thai
    "id": "ind_Latn",  # Indonesian
    "ms": "zsm_Latn",  # Malay
    "tl": "tgl_Latn",  # Tagalog
    "hi": "hin_Deva",  # Hindi
    "bn": "ben_Beng",  # Bengali
    "ur": "urd_Arab",  # Urdu
    "ta": "tam_Taml",  # Tamil
    "te": "tel_Telu",  # Telugu
    "mr": "mar_Deva",  # Marathi

    # Middle Eastern languages
    "ar": "arb_Arab",  # Arabic
    "he": "heb_Hebr",  # Hebrew
    "fa": "pes_Arab",  # Persian
    "tr": "tur_Latn",  # Turkish

    # African languages
    "sw": "swh_Latn",  # Swahili
    "am": "amh_Ethi",  # Amharic
    "ha": "hau_Latn",  # Hausa
    "yo": "yor_Latn",  # Yoruba
    "ig": "ibo_Latn",  # Igbo
    "zu": "zul_Latn",  # Zulu
    "af": "afr_Latn",  # Afrikaans
}


def get_nllb_code(iso_code: str) -> Optional[str]:
    """
    Convert ISO 639-1 language code to NLLB-200 code

    Args:
        iso_code: ISO 639-1 language code (e.g., "en", "es")

    Returns:
        NLLB-200 language code or None if not supported
    """
    return LANGUAGE_CODE_MAP.get(iso_code.lower())


class HTMLTagPreserver:
    """Utility class to preserve HTML tags during translation"""

    @staticmethod
    def extract_tags(content: str) -> Tuple[str, Dict[str, str]]:
        """
        Extract HTML tags and replace with placeholders

        Args:
            content: Original content with HTML tags

        Returns:
            Tuple of (content with placeholders, tag mapping)
        """
        # Pattern to match HTML tags
        tag_pattern = r'<[^>]+>'

        tags = {}
        tag_counter = 0

        def replace_tag(match):
            nonlocal tag_counter
            tag = match.group(0)
            placeholder = f"__TAG_{tag_counter}__"
            tags[placeholder] = tag
            tag_counter += 1
            return placeholder

        # Replace all HTML tags with placeholders
        clean_content = re.sub(tag_pattern, replace_tag, content)

        return clean_content, tags

    @staticmethod
    def restore_tags(translated_content: str, tags: Dict[str, str]) -> str:
        """
        Restore HTML tags in translated content

        Args:
            translated_content: Translated content with placeholders
            tags: Mapping of placeholders to original tags

        Returns:
            Content with restored HTML tags
        """
        result = translated_content

        # Replace placeholders with original tags
        for placeholder, tag in tags.items():
            result = result.replace(placeholder, tag)

        return result


@app.function(
    image=image,
    gpu="T4",
    timeout=300,
    memory=8192,
)
def translate_4b(
    content: str,
    source_lang: str,
    target_lang: str,
    tone: str = "neutral"
) -> Dict[str, any]:
    """
    Translate using NLLB-200 Distilled 600M model (lightweight, fast)

    This model supports 200+ languages with good quality and fast inference.

    Args:
        content: Text to translate
        source_lang: Source language code (ISO 639-1, e.g., "en", "es")
        target_lang: Target language code (ISO 639-1)
        tone: Translation tone (neutral, formal, casual) - Note: NLLB doesn't natively support tone

    Returns:
        Dict with translation, tokens_used, processing_time_ms, and model

    Raises:
        ValueError: If language codes are not supported
        RuntimeError: If translation fails
    """
    from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
    import time

    start_time = time.time()

    try:
        # Validate and convert language codes
        src_nllb_code = get_nllb_code(source_lang)
        tgt_nllb_code = get_nllb_code(target_lang)

        if not src_nllb_code:
            raise ValueError(
                f"Unsupported source language: {source_lang}. "
                f"Supported languages: {', '.join(sorted(LANGUAGE_CODE_MAP.keys()))}"
            )

        if not tgt_nllb_code:
            raise ValueError(
                f"Unsupported target language: {target_lang}. "
                f"Supported languages: {', '.join(sorted(LANGUAGE_CODE_MAP.keys()))}"
            )

        logger.info(f"Translating from {source_lang} ({src_nllb_code}) to {target_lang} ({tgt_nllb_code})")

        # Extract HTML tags
        clean_content, tags = HTMLTagPreserver.extract_tags(content)

        if not clean_content.strip():
            logger.warning("Empty content after tag extraction")
            return {
                "translation": content,  # Return original if empty
                "tokens_used": 0,
                "processing_time_ms": 0,
                "model": "4b"
            }

        # Load NLLB-200 distilled model (cached after first run)
        model_name = "facebook/nllb-200-distilled-600M"
        logger.info(f"Loading model: {model_name}")

        tokenizer = AutoTokenizer.from_pretrained(
            model_name,
            src_lang=src_nllb_code,
            tgt_lang=tgt_nllb_code
        )
        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

        # Tokenize input
        inputs = tokenizer(
            clean_content,
            return_tensors="pt",
            max_length=512,
            truncation=True,
            padding=True
        )

        input_tokens = len(inputs.input_ids[0])
        logger.info(f"Input tokens: {input_tokens}")

        # Generate translation
        # NLLB uses forced_bos_token_id to specify target language
        translated_tokens = model.generate(
            **inputs,
            forced_bos_token_id=tokenizer.lang_code_to_id[tgt_nllb_code],
            max_length=512,
            num_beams=5,
            early_stopping=True,
            no_repeat_ngram_size=3
        )

        # Decode output
        translated = tokenizer.batch_decode(
            translated_tokens,
            skip_special_tokens=True
        )[0]

        # Restore HTML tags
        translated_with_tags = HTMLTagPreserver.restore_tags(translated, tags)

        output_tokens = len(translated_tokens[0])
        total_tokens = input_tokens + output_tokens

        processing_time = int((time.time() - start_time) * 1000)

        logger.info(
            f"Translation completed in {processing_time}ms using {total_tokens} tokens "
            f"(input: {input_tokens}, output: {output_tokens})"
        )

        # Tone adjustment note: NLLB doesn't natively support tone control
        # For production, consider post-processing or fine-tuned models for tone
        if tone != "neutral":
            logger.info(f"Note: Tone '{tone}' requested but NLLB doesn't support tone control")

        return {
            "translation": translated_with_tags,
            "tokens_used": total_tokens,
            "processing_time_ms": processing_time,
            "model": "4b",
            "source_lang": source_lang,
            "target_lang": target_lang
        }

    except ValueError as e:
        logger.error(f"Validation error: {str(e)}")
        raise

    except Exception as e:
        logger.error(f"Translation error: {str(e)}", exc_info=True)
        raise RuntimeError(f"Translation failed: {str(e)}") from e


@app.function(
    image=image,
    gpu="A100",  # More powerful GPU for larger model
    timeout=600,
    memory=40960,
)
def translate_27b(
    content: str,
    source_lang: str,
    target_lang: str,
    tone: str = "neutral"
) -> Dict[str, any]:
    """
    Translate using NLLB-200 3.3B model (higher quality, slower)

    This is the larger NLLB model providing superior translation quality
    for complex texts and nuanced translations.

    Args:
        content: Text to translate
        source_lang: Source language code (ISO 639-1, e.g., "en", "es")
        target_lang: Target language code (ISO 639-1)
        tone: Translation tone (neutral, formal, casual) - Note: NLLB doesn't natively support tone

    Returns:
        Dict with translation, tokens_used, processing_time_ms, and model

    Raises:
        ValueError: If language codes are not supported
        RuntimeError: If translation fails
    """
    from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
    import time
    import torch

    start_time = time.time()

    try:
        # Validate and convert language codes
        src_nllb_code = get_nllb_code(source_lang)
        tgt_nllb_code = get_nllb_code(target_lang)

        if not src_nllb_code:
            raise ValueError(
                f"Unsupported source language: {source_lang}. "
                f"Supported languages: {', '.join(sorted(LANGUAGE_CODE_MAP.keys()))}"
            )

        if not tgt_nllb_code:
            raise ValueError(
                f"Unsupported target language: {target_lang}. "
                f"Supported languages: {', '.join(sorted(LANGUAGE_CODE_MAP.keys()))}"
            )

        logger.info(
            f"Translating with 3.3B model from {source_lang} ({src_nllb_code}) "
            f"to {target_lang} ({tgt_nllb_code})"
        )

        # Extract HTML tags
        clean_content, tags = HTMLTagPreserver.extract_tags(content)

        if not clean_content.strip():
            logger.warning("Empty content after tag extraction")
            return {
                "translation": content,  # Return original if empty
                "tokens_used": 0,
                "processing_time_ms": 0,
                "model": "27b"
            }

        # Load NLLB-200 3.3B model (cached after first run)
        model_name = "facebook/nllb-200-3.3B"
        logger.info(f"Loading model: {model_name}")

        tokenizer = AutoTokenizer.from_pretrained(
            model_name,
            src_lang=src_nllb_code,
            tgt_lang=tgt_nllb_code
        )

        # Load model with 8-bit quantization for memory efficiency
        model = AutoModelForSeq2SeqLM.from_pretrained(
            model_name,
            device_map="auto",
            load_in_8bit=True,
            torch_dtype=torch.float16
        )

        # Tokenize input
        inputs = tokenizer(
            clean_content,
            return_tensors="pt",
            max_length=1024,
            truncation=True,
            padding=True
        )

        # Move inputs to GPU
        inputs = {k: v.to("cuda") for k, v in inputs.items()}

        input_tokens = len(inputs["input_ids"][0])
        logger.info(f"Input tokens: {input_tokens}")

        # Generate translation with higher quality settings
        # NLLB uses forced_bos_token_id to specify target language
        translated_tokens = model.generate(
            **inputs,
            forced_bos_token_id=tokenizer.lang_code_to_id[tgt_nllb_code],
            max_length=1024,
            num_beams=8,  # More beams for higher quality
            early_stopping=True,
            no_repeat_ngram_size=3,
            length_penalty=1.0,
            repetition_penalty=1.2
        )

        # Decode output
        translated = tokenizer.batch_decode(
            translated_tokens,
            skip_special_tokens=True
        )[0]

        # Restore HTML tags
        translated_with_tags = HTMLTagPreserver.restore_tags(translated, tags)

        output_tokens = len(translated_tokens[0])
        total_tokens = input_tokens + output_tokens

        processing_time = int((time.time() - start_time) * 1000)

        logger.info(
            f"Translation completed in {processing_time}ms using {total_tokens} tokens "
            f"(input: {input_tokens}, output: {output_tokens})"
        )

        # Tone adjustment note: NLLB doesn't natively support tone control
        if tone != "neutral":
            logger.info(f"Note: Tone '{tone}' requested but NLLB doesn't support tone control")

        return {
            "translation": translated_with_tags,
            "tokens_used": total_tokens,
            "processing_time_ms": processing_time,
            "model": "27b",
            "source_lang": source_lang,
            "target_lang": target_lang
        }

    except ValueError as e:
        logger.error(f"Validation error: {str(e)}")
        raise

    except Exception as e:
        logger.error(f"Translation error: {str(e)}", exc_info=True)
        raise RuntimeError(f"Translation failed: {str(e)}") from e


@app.function()
@modal.web_endpoint(method="POST")
def translate(request_data: Dict) -> Dict:
    """
    Web endpoint for translation requests

    Request body:
    {
        "content": "Text to translate",
        "source_lang": "en",  // ISO 639-1 code
        "target_lang": "es",  // ISO 639-1 code
        "model": "4b" or "27b",
        "tone": "neutral" (optional, not currently supported by NLLB)
    }

    Returns:
    Success:
    {
        "translation": "Translated text",
        "tokens_used": 123,
        "processing_time_ms": 1500,
        "model": "4b",
        "source_lang": "en",
        "target_lang": "es"
    }

    Error:
    {
        "error": {
            "code": "ERROR_CODE",
            "message": "Error description"
        }
    }
    """
    try:
        # Validate request structure
        if not isinstance(request_data, dict):
            return {
                "error": {
                    "code": "INVALID_REQUEST",
                    "message": "Request body must be a JSON object"
                }
            }

        # Extract and validate parameters
        content = request_data.get("content")
        source_lang = request_data.get("source_lang")
        target_lang = request_data.get("target_lang")
        model = request_data.get("model", "4b")
        tone = request_data.get("tone", "neutral")

        # Validate required fields
        if not content:
            return {
                "error": {
                    "code": "INVALID_REQUEST",
                    "message": "Missing required field: content"
                }
            }

        if not source_lang:
            return {
                "error": {
                    "code": "INVALID_REQUEST",
                    "message": "Missing required field: source_lang"
                }
            }

        if not target_lang:
            return {
                "error": {
                    "code": "INVALID_REQUEST",
                    "message": "Missing required field: target_lang"
                }
            }

        # Validate model parameter
        if model not in ["4b", "27b"]:
            return {
                "error": {
                    "code": "INVALID_MODEL",
                    "message": f"Invalid model: {model}. Must be '4b' or '27b'"
                }
            }

        # Validate content length
        if len(content) > 5000:
            return {
                "error": {
                    "code": "CONTENT_TOO_LONG",
                    "message": "Content exceeds maximum length of 5000 characters"
                }
            }

        logger.info(
            f"Translation request: {source_lang} -> {target_lang}, "
            f"model={model}, content_length={len(content)}"
        )

        # Route to appropriate model
        if model == "27b":
            result = translate_27b.remote(content, source_lang, target_lang, tone)
        else:
            result = translate_4b.remote(content, source_lang, target_lang, tone)

        return result

    except ValueError as e:
        # Language validation errors
        logger.error(f"Validation error: {str(e)}")
        return {
            "error": {
                "code": "UNSUPPORTED_LANGUAGE",
                "message": str(e)
            }
        }

    except RuntimeError as e:
        # Translation runtime errors
        logger.error(f"Runtime error: {str(e)}")
        return {
            "error": {
                "code": "TRANSLATION_FAILED",
                "message": str(e)
            }
        }

    except Exception as e:
        # Unexpected errors
        logger.error(f"Endpoint error: {str(e)}", exc_info=True)
        return {
            "error": {
                "code": "INTERNAL_ERROR",
                "message": "An unexpected error occurred. Please try again."
            }
        }


@app.function()
@modal.web_endpoint(method="GET")
def health() -> Dict:
    """
    Health check endpoint

    Returns:
    {
        "status": "healthy",
        "service": "translate-nllb",
        "models": ["4b", "27b"],
        "supported_languages": 60
    }
    """
    return {
        "status": "healthy",
        "service": "translate-nllb",
        "models": ["4b", "27b"],
        "model_details": {
            "4b": {
                "name": "facebook/nllb-200-distilled-600M",
                "description": "Fast, lightweight translation for 200+ languages",
                "gpu": "T4"
            },
            "27b": {
                "name": "facebook/nllb-200-3.3B",
                "description": "High-quality translation for complex texts",
                "gpu": "A100"
            }
        },
        "supported_languages": len(LANGUAGE_CODE_MAP)
    }


@app.function()
@modal.web_endpoint(method="GET")
def languages() -> Dict:
    """
    Get list of supported languages

    Returns:
    {
        "languages": {
            "en": "eng_Latn",
            "es": "spa_Latn",
            ...
        }
    }
    """
    return {
        "languages": LANGUAGE_CODE_MAP,
        "total": len(LANGUAGE_CODE_MAP)
    }