# src/tools/audio_utils.py
import re
import os
from dotenv import load_dotenv
import zipfile
import logging
logger = logging.getLogger(__name__)
import requests
from indic_transliteration.sanscript import transliterate
from indic_transliteration import sanscript
import json
import time
from google.generativeai import GenerativeModel, configure
from google.generativeai.types import GenerationConfig
try:
    from google import genai
    from google.genai import types
except ImportError:
    genai = None
    types = None
    print("Warning: google-genai package not installed. Gemini TTS will not work. Install with: pip install google-genai")
from src.tools.prompt_utils import get_regional_translation_prompt, get_named_entity_identification_prompt, get_pun_translation_prompt

from typing import Literal

import subprocess
import wave
import io
from pydub import AudioSegment
import tempfile
import shutil
import traceback

# Exports - ensure these names are importable
__all__ = [
    "audio_fn_from_string",
    "audio_fn",
    "pun_audio_fn_from_string",
    "trim_audio_to_max_duration",
    "pad_audio_to_duration",
    "combine_audio_with_video",
    "pun_generate_segment_audios",
    "pun_check_audio_files",
    "pun_get_audio_duration"
]

# Load environment variables and configure Gemini
load_dotenv()
GEMINI_API_KEY = os.environ.get('GEMINI_API_KEY')

# Optionally allow overriding technical terms through env var (comma separated)
_env_terms = os.environ.get("TECHNICAL_TERMS", "")
DEFAULT_TECHNICAL_TERMS = [
    "CNN", "RNN", "LSTM", "GRU", "Transformer", "Transformer-based", "BERT", "GPT",
    "OpenAI", "Gemini", "ffmpeg", "Playwright", "TensorFlow", "PyTorch", "scikit-learn",
    "API", "MSE", "R-squared", "R2", "accuracy", "precision", "recall", "F1", "BLEU",
    "AUC", "ROC", "Linear Regression", "logistic regression", "SGD", "Adam", "epoch",
    "batch", "GPU", "TPU", "embedding", "token", "tokenizer"
]
TECHNICAL_TERMS = [t.strip() for t in (_env_terms.split(",") if _env_terms else []) if t.strip()] or DEFAULT_TECHNICAL_TERMS

# Configure Gemini API globally if key present
if GEMINI_API_KEY:
    try:
        configure(api_key=GEMINI_API_KEY)
        gemini = GenerativeModel('gemini-2.5-flash')
    except Exception:
        gemini = None
else:
    gemini = None
    print('Warning: GEMINI_API_KEY not found in environment variables. Some functions may not work.')

LANGUAGE_TO_SCRIPT_MAP = {
    "hindi": sanscript.DEVANAGARI,
    "punjabi": sanscript.GURMUKHI,
    "gujarati": sanscript.GUJARATI,
    "marathi": sanscript.DEVANAGARI,
    "kannada": sanscript.KANNADA,
}

# ---------------- Helpers ----------------
def get_temp_dir():
    return tempfile.gettempdir()

def pun_remove_special_characters(s: str) -> str:
    if s is None:
        return s
    s = s.replace('`', '').replace('\xa0', '').replace(' ', '')
    s = s.replace('\u200b', '')  # zero-width space
    return s

def _clean_text_for_tts(text: str) -> str:
    """
    Light cleanup for TTS input.
    Removes double-quote characters and backticks only.
    Apostrophes inside contractions (we'll, don't, you're) are preserved so the
    TTS model pronounces them correctly. Only standalone/surrounding apostrophes
    that are acting as quote marks are removed.
    """
    if text is None:
        return ""
    # Remove straight and curly double quotes
    text = re.sub(r'["\u201c\u201d\u201e\u201f]', '', text)
    # Remove backticks
    text = text.replace('`', '')
    # Remove apostrophes used as standalone quote marks (not inside a word).
    # Pattern: apostrophe NOT flanked by word characters on both sides.
    text = re.sub(r"(?<!\w)'|'(?!\w)", '', text)
    # Fix orphaned sentence-terminal punctuation after quote removal.
    # When a quoted phrase like "karo!" is stripped of its quotes, the ! sits
    # mid-sentence (e.g. "...karo! she says..."). Gemini TTS interprets ! as a
    # hard stop and goes silent for the rest of the narration.
    # Detection: ! or ? followed by whitespace then a lowercase letter means it
    # is mid-sentence (not a real sentence boundary). Replace with a comma.
    text = re.sub(r'([!?])(\s+[a-z])', r',\2', text)
    # Remove markdown bold/italic markers that may appear in generated narration
    text = re.sub(r'\*{1,3}([^*]+)\*{1,3}', r'\1', text)
    # Normalize whitespace
    text = ' '.join(text.split())
    # Replace slash with spoken "or"
    text = text.replace('/', ' or ')
    return text.strip()


def _extract_audio_bytes(response) -> bytes:
    """
    Extract audio bytes from Gemini TTS responses.
    Supports multiple SDK response formats.
    """
    chunks = []

    def add_chunk(value):
        if not value:
            return

        if isinstance(value, (bytes, bytearray)):
            chunks.append(bytes(value))
            return

        if isinstance(value, str):
            try:
                import base64
                decoded = base64.b64decode(value)
                if decoded:
                    chunks.append(decoded)
            except Exception:
                pass

    # Candidate → Content → Parts
    for candidate in getattr(response, "candidates", []) or []:
        content = getattr(candidate, "content", None)
        parts = getattr(content, "parts", []) if content else []

        for part in parts:
            inline_data = getattr(part, "inline_data", None)

            if inline_data:
                add_chunk(getattr(inline_data, "data", None))
                add_chunk(getattr(inline_data, "_raw_data", None))

            audio_data = getattr(part, "audio_data", None)

            if audio_data:
                add_chunk(getattr(audio_data, "data", None))
                add_chunk(getattr(audio_data, "_raw_data", None))

    # Top-level fallback
    for attr in ["audio_data", "audio", "data"]:
        obj = getattr(response, attr, None)

        if obj:
            add_chunk(getattr(obj, "data", None))
            add_chunk(getattr(obj, "_raw_data", None))

    return b"".join(chunks)

