agents/tools.py

"""
agents/tools.py — LLM tools + TTS + Personas
"""

from __future__ import annotations
import os
import json
import re
import hashlib
import random
from pathlib import Path
from data.db import SEG_DIR
from openai import OpenAI

# ── LLM client (HuggingFace router → Qwen) ─────────────────────────────────

def _get_llm_client() -> OpenAI | None:
    token = os.getenv("HF_TOKEN")
    if not token:
        return None
    return OpenAI(
        base_url="https://router.huggingface.co/v1",
        api_key=token,
    )

LLM_MODEL_Nemo ="nvidia/Llama-3.1-Nemotron-8B-UltraLong-1M-Instruct:featherless-ai"
LLM_MODEL = "Qwen/Qwen2.5-7B-Instruct:together"


def _llm_chat(messages: list[dict], system: str = "") -> str:
    """Call the LLM. Falls back to mock if no HF_TOKEN."""
    client = _get_llm_client()
    if client is None:
        return _mock_llm(messages[-1]["content"] if messages else "")

    full_messages = []
    if system:
        full_messages.append({"role": "system", "content": system})
    full_messages.extend(messages)

    try:
        resp = client.chat.completions.create(
            model=LLM_MODEL,
            messages=full_messages,
            max_tokens=800,
            temperature=0.85,
        )
        return resp.choices[0].message.content or ""
    except Exception as e:
        print(f"[LLM error] {e}")
        return _mock_llm(messages[-1]["content"] if messages else "")


def _mock_llm(text: str) -> str:
    """Returns a structurally valid mock response for dev/testing."""
    score = round(random.uniform(4.5, 9.5), 1)
    reactions = [
        "I felt this in my entire body.",
        "This is the most relatable thing I've read today.",
        "Whoever wrote this, I see you.",
        "Corporate dystopia summed up perfectly.",
        "I laughed, then I cried, then I refreshed my LinkedIn.",
    ]
    return json.dumps({
        "is_brag": False,
        "confidence": 0.1,
        "roast": "",
        "scores": {
            "hilarious": round(random.uniform(3, 9), 1),
            "tragic": round(random.uniform(2, 8), 1),
            "unhinged": round(random.uniform(3, 9), 1),
            "awkward": round(random.uniform(2, 7), 1),
            "chaotic": round(random.uniform(3, 9), 1),
        },
        "top_category": "hilarious",
        "top_score": score,
        "reaction": random.choice(reactions),
    })


# ── Personas ────────────────────────────────────────────────────────────────

PERSONAS = {
    "hilarious": {
        "name": "Dave",
        "show": "The Weekly Wheeze",
        "desc": "42, raspy, laughs mid-sentence",
        "voice": "Mr_Quagmire",
        "style": "comedic stand-up energy, warm and relatable",
        "sfx": ["drumroll.wav", "crowd_react.wav"],
    },
    "tragic": {
        "name": "Elena",
        "show": "Corporate Tears",
        "desc": "31, soft, melancholic sighs",
        "voice": "Sarah",
        "style": "melancholic NPR storyteller, gentle and empathetic",
        "sfx": ["violin_sting.wav"],
    },
    "unhinged": {
        "name": "Marcus",
        "show": "Officially Unhinged",
        "desc": "55, baritone, barely-contained rage",
        "voice": "Cho",
        "style": "barely restrained fury, dark comedy, gravelly intensity",
        "sfx": ["thunder.wav", "glass_break.wav", "dramatic_sting.wav"],
    },
    "awkward": {
        "name": "Priya",
        "show": "Please Stop Talking",
        "desc": "27, fast talker, nervous giggle",
        "voice": "nicole",
        "style": "rapid nervous energy, second-hand cringe expert",
        "sfx": [],
    },
    "chaotic": {
        "name": "Rex",
        "show": "Total System Failure",
        "desc": "???, manic, breaks 4th wall",
        "voice": "max",
        "style": "completely unhinged, 4th-wall breaking, reality-questioning",
        "sfx": ["static.wav", "dial_up.wav", "system_error.wav", "vhs_rewind.wav"],
    },
}

