miku-discord/bot/utils/danbooru_client.py

# danbooru_client.py
"""
Danbooru API client for fetching Hatsune Miku artwork.
"""

import aiohttp
import random
from typing import Optional, List, Dict
import asyncio

class DanbooruClient:
    """Client for interacting with Danbooru API"""
    
    BASE_URL = "https://danbooru.donmai.us"
    
    def __init__(self):
        self.session: Optional[aiohttp.ClientSession] = None
    
    async def _ensure_session(self):
        """Ensure aiohttp session exists"""
        if self.session is None or self.session.closed:
            self.session = aiohttp.ClientSession()
    
    async def close(self):
        """Close the aiohttp session"""
        if self.session and not self.session.closed:
            await self.session.close()
    
    async def search_miku_images(
        self, 
        tags: List[str] = None,
        rating: List[str] = None,
        limit: int = 100,
        random_page: bool = True
    ) -> List[Dict]:
        """
        Search for Hatsune Miku images on Danbooru.
        
        Args:
            tags: Additional tags to include (e.g., ["solo", "smile"])
            rating: Rating filter. Options: ["g", "s"] for general/sensitive
            limit: Number of results to fetch (max 200)
            random_page: If True, fetch from a random page (more variety)
        
        Returns:
            List of post dictionaries with image data
        """
        await self._ensure_session()
        
        # Build tag string
        tag_list = ["hatsune_miku"]
        if tags:
            tag_list.extend(tags)
        
        # Add rating filter using proper Danbooru syntax
        # We want general (g) and sensitive (s), so exclude questionable and explicit
        if rating and ("g" in rating or "s" in rating):
            # Exclude unwanted ratings
            tag_list.append("-rating:q")  # exclude questionable
            tag_list.append("-rating:e")  # exclude explicit
        
        # Combine tags
        tags_query = " ".join(tag_list)
        
        # Determine page
        page = random.randint(1, 20) if random_page else 1
        
        # Build request params
        params = {
            "tags": tags_query,
            "limit": min(limit, 200),  # Danbooru max is 200
            "page": page
        }
        
        try:
            url = f"{self.BASE_URL}/posts.json"
            print(f"🎨 Danbooru request: {url} with params: {params}")
            async with self.session.get(url, params=params, timeout=10) as response:
                if response.status == 200:
                    posts = await response.json()
                    print(f"🎨 Danbooru: Found {len(posts)} posts (page {page})")
                    return posts
                else:
                    error_text = await response.text()
                    print(f"⚠️ Danbooru API error: {response.status}")
                    print(f"⚠️ Request URL: {response.url}")
                    print(f"⚠️ Error details: {error_text[:500]}")
                    return []
        except asyncio.TimeoutError:
            print(f"⚠️ Danbooru API timeout")
            return []
        except Exception as e:
            print(f"⚠️ Danbooru API error: {e}")
            return []
    
    async def get_random_miku_image(
        self,
        mood: Optional[str] = None,
        exclude_tags: List[str] = None
    ) -> Optional[Dict]:
        """
        Get a single random Hatsune Miku image suitable for profile picture.
        
        Args:
            mood: Current mood to influence tag selection
            exclude_tags: Tags to exclude from search
        
        Returns:
            Post dictionary with image URL and metadata, or None
        """
        # Build tags based on mood
        tags = self._get_mood_tags(mood)
        
        # Add exclusions
        if exclude_tags:
            for tag in exclude_tags:
                tags.append(f"-{tag}")
        
        # Prefer solo images for profile pictures
        tags.append("solo")
        
        # Search with general and sensitive ratings only
        posts = await self.search_miku_images(
            tags=tags,
            rating=["g", "s"],  # general and sensitive only
            limit=50,
            random_page=True
        )
        
        if not posts:
            print("⚠️ No posts found, trying without mood tags")
            # Fallback: try without mood tags
            posts = await self.search_miku_images(
                rating=["g", "s"],
                limit=50,
                random_page=True
            )
        
        if not posts:
            return None
        
        # Filter posts with valid image URLs
        valid_posts = [
            p for p in posts 
            if p.get("file_url") and p.get("image_width", 0) >= 512
        ]
        
        if not valid_posts:
            print("⚠️ No valid posts with sufficient resolution")
            return None
        
        # Pick a random one
        selected = random.choice(valid_posts)
        
        print(f"🎨 Selected Danbooru post #{selected.get('id')} - {selected.get('tag_string_character', 'unknown character')}")
        
        return selected
    
    def _get_mood_tags(self, mood: Optional[str]) -> List[str]:
        """Get Danbooru tags based on Miku's current mood"""
        if not mood:
            return []
        
        mood_tag_map = {
            "bubbly": ["smile", "happy"],
            "sleepy": ["sleepy", "closed_eyes"],
            "curious": ["looking_at_viewer"],
            "shy": ["blush", "embarrassed"],
            "serious": ["serious"],
            "excited": ["happy", "open_mouth"],
            "silly": ["smile", "tongue_out"],
            "melancholy": ["sad", "tears"],
            "flirty": ["blush", "wink"],
            "romantic": ["blush", "heart"],
            "irritated": ["annoyed"],
            "angry": ["angry", "frown"],
            "neutral": [],
            "asleep": ["sleeping", "closed_eyes"],
        }
        
        tags = mood_tag_map.get(mood, [])
        
        # Only return one random tag to avoid over-filtering
        if tags:
            return [random.choice(tags)]
        return []
    
    def extract_image_url(self, post: Dict) -> Optional[str]:
        """Extract the best image URL from a Danbooru post"""
        # Prefer file_url (original), fallback to large_file_url
        return post.get("file_url") or post.get("large_file_url")
    
    def get_post_metadata(self, post: Dict) -> Dict:
        """Extract useful metadata from a Danbooru post"""
        return {
            "id": post.get("id"),
            "rating": post.get("rating"),
            "score": post.get("score"),
            "tags": post.get("tag_string", "").split(),
            "artist": post.get("tag_string_artist", "unknown"),
            "width": post.get("image_width"),
            "height": post.get("image_height"),
            "file_url": self.extract_image_url(post),
            "source": post.get("source", "")
        }


# Global instance
danbooru_client = DanbooruClient()