# Utility: create safe markers for preserving tokens
def _make_marker(term: str) -> str:
    # safe marker unlikely to be touched by translation
    safe = "__KEEPTERM__" + re.sub(r'\W+', '_', term) + "__"
    return safe

def _protect_technical_terms(text: str) -> (str, dict):
    """
    Replace occurrences of technical terms with markers, return (protected_text, mapping)
    mapping: marker -> original term
    """
    mapping = {}
    if not text:
        return text, mapping
    # Sort by length to protect longest terms first
    terms_sorted = sorted(TECHNICAL_TERMS, key=lambda x: -len(x))
    protected = text
    for term in terms_sorted:
        # Word-boundary, case-insensitive
        # allow terms containing spaces (e.g., "Linear Regression")
        pattern = re.compile(re.escape(term), re.IGNORECASE)
        def _repl(m):
            orig = m.group(0)
            marker = _make_marker(term)
            mapping[marker] = orig  # store original matched case
            return marker
        protected = pattern.sub(_repl, protected)
    return protected, mapping

def _restore_markers(text: str, mapping: dict) -> str:
    if not mapping:
        return text
    restored = text
    for marker, orig in mapping.items():
        restored = restored.replace(marker, orig)
    return restored

# Heuristic to decide whether to keep a token as English / technical
def _looks_like_english_token(token: str) -> bool:
    if not token:
        return False
    # If it's explicitly in technical terms (case-insensitive)
    for t in TECHNICAL_TERMS:
        if token.lower() == t.lower():
            return True
    # Mostly ASCII chars
    ascii_ratio = sum(1 for ch in token if ord(ch) < 128) / max(1, len(token))
    if ascii_ratio > 0.85:
        return True
    # if contains digits or common code symbols or CamelCase
    if re.search(r'[A-Z][a-z]+[A-Z]', token) or re.search(r'[\d\(\)\=\+\-_/\.]', token):
        return True
    # Contains dot notation or file extensions
    if '.' in token and len(token) <= 40:
        return True
    return False

# --- Keep English named entities as-is for Hinglish ---
def transliterate_english_to_native_script(sentence, target_language):
    if sentence is None:
        return ""
    if isinstance(sentence, str):
        # If the token looks like an english/technical token, keep as-is
        if _looks_like_english_token(sentence):
            return sentence.strip()
        ascii_ratio = sum(1 for ch in sentence if ord(ch) < 128) / max(1, len(sentence))
        if ascii_ratio > 0.8:
            return sentence.strip()
    script = LANGUAGE_TO_SCRIPT_MAP.get((target_language or "").lower())
    if not script:
        return sentence
    try:
        return transliterate(sentence, sanscript.ITRANS, script)
    except Exception:
        return sentence

# ---------------- Translation helpers (Gemini) ----------------
def translate_to_regional(target_language, text):
    if not GEMINI_API_KEY or not gemini:
        print("Warning: Gemini not configured; translate_to_regional returning original text.")
        return text
    try:
        # protect technical terms with markers so translator won't alter them
        protected_text, mapping = _protect_technical_terms(text or "")
        cleaned_text = re.sub(r'``````', '', protected_text or '', flags=re.DOTALL)
        cleaned_text = re.sub(r'```.*?```', '', cleaned_text, flags=re.DOTALL)
        prompt = get_regional_translation_prompt(target_language, cleaned_text)
        response = gemini.generate_content(
            prompt,
            generation_config=GenerationConfig(response_mime_type="application/json")
        )
        raw_text = response.text
        # attempt parse JSON
        try:
            result = json.loads(raw_text)
            translated = result.get("translation", text)
            # restore protected terms
            return _restore_markers(translated, mapping)
        except json.JSONDecodeError:
            match = re.search(r"\{.*\}", raw_text, re.DOTALL)
            if match:
                try:
                    result = json.loads(match.group(0))
                    translated = result.get("translation", text)
                    return _restore_markers(translated, mapping)
                except Exception:
                    pass
            if (target_language or "").lower() not in ["english", "en"]:
                return _restore_markers(raw_text.strip(), mapping)
            return text
    except Exception as e:
        print(f":x: Gemini error: {str(e)}")
        try:
            if 'raw_text' in locals() and (target_language or "").lower() not in ["english", "en"]:
                return _restore_markers(raw_text.strip(), {})
        except Exception:
            pass
        return text

def pun_translate_to_regional(target_language: str, text: str) -> str:
    if not GEMINI_API_KEY or not gemini:
        print("Warning: Gemini not configured; pun_translate_to_regional returning original text.")
        return text
    try:
        protected_text, mapping = _protect_technical_terms(text or "")
        cleaned_text = re.sub(r'```python.*?```', '', protected_text or '', flags=re.DOTALL)
        cleaned_text = re.sub(r'```.*?```', '', cleaned_text, flags=re.DOTALL)
        prompt = get_pun_translation_prompt(target_language, cleaned_text)
        response = gemini.generate_content(prompt)
        raw_text = response.text.strip()
        match = re.search(r'\{.*\}', raw_text, re.DOTALL)
        if match:
            try:
                result = json.loads(match.group(0))
                translated = result.get('translation', text)
                return _restore_markers(translated, mapping)
            except Exception:
                pass
        return _restore_markers(raw_text, mapping)
    except Exception as e:
        print(f':x: Gemini error: {str(e)}')
        try:
            if 'raw_text' in locals() and (target_language or "").lower() not in ["english", "en"]:
                return _restore_markers(raw_text.strip(), {})
        except Exception:
            pass
        return text