# ── Tool 1: Detect Bragging ─────────────────────────────────────────────────

BRAG_SYSTEM = """You are a LinkedIn Brag Detector. Your job is to identify whether a workplace post 
is a humblebrag, achievement flex, or self-promotional content.

Flag as brag if it contains: achievement announcements, name-drops, revenue flexing, 
"humbled and honored" energy, "I've been selected…" constructions, follower-count mentions,
promotion announcements, award acceptances, or obvious self-congratulation.

Respond ONLY with valid JSON, no markdown, no explanation:
{"is_brag": true/false, "confidence": 0.0-1.0, "roast": "short witty roast if it's a brag, else empty string"}"""


def detect_bragging(text: str) -> dict:
    raw = _llm_chat(
        [{"role": "user", "content": f"Is this a brag?\n\n{text}"}],
        system=BRAG_SYSTEM,
    )
    try:
        clean = re.sub(r"```(?:json)?|```", "", raw).strip()
        return json.loads(clean)
    except Exception:
        return {"is_brag": False, "confidence": 0.0, "roast": ""}


# ── Tool 2: Classify Emotion ────────────────────────────────────────────────

EMOTION_SYSTEM = """You are an emotion classifier for workplace confessions. 
Score this post 0–10 across all 5 categories. Be generous — real stories deserve high scores.

Categories:
- hilarious: laugh-out-loud absurd, relatable workplace comedy
- tragic: real human pain dressed in corporate language
- unhinged: chaos energy, something broke inside this person
- awkward: socially catastrophic, second-hand cringe
- chaotic: entropy incarnate, nothing makes sense

Also write a one-line coworker reaction (what a colleague would DM you after reading this).

Respond ONLY with valid JSON, no markdown:
{"scores": {"hilarious": 0-10, "tragic": 0-10, "unhinged": 0-10, "awkward": 0-10, "chaotic": 0-10}, 
 "top_category": "category_name", "top_score": 0-10, "reaction": "one-line coworker DM"}"""


def classify_emotion(text: str) -> dict:
    raw = _llm_chat(
        [{"role": "user", "content": f"Score this workplace confession:\n\n{text}"}],
        system=EMOTION_SYSTEM,
    )
    try:
        clean = re.sub(r"```(?:json)?|```", "", raw).strip()
        return json.loads(clean)
    except Exception:
        return {
            "scores": {"hilarious": 5, "tragic": 5, "unhinged": 5, "awkward": 5, "chaotic": 5},
            "top_category": "chaotic",
            "top_score": 5.0,
            "reaction": "I don't know what to say.",
        }


# ── Tool 3: Build Podcast Script ─────────────────────────────────────────────

SCRIPT_SYSTEM = """You are {name}, host of "{show}" ({desc}).
Your style: {style}.

Write a podcast episode script based on the provided workplace confessions.
Structure it with these exact delimiters (no extra text between them):

---INTRO---
[Your opening monologue, 3-4 sentences, hook the listener]

---POST_1---
[Your commentary on the first post, 2-3 sentences]

---POST_2---
[Commentary on second post]

[Continue for all posts]

---OUTRO---
[Sign-off, 2-3 sentences]

Use emotion tags naturally: <laugh/> <sigh/> <gasp/> <cry/> <pause ms="500"/>
You can also use <whisper>text</whisper> and <shout>text</shout>.
Keep each section conversational and in character. Reference the post content directly."""


def build_podcast_script(category: str, posts: list[dict]) -> str:
    persona = PERSONAS[category]
    system = SCRIPT_SYSTEM.format(
        name=persona["name"],
        show=persona["show"],
        desc=persona["desc"],
        style=persona["style"],
    )
    posts_text = "\n\n".join(
        [f"POST {i+1}:\n{p['text']}" for i, p in enumerate(posts)]
    )
    return _llm_chat(
        [{"role": "user", "content": f"Here are the top workplace confessions for today's episode:\n\n{posts_text}"}],
        system=system,
    )


