MusicAnalyser/backend/app/services/narrative_service.py

import os
import json
import re
from typing import Dict, Any

try:
    from openai import OpenAI
except ImportError:
    OpenAI = None

try:
    from google import genai
except ImportError:
    genai = None


class NarrativeService:
    def __init__(self, model_name: str = "gpt-5-mini-2025-08-07"):
        self.model_name = model_name
        self.provider = self._detect_provider()
        self.client = self._init_client()

    def _detect_provider(self) -> str:
        openai_key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
        gemini_key = os.getenv("GEMINI_API_KEY")

        if self.model_name.startswith("gpt") and openai_key and OpenAI:
            return "openai"
        elif gemini_key and genai:
            return "gemini"
        elif openai_key and OpenAI:
            return "openai"
        elif gemini_key and genai:
            return "gemini"
        return "none"

    def _init_client(self):
        if self.provider == "openai":
            api_key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
            return OpenAI(api_key=api_key)
        elif self.provider == "gemini":
            api_key = os.getenv("GEMINI_API_KEY")
            return genai.Client(api_key=api_key)
        return None

    def generate_full_narrative(self, stats_json: Dict[str, Any]) -> Dict[str, Any]:
        if not self.client:
            print("WARNING: No LLM client available")
            return self._get_fallback_narrative()

        clean_stats = self._shape_payload(stats_json)
        prompt = self._build_prompt(clean_stats)

        try:
            if self.provider == "openai":
                return self._call_openai(prompt)
            elif self.provider == "gemini":
                return self._call_gemini(prompt)
        except Exception as e:
            print(f"LLM Generation Error: {e}")
            return self._get_fallback_narrative()

        return self._get_fallback_narrative()

    def generate_playlist_theme(self, listening_data: Dict[str, Any]) -> Dict[str, Any]:
        """Generate playlist theme based on daily listening patterns."""
        if not self.client:
            return self._get_fallback_theme()

        prompt = self._build_theme_prompt(listening_data)

        try:
            if self.provider == "openai":
                return self._call_openai_for_theme(prompt)
            elif self.provider == "gemini":
                return self._call_gemini_for_theme(prompt)
        except Exception as e:
            print(f"Theme generation error: {e}")
            return self._get_fallback_theme()

        return self._get_fallback_theme()

    def _call_openai_for_theme(self, prompt: str) -> Dict[str, Any]:
        response = self.client.chat.completions.create(
            model=self.model_name,
            messages=[
                {
                    "role": "system",
                    "content": "You are a specialized music curator. Output only valid JSON.",
                },
                {"role": "user", "content": prompt},
            ],
            response_format={"type": "json_object"},
        )
        return self._clean_and_parse_json(response.choices[0].message.content)

    def _call_gemini_for_theme(self, prompt: str) -> Dict[str, Any]:
        response = self.client.models.generate_content(
            model=self.model_name,
            contents=prompt,
            config=genai.types.GenerateContentConfig(
                response_mime_type="application/json"
            ),
        )
        return self._clean_and_parse_json(response.text)

    def _build_theme_prompt(self, data: Dict[str, Any]) -> str:
        return f"""Analyze this listening data from the last 6 hours and curate a specific "themed" playlist.

**DATA:**
- Peak hour: {data.get("peak_hour")}
- Avg energy: {data.get("avg_energy"):.2f}
- Avg valence: {data.get("avg_valence"):.2f}
- Top artists: {", ".join([a["name"] for a in data.get("top_artists", [])])}
- Total plays: {data.get("total_plays")}

**RULES:**
1. Create a "theme_name" (e.g. "Morning Coffee Jazz", "Midnight Deep Work").
2. Provide a "description" (2-3 sentences explaining why).
3. Identify 10-15 "curated_tracks" (song names only) that fit this vibe and the artists listed.
4. Return ONLY valid JSON.

**REQUIRED JSON:**
{{
  "theme_name": "String",
  "description": "String",
  "curated_tracks": ["Track 1", "Track 2", ...]
}}"""

    def _get_fallback_theme(self) -> Dict[str, Any]:
        return {
            "theme_name": "Daily Mix",
            "description": "A curated mix of your recent favorites.",
            "curated_tracks": [],
        }

    def _call_openai(self, prompt: str) -> Dict[str, Any]:
        response = self.client.chat.completions.create(
            model=self.model_name,
            messages=[
                {
                    "role": "system",
                    "content": "You are a witty music critic. Output only valid JSON.",
                },
                {"role": "user", "content": prompt},
            ],
            response_format={"type": "json_object"},
            max_completion_tokens=4000,
        )
        return self._clean_and_parse_json(response.choices[0].message.content)

    def _call_gemini(self, prompt: str) -> Dict[str, Any]:
        response = self.client.models.generate_content(
            model=self.model_name,
            contents=prompt,
            config=genai.types.GenerateContentConfig(
                response_mime_type="application/json"
            ),
        )
        return self._clean_and_parse_json(response.text)

    def _build_prompt(self, clean_stats: Dict[str, Any]) -> str:
        volume = clean_stats.get("volume", {})
        concentration = volume.get("concentration", {})
        time_habits = clean_stats.get("time_habits", {})
        vibe = clean_stats.get("vibe", {})
        peak_hour = time_habits.get("peak_hour")
        if isinstance(peak_hour, int):
            peak_listening = f"{peak_hour}:00"
        else:
            peak_listening = peak_hour or "N/A"
        concentration_score = (
            round(concentration.get("hhi", 0), 3)
            if concentration and concentration.get("hhi") is not None
            else "N/A"
        )
        playlist_diversity = (
            round(1 - concentration.get("hhi", 0), 3)
            if concentration and concentration.get("hhi") is not None
            else "N/A"
        )
        avg_energy = vibe.get("avg_energy", 0)
        avg_valence = vibe.get("avg_valence", 0)
        top_artists = volume.get("top_artists", [])
        top_artists_str = ", ".join(top_artists) if top_artists else "N/A"
        era_label = clean_stats.get("era", {}).get("musical_age", "N/A")

        return f"""Analyze this Spotify listening data and generate a personalized report.

**RULES:**
1. NO mental health diagnoses. Use behavioral descriptors only.
2. Be specific - reference actual metrics from the data.
3. Be playful but not cruel.
4. Return ONLY valid JSON.

**LISTENING HIGHLIGHTS:**
- Peak listening: {peak_listening}
- Concentration score: {concentration_score}
- Playlist diversity: {playlist_diversity}
- Average energy: {avg_energy:.2f}
- Average valence: {avg_valence:.2f}
- Top artists: {top_artists_str}

**DATA:**
{json.dumps(clean_stats, indent=2)}

**REQUIRED JSON:**
{{
  "vibe_check_short": "1-2 sentence hook for the hero banner.",
  "vibe_check": "2-3 paragraphs describing their overall listening personality.",
  "patterns": ["Observation 1", "Observation 2", "Observation 3"],
  "persona": "A creative label (e.g., 'The Genre Chameleon').",
  "era_insight": "Comment on Musical Age ({era_label}).",
  "roast": "1-2 sentence playful roast.",
  "comparison": "Compare to previous period if data exists."
}}"""

    def _shape_payload(self, stats: Dict[str, Any]) -> Dict[str, Any]:
        s = stats.copy()

        if "volume" in s:
            volume_copy = {
                k: v
                for k, v in s["volume"].items()
                if k not in ["top_tracks", "top_artists", "top_albums", "top_genres"]
            }
            volume_copy["top_tracks"] = [
                t["name"] for t in stats["volume"].get("top_tracks", [])[:5]
            ]
            volume_copy["top_artists"] = [
                a["name"] for a in stats["volume"].get("top_artists", [])[:5]
            ]
            volume_copy["top_genres"] = [
                g["name"] for g in stats["volume"].get("top_genres", [])[:5]
            ]
            s["volume"] = volume_copy

        if "time_habits" in s:
            s["time_habits"] = {
                k: v for k, v in s["time_habits"].items() if k != "heatmap"
            }

        if "sessions" in s:
            s["sessions"] = {
                k: v for k, v in s["sessions"].items() if k != "session_list"
            }

        return s

    def _clean_and_parse_json(self, raw_text: str) -> Dict[str, Any]:
        try:
            return json.loads(raw_text)
        except json.JSONDecodeError:
            pass

        try:
            match = re.search(r"\{.*\}", raw_text, re.DOTALL)
            if match:
                return json.loads(match.group(0))
        except:
            pass

        return self._get_fallback_narrative()

    def _get_fallback_narrative(self) -> Dict[str, Any]:
        return {
            "vibe_check_short": "Your taste is... interesting.",
            "vibe_check": "Data processing error. You're too mysterious to analyze right now.",
            "patterns": [],
            "persona": "The Enigma",
            "era_insight": "Time is a flat circle.",
            "roast": "You broke the machine. Congratulations.",
            "comparison": "N/A",
        }