# ---------------- Named entity identification ----------------
def indentify_named_entities(input_text):
    prompt = get_named_entity_identification_prompt()
    if not GEMINI_API_KEY or not gemini:
        print('Warning: GEMINI_API_KEY not found. Returning input wrapped in triple quotes.')
        return f"'''{input_text}'''"
    try:
        response = gemini.generate_content(f"{prompt}\n\nText to process:\n{input_text}")
        oration = response.text
        return oration
    except Exception as e:
        print(f'Error identifying named entities with Gemini: {e}')
        return f"'''{input_text}'''"

def pun_indentify_named_entities(input_text: str) -> str:
    prompt = get_named_entity_identification_prompt()
    if not GEMINI_API_KEY or not gemini:
        print('Warning: GEMINI_API_KEY not found. Returning input wrapped in triple quotes.')
        return f"'''{input_text}'''"
    try:
        response = gemini.generate_content(f"{prompt}\n\nText to process:\n{input_text}")
        oration = response.text
        return oration
    except Exception as e:
        print(f'Error identifying named entities with Gemini: {e}')
        return f"'''{input_text}'''"

# ---------------- Text -> dictionary parsing (shared) ----------------
def process_text_into_dictionary(input_text):
    main_text_match = re.search(r"'''(.*?)'''", input_text, re.DOTALL)
    if not main_text_match:
        raise ValueError("No text found within triple single quotes.")
    main_text = main_text_match.group(1).strip()
    pattern = r"(<named_entities>(.*?)</named_entities>)|([^<]+)"
    result = {}
    hindi_counter = 1
    english_counter = 1
    for match in re.finditer(pattern, main_text):
        if match.group(2):
            result[f"english_{english_counter}"] = match.group(2).strip()
            english_counter += 1
        elif match.group(3):
            hindi_text = match.group(3).strip()
            if hindi_text:
                result[f"hindi_{hindi_counter}"] = hindi_text
                hindi_counter += 1
    return result

def pun_process_text_into_dictionary(input_text: str) -> dict:
    main_text_match = re.search(r"'''(.*?)'''", input_text, re.DOTALL)
    if not main_text_match:
        main_text = input_text.strip()
    else:
        main_text = main_text_match.group(1).strip()
    pattern = r'(<named_entities>(.*?)</named_entities>)|([^<]+)'
    result = {}
    hindi_counter = 1
    english_counter = 1
    for match in re.finditer(pattern, main_text):
        if match.group(2):
            result[f'english_{english_counter}'] = match.group(2).strip()
            english_counter += 1
        elif match.group(3):
            hindi_text = match.group(3).strip()
            if hindi_text:
                result[f'hindi_{hindi_counter}'] = hindi_text
                hindi_counter += 1
    return result

# ---------------- Dictionary -> final text (Hinglish behavior) ----------------
def process_dictionary(input_dict, target_language):
    outputs = []
    for key, value in sorted(input_dict.items(), key=lambda kv: kv[0]):
        if key.startswith("hindi") or key.startswith("regional") or key.startswith("marathi") or key.startswith("kannada"):
            # Protect technical terms, translate, then restore
            translated = translate_to_regional(target_language, value)
            # After translation, ensure technical terms remain original if translation altered them
            # Best-effort: try restoring markers (translate_to_regional already restores markers)
            outputs.append(translated.strip())
        elif key.startswith("english"):
            # keep english as-is (NER tagged)
            outputs.append(value.strip())
        else:
            outputs.append(value.strip())
    return " ".join([o for o in outputs if o])

def pun_process_dictionary(input_dict: dict, target_language: str) -> str:
    outputs = []
    for key, value in input_dict.items():
        if key.startswith('hindi'):
            outputs.append(pun_translate_to_regional(text=value, target_language=target_language))
        elif key.startswith('english'):
            outputs.append(value.strip())
    return ' '.join(outputs)

def pun_convert_english_to_devanaigiri_script(sentence):
    try:
        devanagari_text = transliterate(sentence, sanscript.ITRANS, sanscript.DEVANAGARI)
        return devanagari_text
    except Exception:
        return sentence

# ---------------- TTS helpers for retry / chunking ----------------

class _TTSContentRejected(Exception):
    """
    Raised when Gemini TTS returns FinishReason.OTHER with no audio payload.
    This means the model actively rejected the input content, so retrying
    with the identical text will not help — the caller must transform the text.
    """
    pass