# ── TTS helpers ──────────────────────────────────────────────────────────────

SFX_TAG_MAP = {
    "<laugh/>": "ha,",
    "<sigh/>": "...",
    "<gasp/>": "oh!",
    "<cry/>": "...",
}

SFX_PATTERN = re.compile(r'<(laugh|sigh|gasp|cry)/>')
PAUSE_PATTERN = re.compile(r'<pause[^/]*/>')
WHISPER_PATTERN = re.compile(r'<whisper>(.*?)</whisper>', re.DOTALL)
SHOUT_PATTERN = re.compile(r'<shout>(.*?)</shout>', re.DOTALL)


def strip_emotion_tags(text: str) -> str:
    """Convert emotion tags to natural spoken equivalents for TTS."""
    text = text.replace("<laugh/>", "ha,")
    text = text.replace("<sigh/>", "...")
    text = text.replace("<gasp/>", "oh!")
    text = text.replace("<cry/>", "...")
    text = PAUSE_PATTERN.sub(" ", text)
    text = WHISPER_PATTERN.sub(r"\1", text)
    text = SHOUT_PATTERN.sub(r"\1", text)
    return text.strip()


def parse_script(script: str) -> dict[str, str]:
    """Parse ---SECTION--- delimited script into dict."""
    sections = {}
    current_key = None
    current_lines = []

    for line in script.split("\n"):
        line = line.strip()
        if line.startswith("---") and line.endswith("---"):
            if current_key and current_lines:
                sections[current_key] = "\n".join(current_lines).strip()
            current_key = line.strip("-").strip()
            current_lines = []
        elif current_key:
            current_lines.append(line)

    if current_key and current_lines:
        sections[current_key] = "\n".join(current_lines).strip()

    return sections


# ─────────────────────────────────────────────────────────────────────────────
# TTS BACKEND SELECTION
#
# Set TTS_BACKEND env var to force a backend:
#   "hf"          → HuggingFace InferenceClient  (hexgrad/Kokoro-82M via API)
#   "chatterbox"  → Chatterbox Turbo              (local, needs GPU for speed)
#   "voxcpm2"     → VoxCPM2                       (local, voice design mode)
#   Fallback "elevenlabs"  → ElevenLabs API                (needs ELEVENLABS_API_KEY)
#
# If TTS_BACKEND is not set, backends are tried in this order:
#   hf → voxcpm2 → chatterbox → elevenlabs
# ─────────────────────────────────────────────────────────────────────────────

# Per-persona voice descriptions used by VoxCPM2 and Chatterbox (no voice files needed)
VOICE_DESCRIPTIONS = {
    "hilarious": "A raspy middle-aged man, 42 years old, warm comedic energy, laughs easily",
    "tragic":    "A soft young woman, 31 years old, gentle and melancholic, empathetic storyteller",
    "unhinged":  "A deep baritone man, 55 years old, barely contained rage, gravelly intensity",
    "awkward":   "A fast-talking young woman, 27 years old, nervous giggle, second-hand cringe energy",
    "chaotic":   "A manic voice of unknown age, breaks the fourth wall, completely unhinged",
}

# HF InferenceClient voice IDs (Kokoro voices, used with hf backend)
HF_VOICE_IDS = {
    "hilarious": "Mr_Quagmire",
    "tragic":    "Sarah",
    "unhinged":  "Cho",
    "awkward":   "nicole",
    "chaotic":   "max",
}


# ── Option 1: HuggingFace InferenceClient → hexgrad/Kokoro-82M ───────────────

