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

This service deploys the GemmaTranslate model on Modal's infrastructure with persistent model storage.
The model is downloaded once and cached in a Modal Volume for fast subsequent loads.

Features:
- Persistent model storage (no re-download on every run)
- HTML tag preservation during translation
- Token usage tracking
- Production-ready error handling
- Comprehensive logging
"""

import modal
import re
from typing import Dict, Tuple
import logging
import time

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

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

# Create persistent volume for model storage
# This volume persists across function calls, so models are downloaded only once
model_volume = modal.Volume.from_name("gemmatranslate-models", create_if_missing=True)

# Docker image with ML dependencies
image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install(
        "transformers>=4.46.0",  # Upgraded for Gemma tokenizer compatibility
        "torch>=2.6.0",  # Required for Gemma3 masking functions
        "sentencepiece==0.2.0",
        "protobuf==5.28.2",
        "accelerate>=0.34.2",
        "safetensors>=0.4.3",
        "pillow",  # Required for AutoImageProcessor
        "fastapi[standard]",
    )
)

# Model cache path in the volume
MODEL_CACHE_PATH = "/models"


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)
        """
        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

        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
        for placeholder, tag in tags.items():
            result = result.replace(placeholder, tag)
        return result


@app.function(
    image=image,
    gpu="T4",  # T4 GPU for TranslateGemma 4B
    timeout=300,  # 5 minutes for translation
    memory=20480,  # 20GB memory for 4B model
    volumes={MODEL_CACHE_PATH: model_volume},  # Mount persistent volume
    min_containers=1,  # Keep 1 container warm to avoid cold starts
    scaledown_window=60,  # Keep container alive for 60 seconds after last request
)
def translate_gemma(
    content: str,
    source_lang: str,
    target_lang: str,
    tone: str = "neutral"
) -> Dict[str, any]:
    """
    Translate using Google's TranslateGemma 4B instruction-tuned model

    TranslateGemma is a family of open machine translation models from Google,
    based on Gemma 2 and trained on a mix of parallel and monolingual examples.

    The model is stored in a Modal Volume and only downloaded once (~10GB).
    Subsequent runs load from the cached model in 1-2 seconds.

    Args:
        content: Text to translate
        source_lang: Source language code (e.g., "en", "es", "fr")
        target_lang: Target language code
        tone: Translation tone (neutral, formal, casual) - note: model doesn't explicitly support tone

    Returns:
        Dict with translation, tokens_used, processing_time_ms, and model
    """
    from transformers import AutoProcessor, AutoModelForImageTextToText
    import os
    import torch

    start_time = time.time()

    try:
        logger.info(f"Translating from {source_lang} to {target_lang} using TranslateGemma 4B")

        # 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,
                "tokens_used": 0,
                "processing_time_ms": 0,
                "model": "gemmatranslate"
            }

        # TranslateGemma 4B Instruction-Tuned model from Google
        model_name = "google/translategemma-4b-it"
        hf_token = "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"

        # Check if model is already cached in volume
        model_cache_dir = os.path.join(MODEL_CACHE_PATH, model_name.replace("/", "--"))

        logger.info(f"Model cache directory: {model_cache_dir}")

        if os.path.exists(model_cache_dir):
            logger.info(f"Loading cached model from {model_cache_dir}")
            load_start = time.time()

            processor = AutoProcessor.from_pretrained(
                model_cache_dir,
                local_files_only=True
            )
            model = AutoModelForImageTextToText.from_pretrained(
                model_cache_dir,
                local_files_only=True,
                torch_dtype=torch.bfloat16,
                device_map="auto"
            )

            logger.info(f"Model loaded from cache in {time.time() - load_start:.2f}s")
        else:
            logger.info(f"Downloading {model_name} model (first time only, ~10GB)")
            logger.info("This will take 5-10 minutes depending on your connection")
            download_start = time.time()

            # Download model with HuggingFace token
            processor = AutoProcessor.from_pretrained(
                model_name,
                token=hf_token
            )
            model = AutoModelForImageTextToText.from_pretrained(
                model_name,
                token=hf_token,
                torch_dtype=torch.bfloat16,
                device_map="auto"
            )

            # Save to persistent volume
            os.makedirs(model_cache_dir, exist_ok=True)
            processor.save_pretrained(model_cache_dir)
            model.save_pretrained(model_cache_dir)

            # Commit changes to volume so they persist
            model_volume.commit()

            logger.info(f"Model downloaded and cached in {time.time() - download_start:.2f}s")

        # Prepare translation using TranslateGemma chat template
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "source_lang_code": source_lang,
                        "target_lang_code": target_lang,
                        "text": clean_content
                    }
                ]
            }
        ]

        # Apply chat template
        inputs = processor.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_dict=True,
            return_tensors="pt"
        ).to(model.device, dtype=torch.bfloat16)

        input_tokens = inputs['input_ids'].shape[1]
        logger.info(f"Input tokens: {input_tokens}")

        # Generate translation
        with torch.inference_mode():
            output = model.generate(
                **inputs,
                max_new_tokens=512,
                do_sample=False,
            )

        # Decode output - CRITICAL: Remove input tokens first (official HuggingFace approach)
        input_len = len(inputs['input_ids'][0])
        generation = output[0][input_len:]
        translated = processor.decode(generation, skip_special_tokens=True)

        translated = translated.strip()

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

        output_tokens = len(generation)
        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})"
        )

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

    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)