def _normalize_text_for_tts_retry(text: str, attempt: int) -> str:
    """
    Apply progressive, generic text normalizations before a TTS retry
    that follows a content-rejection (FinishReason.OTHER).

    Rules are format-based — nothing is hardcoded to specific words.
    Each attempt produces meaningfully different text so the model has a real
    chance to accept content it rejected on the previous attempt.
    """
    normalized = text
    # Remove comma grouping from numbers: 5,999 → 5999, 1,00,000 → 100000
    # This is the most common trigger for Gemini TTS content rejection on financial text.
    normalized = re.sub(r'(?<!\d)(\d{1,3}),(\d{3})(?!\d)', r'\1\2', normalized)
    # Indian lakh-style formatting: 1,00,000
    normalized = re.sub(r'(?<!\d)(\d{1,2}),(\d{2}),(\d{3})(?!\d)', r'\1\2\3', normalized)
    if attempt >= 2:
        # Strip Devanagari Unicode block (U+0900–U+097F).
        # Mixed Devanagari + Romanized Hindi in the same utterance can trigger
        # Gemini TTS content rejection even when each language alone would pass.
        # The Romanized Hinglish tokens carry the same spoken meaning, so
        # removing the Devanagari characters does not lose the narration intent.
        normalized = re.sub(r'[\u0900-\u097F]+', '', normalized)
        # Strip markdown bold/italic markers (**text** or *text*) that the
        # script generator occasionally emits — these can confuse the model.
        normalized = re.sub(r'\*{1,3}([^*]+)\*{1,3}', r'\1', normalized)
        # Re-normalize whitespace after character removal
        normalized = ' '.join(normalized.split())
    if attempt >= 3:
        # Last full-text attempt: also replace ? in Hinglish questions with a period
        # so the model treats them as statements rather than interrogative sequences.
        normalized = re.sub(r'\?', '.', normalized)
    return normalized


def _generate_tts_in_chunks(
    text: str,
    voice_name: str,
    output_path: str,
    tts_gender: str = None
) -> str:
    """
    Sentence-level chunked TTS generation used as a last-resort fallback.

    When the full text is consistently rejected by Gemini TTS, split it at
    sentence boundaries and generate audio for each sentence independently,
    then concatenate. Sentences that still fail are silently skipped so the
    rest of the beat audio is preserved rather than crashing the pipeline.
    """
    sentences = [
        s.strip()
        for s in re.split(r'(?<=[.!?])\s+', text.strip())
        if s.strip()
    ]

    if len(sentences) <= 1:
        # Cannot split further — apply normalization and try once more.
        normalized = _normalize_text_for_tts_retry(text, 3)
        return generate_english_tts_with_gemini(normalized, voice_name, output_path, tts_gender)

    temp_dir = tempfile.mkdtemp()
    chunk_paths = []
    try:
        for i, sentence in enumerate(sentences):
            chunk_path = os.path.join(temp_dir, f"tts_chunk_{i:03d}.mp3")
            norm_sentence = _normalize_text_for_tts_retry(sentence, 3)
            for chunk_attempt in range(2):
                try:
                    result = generate_english_tts_with_gemini(
                        norm_sentence, voice_name, chunk_path, tts_gender
                    )
                    if result and os.path.exists(result):
                        chunk_paths.append(chunk_path)
                        break
                except _TTSContentRejected:
                    logger.warning(
                        "[TTS] Chunk %d content-rejected, skipping: %s...",
                        i, sentence[:60]
                    )
                    break
                except Exception:
                    time.sleep(1.0)

        if not chunk_paths:
            raise ValueError("All sentence chunks failed TTS generation in chunked fallback.")

        # Concatenate all successful chunk audio files
        combined = AudioSegment.from_file(chunk_paths[0])
        for p in chunk_paths[1:]:
            combined = combined + AudioSegment.from_file(p)
        combined.export(output_path, format="mp3", bitrate="128k")
        logger.info(
            "[TTS] Chunked TTS produced %d/%d sentences → %s",
            len(chunk_paths), len(sentences), output_path
        )
        return output_path
    finally:
        shutil.rmtree(temp_dir, ignore_errors=True)


# ---------------- TTS: Gemini English + regional path ----------------
def _extract_english_fallback(text: str) -> str:
    """
    Produce a clean, English-safe version of the narration as an absolute
    last-resort TTS input when all retry attempts and the chunked fallback
    have been exhausted.

    Transformations applied (all generic, no hardcoded words):
    - Strip Devanagari Unicode block so mixed-script content doesn't trigger
      the TTS content filter.
    - Remove markdown bold/italic markers.
    - Expand comma-grouped numbers (5,999 → 5999 and lakh-format 1,00,000 → 100000).
    - Replace ? with . to avoid interrogative prosody in borderline content.
    - Normalize whitespace.

    Returns the sanitized string, or the original if nothing meaningfully changed.
    """
    if not text:
        return text
    sanitized = text
    # Strip Devanagari
    sanitized = re.sub(r'[\u0900-\u097F]+', '', sanitized)
    # Strip markdown
    sanitized = re.sub(r'\*{1,3}([^*]+)\*{1,3}', r'\1', sanitized)
    # Expand comma-grouped numbers
    sanitized = re.sub(r'(?<!\d)(\d{1,3}),(\d{3})(?!\d)', r'\1\2', sanitized)
    sanitized = re.sub(r'(?<!\d)(\d{1,2}),(\d{2}),(\d{3})(?!\d)', r'\1\2\3', sanitized)
    # Replace ? with .
    sanitized = re.sub(r'\?', '.', sanitized)
    # Normalize whitespace
    sanitized = ' '.join(sanitized.split())
    return sanitized.strip()



