"""
GovBridge India — Translation Module (Sprint 18 v2)

Translation Architecture (prioritized fallback chain):
  Layer 1: Bhashini API (government-grade, when API key is approved)
  Layer 2: Groq LLM Translation (llama-3.3-70b-versatile — already in our stack)
  Layer 3: Graceful degradation (return original text)

WHY GROQ LLM INSTEAD OF IndicTrans2:
  - IndicTrans2 (200M params) is fundamentally incompatible with transformers>=4.48
    (required by ModernBERT/Ettin Reranker). Config, tokenizer, and ONNX layers all break.
  - Groq's llama-3.3-70b-versatile (70B params) provides SUPERIOR translation quality
    for all 10 supported Indic languages — 350x more parameters.
  - Zero additional cost: we already have the Groq API key.
  - Saves 1.6GB RAM on the 16GB HF Spaces instance.
  - Zero dependency conflicts. Zero model loading time.
  - Translation latency: ~200ms via Groq (vs ~2-4s for local IndicTrans2 on CPU).
"""

import httpx
import os
from groq import Groq

LANGUAGE_CODES = {
    "hindi": "hi", "tamil": "ta", "bengali": "bn",
    "telugu": "te", "marathi": "mr", "gujarati": "gu",
    "kannada": "kn", "malayalam": "ml", "punjabi": "pa",
    "odia": "or",
    "english": "en"
}

# Full language names for LLM prompt clarity
LANGUAGE_NAMES = {
    "hi": "Hindi", "ta": "Tamil", "bn": "Bengali",
    "te": "Telugu", "mr": "Marathi", "gu": "Gujarati",
    "kn": "Kannada", "ml": "Malayalam", "pa": "Punjabi",
    "or": "Odia", "en": "English",
    "hindi": "Hindi", "tamil": "Tamil", "bengali": "Bengali",
    "telugu": "Telugu", "marathi": "Marathi", "gujarati": "Gujarati",
    "kannada": "Kannada", "malayalam": "Malayalam", "punjabi": "Punjabi",
    "odia": "Odia", "english": "English"
}

BHASHINI_USER_ID = os.environ.get("BHASHINI_USER_ID")
BHASHINI_API_KEY = os.environ.get("BHASHINI_API_KEY")
BHASHINI_INFERENCE_URL = "https://dhruva-api.bhashini.gov.in/services/inference/pipeline"

# Groq client for LLM translation (Layer 2)
_groq_client = None

def _get_groq_client() -> Groq:
    """Lazy-init Groq client for translation."""
    global _groq_client
    if _groq_client is None:
        api_key = os.environ.get("GROQ_API_KEY")
        if not api_key:
            raise RuntimeError("GROQ_API_KEY not set")
        _groq_client = Groq(api_key=api_key)
    return _groq_client


def _groq_translate(text: str, source_lang: str, target_lang: str) -> str:
    """
    Translate text using Groq's llama-3.3-70b-versatile.
    
    ARCHITECTURAL NOTE: The LLM is used STRICTLY as a translator here.
    It receives a system prompt that constrains it to output ONLY the translation,
    with no explanations, no additions, no hallucinations.
    """
    src_name = LANGUAGE_NAMES.get(source_lang.lower(), source_lang)
    tgt_name = LANGUAGE_NAMES.get(target_lang.lower(), target_lang)
    
    client = _get_groq_client()
    
    response = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=[
            {
                "role": "system",
                "content": (
                    f"You are a professional translator. Translate the following text "
                    f"from {src_name} to {tgt_name}. "
                    f"Output ONLY the translated text. No explanations, no notes, "
                    f"no quotation marks, no prefixes like 'Translation:'. "
                    f"Preserve the original meaning, tone, and formatting exactly."
                )
            },
            {
                "role": "user",
                "content": text
            }
        ],
        temperature=0.1,  # Low temperature for consistent, accurate translations
        max_tokens=1024,
        stream=False
    )
    
    result = (response.choices[0].message.content or "").strip()
    
    # Clean up any accidental prefixes the LLM might add
    for prefix in ["Translation:", "translation:", "Translated:", "translated:"]:
        if result.startswith(prefix):
            result = result[len(prefix):].strip()
    
    return result


async def translate_text(
    text: str,
    source_lang: str,
    target_lang: str,
    fallback: bool = True
) -> str:
    """
    Main translation entry point. Used by api.py for the full pipeline:
    User Query (Indic) → English → RAG Search → English Answer → Indic Response
    
    Fallback chain:
      1. Bhashini API (when approved)
      2. Groq LLM (llama-3.3-70b — immediate, free, high quality)
      3. Return original text (graceful degradation)
    """
    if source_lang.lower() == target_lang.lower():
        return text

    # Layer 1: Bhashini API (government-grade translation)
    if BHASHINI_USER_ID and BHASHINI_API_KEY:
        try:
            src = LANGUAGE_CODES.get(source_lang.lower().strip(), source_lang)
            tgt = LANGUAGE_CODES.get(target_lang.lower().strip(), target_lang)
            async with httpx.AsyncClient(timeout=3.5) as client:
                payload = {
                    "pipelineTasks": [{"taskType": "translation",
                        "config": {"language": {
                            "sourceLanguage": src,
                            "targetLanguage": tgt
                        }}}],
                    "inputData": {"input": [{"source": text}]}
                }
                headers = {
                    "userID": BHASHINI_USER_ID,
                    "ulcaApiKey": BHASHINI_API_KEY,
                    "Content-Type": "application/json"
                }
                resp = await client.post(
                    BHASHINI_INFERENCE_URL,
                    json=payload,
                    headers=headers
                )
                resp.raise_for_status()
                result = resp.json()["pipelineResponse"][0]["output"][0]["target"]
                print(f"✅ Bhashini: {source_lang} → {target_lang}")
                return result
        except Exception as e:
            print(f"⚠️ Bhashini failed: {e}, trying Groq LLM translation")

    # Layer 2: Groq LLM Translation (zero cost — already in our stack)
    try:
        result = _groq_translate(text, source_lang, target_lang)
        print(f"✅ Groq LLM: {source_lang} → {target_lang} | '{text[:30]}' → '{result[:30]}'")
        return result
    except Exception as e:
        print(f"⚠️ Groq translation failed: {e}")

    # Layer 3: Graceful degradation
    print(f"⚠️ All translation failed — returning original")
    return text