def synthesize_with_hf(text: str, voice: str = "am_fenrir") -> bytes | None:
    """
    TTS via HuggingFace InferenceClient using hexgrad/Kokoro-82M.
    Requires HF_TOKEN env var. Returns MP3 bytes or None.
    voice: a Kokoro voice ID like 'am_fenrir', 'af_sarah', etc.
    """
    token = os.getenv("HF_TOKEN")
    if not token:
        print("[HF TTS] No HF_TOKEN set, skipping.")
        return None
    try:
        from huggingface_hub import InferenceClient
        client = InferenceClient(provider="auto", api_key=token)
        # Returns bytes (MP3)
        audio: bytes = client.text_to_speech(
            text,
            model="hexgrad/Kokoro-82M",
            # Pass voice as extra_body since the SDK may not expose it directly
            extra_body={"voice": voice},
        )
        print(f"[HF TTS] OK — {len(audio)} bytes")
        return audio
    except Exception as e:
        print(f"[HF TTS error] {e}")
        return None


# ── Option 2: Chatterbox Turbo (local) ───────────────────────────────────────

def synthesize_with_chatterbox(
    text: str,
    voice_description: str = "A warm conversational voice",
    reference_wav_path: str | None = None,
    temperature: float = 0.8,
) -> bytes | None:
    """
    TTS via Chatterbox Turbo (local model).
    - If reference_wav_path is provided: clones that voice.
    - Otherwise: uses voice_description as a style guide (model ignores it
      but keeps output neutral; swap in a reference wav per-persona for best results).
    Returns WAV bytes or None.
    Install: pip install chatterbox-tts
    """
    try:
        import io
        import numpy as np
        import soundfile as sf
        import torch
        from chatterbox.tts_turbo import ChatterboxTurboTTS

        device = "cuda" if torch.cuda.is_available() else "cpu"

        # Lazy-load model (heavy, cache it)
        if not hasattr(synthesize_with_chatterbox, "_model"):
            print("[Chatterbox] Loading model...")
            synthesize_with_chatterbox._model = ChatterboxTurboTTS.from_pretrained(device)
            print("[Chatterbox] Model ready.")

        model = synthesize_with_chatterbox._model

        wav = model.generate(
            text,
            audio_prompt_path=reference_wav_path,  # None = no cloning
            temperature=temperature,
            min_p=0.05,
            top_p=0.9,
            top_k=50,
            repetition_penalty=1.1,
            norm_loudness=True,
        )

        # wav is a tensor shape (1, samples) or (samples,)
        audio_np = wav.squeeze(0).numpy() if hasattr(wav, "numpy") else np.array(wav)
        buf = io.BytesIO()
        sf.write(buf, audio_np, model.sr, format="WAV")
        print(f"[Chatterbox] OK — {len(buf.getvalue())} bytes")
        return buf.getvalue()

    except Exception as e:
        print(f"[Chatterbox error] {e}")
        return None


# ── Option 3: VoxCPM2 (local, voice design mode) ─────────────────────────────

def synthesize_with_voxcpm2(
    text: str,
    voice_description: str = "A warm conversational voice",
    reference_wav_path: str | None = None,
    cfg_value: float = 2.0,
    inference_timesteps: int = 10,
) -> bytes | None:
    """
    TTS via VoxCPM2.
    - If reference_wav_path provided: voice cloning mode.
    - Otherwise: voice design mode — prepend (voice_description) to text.
    Returns WAV bytes or None.
    Install: pip install voxcpm
    """
    try:
        import io
        import soundfile as sf
        from voxcpm import VoxCPM

        # Lazy-load model
        if not hasattr(synthesize_with_voxcpm2, "_model"):
            print("[VoxCPM2] Loading model...")
            synthesize_with_voxcpm2._model = VoxCPM.from_pretrained(
                "openbmb/VoxCPM2", load_denoiser=False
            )
            print("[VoxCPM2] Model ready.")

        model = synthesize_with_voxcpm2._model

        if reference_wav_path:
            # Controllable voice cloning with optional style prefix
            styled_text = f"({voice_description}){text}" if voice_description else text
            wav = model.generate(
                text=styled_text,
                reference_wav_path=reference_wav_path,
                cfg_value=cfg_value,
                inference_timesteps=inference_timesteps,
            )
        else:
            # Pure voice design — description drives the voice character
            styled_text = f"({voice_description}){text}"
            wav = model.generate(
                text=styled_text,
                cfg_value=cfg_value,
                inference_timesteps=inference_timesteps,
            )

        buf = io.BytesIO()
        sf.write(buf, wav, model.tts_model.sample_rate, format="WAV")
        print(f"[VoxCPM2] OK — {len(buf.getvalue())} bytes")
        return buf.getvalue()

    except Exception as e:
        print(f"[VoxCPM2 error] {e}")
        return None