def generate_english_tts_with_gemini(text: str, voice_name: str, output_path: str, tts_gender: str = None) -> str:
    if not genai or not types:
        raise ImportError("google-genai package is required for Gemini TTS. Install with: pip install google-genai")
    if not GEMINI_API_KEY:
        raise ValueError("GEMINI_API_KEY not found in environment variables")
    # Voice mapping: only voices supported by Gemini TTS API
    # Based on API test results: charon, fenrir, puck, aoede work
    # Removed: titan, lyra, aura, sol (not supported by API)
    # Additional supported voices: achernar, achird, algenib, algieba, alnilam, autonoe, 
    # callirrhoe, despina, enceladus, erinome, gacrux, iapetus, kore, laomedeia, leda, 
    # orus, pulcherrima, rasalgethi, sadachbia, sadaltager, schedar, sulafat, umbriel, 
    # vindemiatrix, zephyr, zubenelgenubi
    # Map user-friendly voice names to Gemini voice keys (male/female map to same prebuilt voice)
    voice_map = {
        "achernar":    {"male": "achernar",    "female": "achernar"},
        "achird":      {"male": "achird",      "female": "achird"},
        "algenib":     {"male": "algenib",     "female": "algenib"},
        "algieba":     {"male": "algieba",     "female": "algieba"},
        "alnilam":     {"male": "alnilam",     "female": "alnilam"},
        "aoede":       {"male": "aoede",       "female": "aoede"},
        "autonoe":     {"male": "autonoe",     "female": "autonoe"},
        "callirrhoe":  {"male": "callirrhoe",  "female": "callirrhoe"},
        "charon":      {"male": "charon",      "female": "charon"},
        "despina":     {"male": "despina",     "female": "despina"},
        "enceladus":   {"male": "enceladus",   "female": "enceladus"},
        "erinome":     {"male": "erinome",     "female": "erinome"},
        "fenrir":      {"male": "fenrir",      "female": "fenrir"},
        "gacrux":      {"male": "gacrux",      "female": "gacrux"},
        "iapetus":     {"male": "iapetus",     "female": "iapetus"},
        "kore":        {"male": "kore",        "female": "kore"},
        "laomedeia":   {"male": "laomedeia",   "female": "laomedeia"},
        "leda":        {"male": "leda",        "female": "leda"},
        "orus":        {"male": "orus",        "female": "orus"},
        "puck":        {"male": "puck",        "female": "puck"},
        "pulcherrima": {"male": "pulcherrima", "female": "pulcherrima"},
        "rasalgethi":  {"male": "rasalgethi",  "female": "rasalgethi"},
        "sadachbia":   {"male": "sadachbia",   "female": "sadachbia"},
        "sadaltager":  {"male": "sadaltager",  "female": "sadaltager"},
        "schedar":     {"male": "schedar",     "female": "schedar"},
        "sulafat":     {"male": "sulafat",     "female": "sulafat"},
        "umbriel":     {"male": "umbriel",     "female": "umbriel"},
        "vindemiatrix":{"male": "vindemiatrix", "female": "vindemiatrix"},
        "zephyr":      {"male": "zephyr",      "female": "zephyr"},
        "zubenelgenubi": {"male": "zubenelgenubi", "female": "zubenelgenubi"},
    }

    # Allowed voices list based on actual Gemini TTS API supported voices (lowercase)
    allowed_voices = {
        "achernar", "achird", "algenib", "algieba", "alnilam", "aoede", "autonoe", 
        "callirrhoe", "charon", "despina", "enceladus", "erinome", "fenrir", "gacrux", 
        "iapetus", "kore", "laomedeia", "leda", "orus", "puck", "pulcherrima", 
        "rasalgethi", "sadachbia", "sadaltager", "schedar", "sulafat", "umbriel", 
        "vindemiatrix", "zephyr", "zubenelgenubi"
    }

    requested = (voice_name or "").strip()
    # Default to 'female' if no gender provided to match wrapper defaults (audio_fn_from_string default is 'female')
    gender_key = 'male' if str(tts_gender or '').lower().startswith('m') else 'female'

    if not requested:
        gemini_voice = 'aoede'
    else:
        mapping = voice_map.get(requested.lower())
        if mapping:
            gemini_voice = mapping.get(gender_key)
        else:
            # If user supplied a direct Gemini voice name, accept it (lowercase)
            gemini_voice = requested.lower()

    # Validate against allowed voices (case-insensitive); if invalid, fallback to a safe default.
    if gemini_voice.lower() not in allowed_voices:
        print(f"[WARN] Requested Gemini voice '{gemini_voice}' not in allowed list; falling back to 'aoede'.")
        gemini_voice = 'aoede'
    else:
        # Ensure voice name is lowercase as API expects lowercase
        gemini_voice = gemini_voice.lower()
    
    try:
        client = genai.Client(api_key=GEMINI_API_KEY)

        logger.info(
            "[TTS] Generating audio | voice=%s | chars=%s",
            gemini_voice,
            len(text)
        )

        logger.info(
            "[TTS] Preview: %s",
            text[:250]
        )

        response = client.models.generate_content(
            model="gemini-2.5-flash-preview-tts",
            contents=text,
            config=types.GenerateContentConfig(
                response_modalities=["AUDIO"],
                speech_config=types.SpeechConfig(
                    voice_config=types.VoiceConfig(
                        prebuilt_voice_config=types.PrebuiltVoiceConfig(
                            voice_name=gemini_voice
                        )
                    )
                )
            )
        )

        audio_data = _extract_audio_bytes(response)

        if not audio_data:
            logger.error(
                "[TTS] Empty audio payload.\nResponse=%s",
                str(response)
            )
            # Distinguish content rejection (FinishReason.OTHER) from a transient failure.
            # When the model rejects content, retrying with the same text will never work —
            # the caller must transform the text before retrying.
            candidates = getattr(response, "candidates", []) or []
            finish_reasons = [
                str(getattr(c, "finish_reason", "")) for c in candidates
            ]
            if any("OTHER" in r for r in finish_reasons):
                raise _TTSContentRejected(
                    f"Gemini TTS rejected content (FinishReason.OTHER) "
                    f"for voice={voice_name}. Text must be modified before retry."
                )
            raise ValueError(
                "No audio data received from Gemini API."
            )

        temp_wav = output_path.replace(
            ".mp3",
            "_temp.wav"
        )

        with wave.open(temp_wav, "wb") as wav_file:
            wav_file.setnchannels(1)
            wav_file.setsampwidth(2)
            wav_file.setframerate(24000)
            wav_file.writeframes(audio_data)

        try:
            subprocess.run(
                [
                    "ffmpeg",
                    "-y",
                    "-i",
                    temp_wav,
                    "-acodec",
                    "libmp3lame",
                    "-ar",
                    "44100",
                    "-ac",
                    "2",
                    "-b:a",
                    "128k",
                    output_path
                ],
                check=True,
                capture_output=True,
                text=True
            )

            if os.path.exists(temp_wav):
                os.remove(temp_wav)

            return output_path

        except subprocess.CalledProcessError as e:

            logger.error(
                "[TTS] FFmpeg conversion failed: %s",
                e.stderr
            )

            with open(output_path, "wb") as f:
                f.write(audio_data)

            return output_path

    except Exception:
        traceback.print_exc()
        raise