@modal.fastapi_endpoint(method="POST")
def translate(request_data: Dict) -> Dict:
    """
    Web endpoint for translation requests

    Request body:
    {
        "content": "Text to translate",
        "source_lang": "en",
        "target_lang": "es",
        "model": "gemmatranslate",
        "tone": "neutral" (optional)
    }

    Returns:
    {
        "translation": "Translated text",
        "tokens_used": 123,
        "processing_time_ms": 1500,
        "model": "gemmatranslate"
    }
    """
    try:
        # Validate request structure
        if not isinstance(request_data, dict):
            return {
                "error": {
                    "code": "INVALID_REQUEST",
                    "message": "Request body must be a JSON object"
                }
            }

        # Extract parameters
        content = request_data.get("content")
        source_lang = request_data.get("source_lang")
        target_lang = request_data.get("target_lang")
        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 or not target_lang:
            return {
                "error": {
                    "code": "INVALID_REQUEST",
                    "message": "Missing required fields: source_lang and target_lang"
                }
            }

        # 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"content_length={len(content)}"
        )

        # Call translation function
        result = translate_gemma.remote(content, source_lang, target_lang, tone)

        return result

    except Exception as e:
        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(image=image)
@modal.fastapi_endpoint(method="GET")
def health() -> Dict:
    """
    Health check endpoint

    Returns service status and model information
    """
    return {
        "status": "healthy",
        "service": "translate-gemma",
        "model": {
            "name": "google/translategemma-4b-it",
            "description": "Google TranslateGemma 4B instruction-tuned translation model",
            "size": "~10GB",
            "gpu": "T4",
            "memory": "20GB",
            "persistent_storage": True,
            "cache_location": MODEL_CACHE_PATH
        },
        "features": [
            "HTML tag preservation",
            "Persistent model caching",
            "Multi-language support",
            "Production-ready error handling"
        ]
    }


@app.function(
    volumes={MODEL_CACHE_PATH: model_volume},
)
def check_model_cache() -> Dict:
    """
    Check what models are cached in the volume

    Returns:
    {
        "cached_models": [...],
        "cache_size_mb": 1234
    }
    """
    import os

    cached_models = []
    total_size = 0

    if os.path.exists(MODEL_CACHE_PATH):
        for item in os.listdir(MODEL_CACHE_PATH):
            item_path = os.path.join(MODEL_CACHE_PATH, item)
            if os.path.isdir(item_path):
                # Calculate directory size
                dir_size = sum(
                    os.path.getsize(os.path.join(dirpath, filename))
                    for dirpath, dirnames, filenames in os.walk(item_path)
                    for filename in filenames
                )
                total_size += dir_size
                cached_models.append({
                    "name": item.replace("--", "/"),
                    "size_mb": round(dir_size / (1024 * 1024), 2)
                })

    return {
        "cached_models": cached_models,
        "cache_size_mb": round(total_size / (1024 * 1024), 2),
        "cache_path": MODEL_CACHE_PATH
    }


# Local entrypoint for testing
@app.local_entrypoint()
def test_translation():
    """
    Test the translation service locally

    Run with: modal run gemmatranslate_service.py

    This will download the model on first run (~10GB, 5-10 minutes).
    Subsequent runs load from cache in 1-2 seconds.
    """
    test_text = "The quick brown fox jumps over the lazy dog. This is a simple test sentence."

    print(f"\n{'='*70}")
    print(f"  TranslateGemma 4B Translation Test")
    print(f"{'='*70}\n")

    print(f"📝 Input text:")
    print(f"   {test_text}\n")
    print(f"🌐 Translation: en -> es\n")

    # Check cache first
    print(f"📦 Checking model cache...")
    cache_info = check_model_cache.remote()

    if cache_info['cached_models']:
        print(f"✅ Model cached ({cache_info['cache_size_mb']} MB)")
        print(f"   Expected time: 5-10 seconds\n")
    else:
        print(f"⚠️  No cached model found")
        print(f"   First run will download ~10GB model (5-10 minutes)")
        print(f"   Please wait...\n")

    # Run translation
    print(f"🔄 Starting translation...")
    start = time.time()

    result = translate_gemma.remote(test_text, "en", "es", "neutral")

    elapsed = time.time() - start

    print(f"\n{'='*70}")
    print(f"  Translation Results")
    print(f"{'='*70}\n")

    print(f"✅ Translation successful!\n")
    print(f"📝 Original:")
    print(f"   {test_text}\n")
    print(f"🌐 Translated:")
    print(f"   {result['translation']}\n")
    print(f"📊 Metrics:")
    print(f"   Tokens used: {result['tokens_used']}")
    print(f"   Processing time (server): {result['processing_time_ms']}ms")
    print(f"   Total time (with network): {elapsed*1000:.0f}ms")
    print(f"   Model: {result['model']}\n")

    print(f"{'='*70}")
    print(f"  Test completed successfully! ✅")
    print(f"{'='*70}\n")