UmbraAudio/umbraaudio/mapping/emotion.py

from __future__ import annotations

import json
import logging
from typing import Any

from umbraaudio.voices.lmstudio_gate import queued_chat_completion

logger = logging.getLogger(__name__)

EMOTION_PROMPT = """You are an audiobook director analyzing dialogue and narration for TTS voice direction.

For EACH utterance below, determine the emotional tone and speaking style.

## Emotion Categories
Use these primary emotions (combine if needed): neutral, angry, sad, happy, fearful, surprised, whisper, pleading, sarcastic, authoritative, thoughtful, excited, anxious, tired, romantic, disgusted, amused, nostalgic

## Pace
slow, moderate, fast

## Output Format
Return a JSON array with one object per utterance, in the same order:
[
  {{"idx": 0, "emotion": "angry", "pace": "fast", "emphasis": ["word1", "word2"], "instruct": "Speak with rising anger, voice tightening"}},
  {{"idx": 1, "emotion": "neutral", "pace": "moderate", "emphasis": [], "instruct": "Read steadily and evenly"}}
]

## Rules
- For NARRATOR text: emotion="neutral", pace="moderate", instruct="Read steadily and evenly. Low emphasis, consistent tone."
- For DIALOGUE: infer from context clues (exclamation marks, word choice, scene tension)
- QUESTIONS: if the utterance ends with a question mark, it is a question — include rising intonation at the end in the instruct (e.g. "End with rising question intonation")
- Keep instruct strings under 80 words
- emphasis: max 3 key words per utterance
- Return ONLY the JSON array

## Emotion Disambiguation (CRITICAL)
DO NOT default to "tense" for everything. Use these guidelines:

- **tense**: ONLY for active danger, confrontation, or high-stakes moments where characters are on edge (e.g., under attack, being chased, in a standoff). NOT for routine planning or conversation.
- **neutral**: Routine conversation, information exchange, casual discussion, planning without urgency
- **authoritative**: Giving orders, making decisions, issuing commands, professional briefing
- **thoughtful**: Reflecting, considering options, analyzing a situation, thinking aloud
- **urgent**: Time-critical situations, warnings, emergency communication
- **anxious**: Worry about uncertain outcomes, nervous anticipation, unease
- **excited**: Positive anticipation, discovery, good news, celebration
- **amused**: Humor, light-hearted moments, wit, playful banter
- **sarcastic**: Ironic commentary, dry wit, mocking
- **sad**: Loss, grief, disappointment, melancholy
- **happy**: Joy, satisfaction, relief, warmth
- **angry**: Frustration, rage, indignation, confrontation
- **fearful**: Terror, dread, being scared, panic
- **pleading**: Begging, desperate requests, imploring
- **romantic**: Intimacy, tenderness, affection, love
- **whispered**: Secretive, private, intimate communication
- **tired**: Exhaustion, weariness, drained delivery
- **nostalgic**: Reminiscing, wistful reflection, looking back
- **disgusted**: Revulsion, rejection, distaste
- **surprised**: Shock, astonishment, unexpected discovery

## Scene Context Matters
- Battle/action scenes: angry, urgent, fearful, tense (appropriate here)
- Planning/briefing scenes: authoritative, thoughtful, neutral
- Personal/emotional scenes: sad, happy, romantic, nostalgic
- Humor/banter scenes: amused, sarcastic, happy
- Routine operations: neutral, authoritative, thoughtful
- Post-action recovery: tired, thoughtful, sad, relieved

## Character Awareness
- A ship AI giving reports: authoritative or neutral (not tense)
- A princess making decisions: authoritative or thoughtful (not tense)
- Casual conversation between crew: neutral or amused (not tense)
- A character alone reflecting: thoughtful or nostalgic (not tense)

## Utterances
{utterances}"""