# ---------------- Legacy batching (no-op kept) ----------------
def process_tts_in_batches(text, target_language, gender, final_speech_file_path, char_limit=300):
    return None

# ---------------- Higher-level generation helpers (Hinglish by default) ----------------
def generate_tts_static_words(input_text, target_language, gender, speech_file_path):
    text = indentify_named_entities(input_text)
    dictionary = process_text_into_dictionary(text)
    final_text = process_dictionary(dictionary, target_language)
    return process_tts_in_batches(final_text, target_language, gender, speech_file_path)

def generate_tts_indian(input_text, target_language, gender, speech_file_path):
    regional_text = translate_to_regional(target_language, input_text)
    if regional_text is None:
        return None
    return process_tts_in_batches(regional_text, target_language, gender, speech_file_path)

def generate_tts_hinglish(input_text, target_language, gender, speech_file_path):
    return generate_tts_static_words(input_text, target_language, gender, speech_file_path)

def generate_tts_hinglish(input_text, target_language, gender, speech_file_path):
    """Generate TTS for Hinglish (mix of Hindi/regional + English named entities)."""
    return generate_tts_static_words(input_text, target_language, gender, speech_file_path)


# ---------------- Main Serverless TTS ----------------
def _generate_speech(
    input_text,
    speech_file_path,
    tts_client,
    target_language,
    tts_gender,
    tts_voice_name,
    toggle_hinglish
):

    if not input_text or not input_text.strip():
        print(
            f"[WARN] Empty input text for {speech_file_path}"
        )
        return None

    clean_text = _clean_text_for_tts(input_text)

    target_language_lower = (
        target_language or "english"
    ).lower()

    use_hinglish = str(toggle_hinglish).lower() in {
        "true",
        "1",
        "yes",
        "y"
    }

    # English TTS path — pass multilingual text directly to Gemini TTS.
    # gemini-2.5-flash-preview-tts handles mixed-script content natively:
    # Devanagari, Romanized Hindi (Hinglish), and English all render with correct
    # pronunciation without any pre-transliteration on our side.
    if target_language_lower in {"english", "en"}:
        pass  # No text transformation; the model handles mixed scripts natively.

    # Regional TTS path
    else:

        try:
            if use_hinglish:

                text_with_entities = indentify_named_entities(
                    clean_text
                )

                text_dict = process_text_into_dictionary(
                    text_with_entities
                )

                clean_text = process_dictionary(
                    text_dict,
                    target_language
                )

            else:

                clean_text = translate_to_regional(
                    target_language,
                    clean_text
                )

        except Exception as e:

            logger.warning(
                "[TTS] Translation failed. Using original text. %s",
                str(e)
            )

    _content_rejected = False
    for attempt in range(1, 4):

        try:
            # After a content rejection, apply progressive text normalization
            # (e.g. remove comma grouping from numbers) so the retry uses
            # transformed text, not the exact same input that was just rejected.
            attempt_text = (
                _normalize_text_for_tts_retry(clean_text, attempt)
                if _content_rejected
                else clean_text
            )

            result = generate_english_tts_with_gemini(
                attempt_text,
                tts_voice_name or "aoede",
                speech_file_path,
                tts_gender
            )

            if result and os.path.exists(result):
                return result

        except _TTSContentRejected as e:
            _content_rejected = True
            print(f"[ERROR] TTS attempt {attempt} content rejected: {e}")

            if attempt == 3:
                # Last resort pass 1: generate audio sentence-by-sentence and stitch.
                # Avoids crashing the full pipeline when one beat's text
                # consistently triggers Gemini's content filter.
                try:
                    logger.info("[TTS] Falling back to chunked sentence-level TTS")
                    result = _generate_tts_in_chunks(
                        clean_text,
                        tts_voice_name or "aoede",
                        speech_file_path,
                        tts_gender
                    )
                    if result and os.path.exists(result):
                        return result
                except Exception as chunk_e:
                    print(f"[ERROR] TTS chunked fallback also failed: {chunk_e}")

                # Last resort pass 2: English-only sanitized version of the text.
                # Strips Devanagari, markdown, and problematic punctuation so the
                # TTS model receives a clean, unambiguous English utterance.
                # This guarantees the beat has audio and the pipeline continues.
                try:
                    logger.info("[TTS] Final fallback: English-only sanitized TTS for beat")
                    eng_fallback = _extract_english_fallback(clean_text)
                    if eng_fallback and eng_fallback.strip():
                        result = _generate_tts_in_chunks(
                            eng_fallback,
                            tts_voice_name or "aoede",
                            speech_file_path,
                            tts_gender
                        )
                        if result and os.path.exists(result):
                            logger.info("[TTS] English-only final fallback succeeded for beat")
                            return result
                except Exception as final_e:
                    logger.error("[TTS] English-only final fallback failed: %s", final_e)

        except Exception as e:

            print(
                f"[ERROR] TTS attempt {attempt} failed: {e}"
            )

        time.sleep(1.5)

    raise RuntimeError(
        f"[FAIL] TTS failed after all attempts for text: "
        f"{clean_text[:100]}..."
    )