# ── Option 4: ElevenLabs API (cloud fallback) ─────────────────────────────────

def synthesize_with_elevenlabs(text: str, voice_id: str = "21m00Tcm4TlvDq8ikWAM") -> bytes | None:
    """
    TTS via ElevenLabs API. Requires ELEVENLABS_API_KEY env var.
    Returns MP3 bytes or None.
    Default voice_id = Rachel (neutral female). Swap per-persona as needed.
    """
    api_key = os.getenv("ELEVENLABS_API_KEY")
    if not api_key:
        print("[ElevenLabs] No ELEVENLABS_API_KEY set, skipping.")
        return None
    try:
        import requests
        url = f"https://api.elevenlabs.io/v1/text-to-speech/{voice_id}"
        headers = {"xi-api-key": api_key, "Content-Type": "application/json"}
        payload = {
            "text": text,
            "model_id": "eleven_monolingual_v1",
            "voice_settings": {"stability": 0.5, "similarity_boost": 0.75},
        }
        r = requests.post(url, json=payload, headers=headers, timeout=30)
        r.raise_for_status()
        print(f"[ElevenLabs] OK — {len(r.content)} bytes")
        return r.content
    except Exception as e:
        print(f"[ElevenLabs error] {e}")
        return None


# ── Main entry point: try backends in configured order ────────────────────────

def synthesize_segment(
    text: str,
    voice: str = "am_fenrir",
    speed: float = 1.0,
    category: str = "hilarious",
) -> bytes | None:
    """
    Synthesize a TTS segment using the configured backend (TTS_BACKEND env var).
    Falls back through all backends if none is forced.

    voice        — Kokoro voice ID, used by the HF backend
    category     — persona key, used to pick voice description for VoxCPM2/Chatterbox
    speed        — currently used by HF/Kokoro only
    """
    backend = os.getenv("TTS_BACKEND", "").lower().strip()
    voice_desc = VOICE_DESCRIPTIONS.get(category, "A clear conversational voice")

    # Forced backend
    if backend == "hf":
        return synthesize_with_hf(text, voice=voice)
    if backend == "chatterbox":
        return synthesize_with_chatterbox(text, voice_description=voice_desc)
    if backend == "voxcpm2":
        return synthesize_with_voxcpm2(text, voice_description=voice_desc)
    if backend == "elevenlabs":
        return synthesize_with_elevenlabs(text)

    # Auto: try in order — HF is fastest (API), then local models, then ElevenLabs
    print(f"[TTS] Auto mode — trying HF first...")
    audio = synthesize_with_hf(text, voice=voice)
    if audio:
        return audio

    print("[TTS] HF failed, trying VoxCPM2...")
    audio = synthesize_with_voxcpm2(text, voice_description=voice_desc)
    if audio:
        return audio

    print("[TTS] VoxCPM2 failed, trying Chatterbox...")
    audio = synthesize_with_chatterbox(text, voice_description=voice_desc)
    if audio:
        return audio

    print("[TTS] Chatterbox failed, trying ElevenLabs...")
    return synthesize_with_elevenlabs(text)

# ── Segment cache path ────────────────────────────────────────────────────────
def get_segment_cache_path(category: str, text: str) -> Path:
    import hashlib
    h = hashlib.md5(f"{category}{text}".encode()).hexdigest()
    return SEG_DIR / f"segment_{h}.wav"