class EmotionMapper:
    def __init__(self, llm_client: Any = None, model: str = "") -> None:
        self.llm = llm_client
        self.model = model

    def _call_llm(self, prompt: str) -> list[dict[str, Any]]:
        if not self.llm:
            return []

        max_attempts = 3
        for attempt in range(max_attempts):
            try:
                response = queued_chat_completion(
                    self.llm,
                    model=self.model,
                    messages=[{"role": "user", "content": prompt}],
                    temperature=0.2,
                    max_tokens=8192,
                )
                break
            except Exception as e:
                logger.warning(
                    "Emotion LLM call failed (attempt %d/%d): %s",
                    attempt + 1,
                    max_attempts,
                    e,
                )
                if attempt < max_attempts - 1:
                    import time

                    time.sleep(3 * (attempt + 1))
                else:
                    return []

        content = response.choices[0].message.content or ""
        content = content.strip()

        return self._parse_response(content)

    def _parse_response(self, content: str) -> list[dict[str, Any]]:
        if not content:
            logger.warning("Empty emotion LLM response")
            return []

        stripped = content.strip()

        if stripped.startswith("```"):
            first_newline = stripped.find("\n")
            if first_newline > 0:
                stripped = stripped[first_newline + 1 :]
            if stripped.endswith("```"):
                stripped = stripped[:-3].strip()

        result = self._try_json_parse(stripped)
        if result is not None:
            return result

        json_start = stripped.find("[")
        json_end = stripped.rfind("]") + 1
        if json_start >= 0 and json_end > json_start:
            candidate = stripped[json_start:json_end]
            result = self._try_json_parse(candidate)
            if result is not None:
                return result

        logger.warning("Failed to parse emotion LLM response")
        return []

    def _try_json_parse(self, text: str) -> list[dict[str, Any]] | None:
        try:
            result = json.loads(text)
            if isinstance(result, list):
                return [dict(x) for x in result if isinstance(x, dict)]
            return None
        except json.JSONDecodeError:
            pass

        fixed = text
        for _ in range(10):
            try:
                result = json.loads(fixed)
                if isinstance(result, list):
                    return [dict(x) for x in result if isinstance(x, dict)]
                return None
            except json.JSONDecodeError:
                last_brace = fixed.rfind("}")
                last_bracket = fixed.rfind("]")
                last_comma = fixed.rfind(",")
                cut = max(last_brace, last_bracket, last_comma)
                if cut > 0:
                    fixed = fixed[:cut]
                    if fixed.endswith(","):
                        fixed = fixed[:-1]
                else:
                    break

        return None

    def analyze_utterances(
        self,
        utterances: list[dict[str, Any]],
        chapter_context: str = "",
    ) -> list[dict[str, Any]]:
        if not utterances:
            return []

        if not self.llm:
            return [self._default_result(u) for u in utterances]

        batch_size = 30
        all_results: list[dict[str, Any]] = []

        for start in range(0, len(utterances), batch_size):
            batch = utterances[start : start + batch_size]

            lines: list[str] = []
            for i, u in enumerate(batch):
                utt_type = u.get("type", "narration")
                speaker = u.get("character_name", "narrator")
                text = u.get("text", "")[:300]
                lines.append(f"[{i}] ({utt_type}/{speaker}) {text}")

            utterances_text = "\n".join(lines)
            context_snippet = (
                f"Chapter context: {chapter_context[:500]}\n\n" if chapter_context else ""
            )

            prompt = EMOTION_PROMPT.format(
                utterances=context_snippet + utterances_text
            )

            results = self._call_llm(prompt)

            if results and len(results) == len(batch):
                for r in results:
                    all_results.append({
                        "emotion": r.get("emotion", "neutral"),
                        "pace": r.get("pace", "moderate"),
                        "emphasis_words": r.get("emphasis", []),
                        "instruct": r.get("instruct", ""),
                        "pause_before_ms": 500,
                        "pause_after_ms": 500,
                    })
            else:
                for u in batch:
                    all_results.append(self._default_result(u))

        return all_results

    @staticmethod
    def _default_result(utt: dict[str, Any]) -> dict[str, Any]:
        text = (utt.get("text") or "").strip()
        is_question = text.endswith("?") or text.endswith("?\"")
        if utt.get("type") == "narration":
            return {
                "emotion": "neutral",
                "pace": "moderate",
                "emphasis_words": [],
                "instruct": "Read steadily and evenly. Low emphasis, consistent tone.",
                "pause_before_ms": 500,
                "pause_after_ms": 500,
            }
        if is_question:
            return {
                "emotion": "neutral",
                "pace": "moderate",
                "emphasis_words": [],
                "instruct": "Speak naturally with rising question intonation at the end.",
                "pause_before_ms": 500,
                "pause_after_ms": 500,
            }
        return {
            "emotion": "neutral",
            "pace": "moderate",
            "emphasis_words": [],
            "instruct": "Speak naturally.",
            "pause_before_ms": 500,
            "pause_after_ms": 500,
        }

    def batch_generate(
        self, utterances: list[dict[str, Any]], chapter_context: str = ""
    ) -> list[dict[str, Any]]:
        return self.analyze_utterances(utterances, chapter_context)