# ---------------- Wrappers (Hinglish default) ----------------
def audio_fn_from_string(input_text, folder_path, target_language='english', tts_gender='female', tts_voice_name='Puck', toggle_hinglish=True, client=None, file_name_prefix="speech"):
    folder_path = folder_path if folder_path and os.access(folder_path, os.W_OK) else get_temp_dir()
    os.makedirs(folder_path, exist_ok=True)
    speech_file_path = os.path.join(folder_path, f"{file_name_prefix}.mp3")
    return _generate_speech(input_text, speech_file_path, client, target_language, tts_gender, tts_voice_name, toggle_hinglish)

def audio_fn(text_file_path, target_language='english', tts_gender='female', tts_voice_name='Puck', toggle_hinglish=True, client=None):
    with open(text_file_path, "r", encoding="utf-8") as f:
        input_text = f.read()
    folder_path = os.path.dirname(text_file_path) or get_temp_dir()
    base_filename = os.path.splitext(os.path.basename(text_file_path))[0]
    speech_file_path = os.path.join(folder_path, f"{base_filename}.mp3")
    return _generate_speech(input_text, speech_file_path, client, target_language, tts_gender, tts_voice_name, toggle_hinglish)

# ---------------- pun_ wrappers (consistent behavior) ----------------
def pun_process_dictionary(input_dict: dict, target_language: str) -> str:
    outputs = []
    for key, value in input_dict.items():
        if key.startswith('hindi'):
            outputs.append(pun_translate_to_regional(text=value, target_language=target_language))
        elif key.startswith('english'):
            outputs.append(value.strip())
    return ' '.join(outputs)

def pun_audio_fn_from_string(text_input: str, client = None,
    target_language: str = 'english', tts_gender: str = 'female', tts_voice_name: str = 'nova',
    toggle_hinglish: bool = True, text_source_type: Literal['file', 'string']='file',
    file_name_prefix: str='speech', output_dir: str | None = None) -> (str | None):
    print('==============================================')
    try:
        if text_source_type == 'file':
            with open(text_input, 'r', encoding='utf-8') as file:
                input_text = file.read()
            if not input_text.strip():
                print(f'Warning: Input text file is empty, skipping TTS for {text_input}')
                return None
            directory = os.path.dirname(text_input)
            base_filename = os.path.splitext(os.path.basename(text_input))[0]
            speech_file_path = os.path.join(directory, f'{base_filename}.mp3')
        elif text_source_type == 'string':
            input_text = text_input
            if not output_dir:
                raise ValueError("output_dir must be provided when text_source_type is 'string'")
            os.makedirs(output_dir, exist_ok=True)
            speech_file_path = os.path.join(output_dir, f'{file_name_prefix}.mp3')
        else:
            print(f"Invalid text_source_type: {text_source_type}")
            return None
        return _generate_speech(input_text, speech_file_path, None, target_language, tts_gender, tts_voice_name, toggle_hinglish)
    except FileNotFoundError:
        print(f'Error: Text file not found at {text_input}')
        return None
    except Exception as e:
        print(f'Error generating audio: {e}')
        return None

# ---------------- Utilities: PowerPoint audio inspection & trimming ----------------
def pun_check_audio_files(pptx_path, output_dir):
    print(f'Checking audio files in: {pptx_path}')
    os.makedirs(output_dir, exist_ok=True)
    audio_mapping = {}
    with zipfile.ZipFile(pptx_path, 'r') as z:
        media_files = [f for f in z.namelist() if f.startswith('ppt/media/')]
        audio_files = [f for f in media_files if f.endswith('.mp3')]
        slide_files = [f for f in z.namelist() if f.startswith('ppt/slides/slide') and f.endswith('.xml')]
        slide_files.sort(key=lambda x: int(re.search('slide(\\d+)\\.xml$', x).group(1)))
        for audio_file in audio_files:
            match = re.search('media(\\d+)\\.mp3$', audio_file)
            if match:
                slide_num = int(match.group(1))
                if slide_num <= len(slide_files):
                    audio_mapping[slide_num] = audio_file
                    output_path = os.path.join(output_dir, f'slide_{slide_num}_{os.path.basename(audio_file)}')
                    with open(output_path, 'wb') as f:
                        f.write(z.read(audio_file))
    return audio_mapping

def pun_get_audio_duration(file_path: str) -> float:
    if not os.path.exists(file_path):
        return 0.0
    try:
        audio = AudioSegment.from_file(file_path)
        return len(audio) / 1000.0
    except Exception:
        return 0.0

def trim_audio_to_max_duration(input_path: str, max_seconds: float) -> str:
    if not os.path.exists(input_path):
        return input_path
    try:
        duration = pun_get_audio_duration(input_path)
        if duration <= max_seconds:
            return input_path
    except Exception:
        pass
    temp_out = f"{input_path}.trim.mp3"
    ffmpeg_bin = shutil.which("ffmpeg") or "ffmpeg"
    cmd = [ffmpeg_bin, "-y", "-i", input_path, "-ss", "0", "-t", str(max_seconds), "-c:a", "libmp3lame", "-b:a", "128k", temp_out]
    try:
        subprocess.run(cmd, check=True, capture_output=True, text=True)
        os.replace(temp_out, input_path)
        return input_path
    except Exception:
        if os.path.exists(temp_out):
            os.remove(temp_out)
        return input_path

def pad_audio_to_duration(input_path: str, target_seconds: float) -> str:
    """Ensure `input_path` audio is at least `target_seconds` long.
    If shorter, append silence to reach target duration. Returns the path to the (possibly modified) file.
    Operates in-place (creates a temporary file and replaces original).
    """
    if not os.path.exists(input_path):
        return input_path
    try:
        audio = AudioSegment.from_file(input_path)
        current = len(audio) / 1000.0
    except Exception:
        return input_path
    if current >= float(target_seconds):
        return input_path
    silence_ms = int((float(target_seconds) - current) * 1000)
    padding = AudioSegment.silent(duration=silence_ms)
    new_audio = audio + padding
    temp_out = f"{input_path}.pad.mp3"
    try:
        new_audio.export(temp_out, format='mp3', bitrate='128k')
        os.replace(temp_out, input_path)
        return input_path
    except Exception:
        if os.path.exists(temp_out):
            os.remove(temp_out)
        return input_path

# ---------------- Media combine ----------------
def get_media_duration(file_path: str) -> float:
    try:
        cmd = ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", file_path]
        result = subprocess.run(cmd, check=True, capture_output=True, text=True)
        return float(result.stdout.strip())
    except Exception:
        return 0.0

def combine_audio_with_video(video_path: str, audio_path: str, output_path: str) -> bool:
    try:
        if not os.path.exists(video_path) or not os.path.exists(audio_path):
            return False
        video_duration = get_media_duration(video_path)
        audio_duration = get_media_duration(audio_path)
        ffmpeg_bin = shutil.which("ffmpeg") or "ffmpeg"
        if audio_duration > video_duration:
            pad = audio_duration - video_duration
            command = [
                ffmpeg_bin, "-y",
                "-i", video_path,
                "-i", audio_path,
                "-filter_complex", f"[0:v]tpad=stop_mode=clone:stop_duration={pad}[v]",
                "-map", "[v]",
                "-map", "1:a:0",
                "-c:v", "libx264",
                "-preset", "medium",
                "-crf", "23",
                "-c:a", "aac",
                "-shortest",
                output_path
            ]
        else:
            command = [
                ffmpeg_bin, "-y",
                "-i", video_path,
                "-i", audio_path,
                "-map", "0:v:0",
                "-map", "1:a:0",
                "-c:v", "copy",
                "-c:a", "aac",
                "-shortest",
                output_path
            ]
        subprocess.run(command, check=True, capture_output=True, text=True)
        return os.path.exists(output_path) and os.path.getsize(output_path) > 0
    except subprocess.CalledProcessError as e:
        print(f"FFmpeg failed: {e.stderr}")
        return False
    except Exception as e:
        print(f"Unexpected combine error: {e}")
        return False

# ---------------- Segment-level audio assembly (pun_) ----------------
def pun_generate_segment_audios(code_segments: list, temp_audio_dir: str,
    client = None, txt_src_typ: Literal['file', 'string']='file',
    target_language: str='english', tts_gender: str='female',
    tts_voice_name: str='nova', toggle_hinglish: bool=False, delay_ms: int=1000, default_chars_per_second: int=25, extra_padding_s: float=0.5) -> tuple[str | None, float]:
    os.makedirs(temp_audio_dir, exist_ok=True)
    final_audio_clips_for_master = []
    total_duration = 0.0
    for i, segment in enumerate(code_segments):
        if segment.explanation:
            try:
                generated_audio_path = pun_audio_fn_from_string(
                    text_input=pun_remove_special_characters(segment.explanation),
                    client=client,
                    target_language=target_language,
                    tts_gender=tts_gender,
                    tts_voice_name=tts_voice_name,
                    toggle_hinglish=toggle_hinglish,
                    file_name_prefix=f'segment_{i:02d}',
                    text_source_type=txt_src_typ
                )
                if generated_audio_path:
                    segment.audio_path = generated_audio_path
                    segment.audio_duration = pun_get_audio_duration(generated_audio_path)
                else:
                    segment.audio_duration = 0.0
            except Exception:
                segment.audio_duration = 0.0
        else:
            segment.audio_duration = 0.0
        typing_duration = len(segment.code_snippet) / default_chars_per_second if default_chars_per_second > 0 else 0.1
        segment_visual_duration = max(segment.audio_duration, typing_duration) + extra_padding_s
        segment_visual_duration_ms = int(segment_visual_duration * 1000)
        segment_audio_block = AudioSegment.silent(duration=segment_visual_duration_ms)
        if segment.audio_duration > 0 and os.path.exists(segment.audio_path):
            try:
                explanation_audio = AudioSegment.from_file(segment.audio_path)
                segment_audio_block = segment_audio_block.overlay(explanation_audio, position=0)
            except Exception:
                pass
        final_audio_clips_for_master.append(segment_audio_block)
        total_duration += segment_visual_duration
        if i < len(code_segments) - 1:
            final_audio_clips_for_master.append(AudioSegment.silent(duration=delay_ms))
            total_duration += delay_ms / 1000.0
    master_audio_path = os.path.join(temp_audio_dir, 'master_explanation.mp3')
    if final_audio_clips_for_master:
        try:
            combined_audio = sum(final_audio_clips_for_master)
            combined_audio.export(master_audio_path, format='mp3')
            return master_audio_path, total_duration
        except Exception:
            return None, 0.0
    else:
        return None, 0.0