tools/wikipedia_tool.py

"""
Wikipedia Specialized Tool for GAIA Agent
Direct Wikipedia API integration with advanced search and data extraction
"""

import os
import logging
import re
from typing import Dict, List, Any, Optional, Union, Tuple
from dataclasses import dataclass
from datetime import datetime
import json

try:
    import wikipedia
    import requests
    WIKIPEDIA_AVAILABLE = True
except ImportError:
    WIKIPEDIA_AVAILABLE = False

logger = logging.getLogger(__name__)

@dataclass
class WikipediaArticle:
    """Structured Wikipedia article data."""
    title: str
    url: str
    content: str
    summary: str
    categories: List[str]
    infobox: Dict[str, Any]
    references: List[str]
    images: List[str]
    last_modified: Optional[str] = None
    page_id: Optional[int] = None
    featured_status: Optional[str] = None

@dataclass
class WikipediaSearchResult:
    """Wikipedia search result with metadata."""
    title: str
    snippet: str
    page_id: int
    url: str
    score: float = 0.0


class WikipediaSpecializedTool:
    """
    Specialized Wikipedia tool with advanced search and data extraction capabilities.
    
    Features:
    - Direct Wikipedia API integration
    - Category and article search
    - Historical data extraction
    - Featured article tracking
    - Structured data parsing
    - Infobox extraction
    - Timeline and date-based queries
    """
    
    def __init__(self, language: str = 'en'):
        """Initialize the Wikipedia specialized tool."""
        self.language = language
        self.base_api_url = f"https://{language}.wikipedia.org/api/rest_v1"
        self.api_url = f"https://{language}.wikipedia.org/w/api.php"
        
        if WIKIPEDIA_AVAILABLE:
            wikipedia.set_lang(language)
            logger.info(f"✅ Wikipedia tool initialized for language: {language}")
        else:
            logger.warning("⚠️ Wikipedia dependencies not available")
        
        # Cache for frequently accessed data
        self._cache = {}
        self._featured_articles_cache = {}
    
    def search_articles(self, query: str, limit: int = 10) -> List[WikipediaSearchResult]:
        """
        Search Wikipedia articles with advanced filtering.
        
        Args:
            query: Search query
            limit: Maximum number of results
            
        Returns:
            List of WikipediaSearchResult objects
        """
        if not WIKIPEDIA_AVAILABLE:
            logger.warning("⚠️ Wikipedia not available")
            return []
        
        try:
            logger.info(f"🔍 Searching Wikipedia for: {query}")
            
            # Use Wikipedia API for search
            params = {
                'action': 'query',
                'format': 'json',
                'list': 'search',
                'srsearch': query,
                'srlimit': limit,
                'srprop': 'snippet|titlesnippet|size|wordcount|timestamp'
            }
            
            response = requests.get(self.api_url, params=params)
            response.raise_for_status()
            data = response.json()
            
            results = []
            if 'query' in data and 'search' in data['query']:
                for item in data['query']['search']:
                    result = WikipediaSearchResult(
                        title=item['title'],
                        snippet=item.get('snippet', ''),
                        page_id=item['pageid'],
                        url=f"https://{self.language}.wikipedia.org/wiki/{item['title'].replace(' ', '_')}",
                        score=self._calculate_search_score(item, query)
                    )
                    results.append(result)
            
            # Sort by relevance score
            results.sort(key=lambda x: x.score, reverse=True)
            
            logger.info(f"✅ Found {len(results)} Wikipedia articles")
            return results
            
        except Exception as e:
            logger.error(f"❌ Wikipedia search error: {e}")
            return []
    
    def get_article(self, title: str, include_content: bool = True) -> Optional[WikipediaArticle]:
        """
        Get detailed Wikipedia article information.
        
        Args:
            title: Article title
            include_content: Whether to include full content
            
        Returns:
            WikipediaArticle object or None
        """
        if not WIKIPEDIA_AVAILABLE:
            return None
        
        try:
            # Check cache first
            cache_key = f"article_{title}_{include_content}"
            if cache_key in self._cache:
                return self._cache[cache_key]
            
            logger.info(f"📖 Fetching Wikipedia article: {title}")
            
            # Get basic page info
            page = wikipedia.page(title)
            
            # Get additional metadata via API
            metadata = self._get_article_metadata(page.pageid)
            
            # Extract infobox data
            infobox = self._extract_infobox(page.content)
            
            # Get categories
            categories = self._get_article_categories(page.pageid)
            
            # Create article object
            article = WikipediaArticle(
                title=page.title,
                url=page.url,
                content=page.content if include_content else "",
                summary=page.summary,
                categories=categories,
                infobox=infobox,
                references=page.references if hasattr(page, 'references') else [],
                images=page.images if hasattr(page, 'images') else [],
                page_id=page.pageid,
                last_modified=metadata.get('last_modified'),
                featured_status=metadata.get('featured_status')
            )
            
            # Cache the result
            self._cache[cache_key] = article
            
            logger.info(f"✅ Retrieved article: {title}")
            return article
            
        except wikipedia.exceptions.DisambiguationError as e:
            logger.warning(f"⚠️ Disambiguation needed for '{title}': {e.options[:5]}")
            # Try the first option
            if e.options:
                return self.get_article(e.options[0], include_content)
            return None
            
        except wikipedia.exceptions.PageError:
            logger.warning(f"⚠️ Wikipedia page not found: {title}")
            return None
            
        except Exception as e:
            logger.error(f"❌ Error fetching Wikipedia article '{title}': {e}")
            return None
    
    def search_by_category(self, category: str, limit: int = 20) -> List[str]:
        """
        Search articles by Wikipedia category.
        
        Args:
            category: Category name (e.g., "Studio albums")
            limit: Maximum number of articles
            
        Returns:
            List of article titles
        """
        try:
            logger.info(f"🏷️ Searching category: {category}")
            
            params = {
                'action': 'query',
                'format': 'json',
                'list': 'categorymembers',
                'cmtitle': f'Category:{category}',
                'cmlimit': limit,
                'cmtype': 'page'
            }
            
            response = requests.get(self.api_url, params=params)
            response.raise_for_status()
            data = response.json()
            
            articles = []
            if 'query' in data and 'categorymembers' in data['query']:
                articles = [item['title'] for item in data['query']['categorymembers']]
            
            logger.info(f"✅ Found {len(articles)} articles in category '{category}'")
            return articles
            
        except Exception as e:
            logger.error(f"❌ Category search error: {e}")
            return []
    
    def get_featured_articles(self, date: Optional[str] = None) -> List[Dict[str, Any]]:
        """
        Get featured articles for a specific date or current featured articles.
        
        Args:
            date: Date in YYYY-MM-DD format (optional)
            
        Returns:
            List of featured article information
        """
        try:
            cache_key = f"featured_{date or 'current'}"
            if cache_key in self._featured_articles_cache:
                return self._featured_articles_cache[cache_key]
            
            if date:
                logger.info(f"🌟 Getting featured articles for date: {date}")
                # Get featured article for specific date
                url = f"https://en.wikipedia.org/api/rest_v1/feed/featured/{date.replace('-', '/')}"
            else:
                logger.info("🌟 Getting current featured articles")
                # Get today's featured article
                today = datetime.now().strftime("%Y/%m/%d")
                url = f"https://en.wikipedia.org/api/rest_v1/feed/featured/{today}"
            
            response = requests.get(url)
            response.raise_for_status()
            data = response.json()
            
            featured_articles = []
            
            # Extract featured article of the day
            if 'tfa' in data:
                tfa = data['tfa']
                featured_articles.append({
                    'type': 'featured_article',
                    'title': tfa.get('title', ''),
                    'extract': tfa.get('extract', ''),
                    'url': tfa.get('content_urls', {}).get('desktop', {}).get('page', ''),
                    'date': date or datetime.now().strftime("%Y-%m-%d")
                })
            
            # Cache the result
            self._featured_articles_cache[cache_key] = featured_articles
            
            logger.info(f"✅ Retrieved {len(featured_articles)} featured articles")
            return featured_articles
            
        except Exception as e:
            logger.error(f"❌ Featured articles error: {e}")
            return []
    
    def search_by_date_range(self, start_date: str, end_date: str, query: str = "") -> List[str]:
        """
        Search articles created or modified within a date range.
        
        Args:
            start_date: Start date (YYYY-MM-DD)
            end_date: End date (YYYY-MM-DD)
            query: Optional search query
            
        Returns:
            List of article titles
        """
        try:
            logger.info(f"📅 Searching articles from {start_date} to {end_date}")
            
            # Convert dates to Wikipedia timestamp format
            start_ts = start_date.replace('-', '') + '000000'
            end_ts = end_date.replace('-', '') + '235959'
            
            params = {
                'action': 'query',
                'format': 'json',
                'list': 'recentchanges',
                'rcstart': end_ts,
                'rcend': start_ts,
                'rcnamespace': 0,  # Main namespace only
                'rctype': 'new|edit',
                'rclimit': 100
            }
            
            if query:
                # If query provided, search within the results
                params['list'] = 'search'
                params['srsearch'] = f'{query} incategory:"Articles created in {start_date[:4]}"'
                del params['rcstart']
                del params['rcend']
                del params['rcnamespace']
                del params['rctype']
            
            response = requests.get(self.api_url, params=params)
            response.raise_for_status()
            data = response.json()
            
            articles = []
            if query and 'query' in data and 'search' in data['query']:
                articles = [item['title'] for item in data['query']['search']]
            elif 'query' in data and 'recentchanges' in data['query']:
                articles = [item['title'] for item in data['query']['recentchanges']]
            
            logger.info(f"✅ Found {len(articles)} articles in date range")
            return articles
            
        except Exception as e:
            logger.error(f"❌ Date range search error: {e}")
            return []
    
    def extract_discography_info(self, artist_name: str, album_type: str = "studio") -> List[Dict[str, Any]]:
        """
        Extract discography information for an artist.
        
        Args:
            artist_name: Name of the artist
            album_type: Type of albums (studio, live, compilation)
            
        Returns:
            List of album information
        """
        try:
            logger.info(f"🎵 Extracting {album_type} albums for: {artist_name}")
            
            # Search for discography page
            discography_queries = [
                f"{artist_name} discography",
                f"{artist_name} albums",
                f"List of {artist_name} albums"
            ]
            
            discography_article = None
            for query in discography_queries:
                search_results = self.search_articles(query, limit=5)
                for result in search_results:
                    if any(word in result.title.lower() for word in ['discography', 'albums', 'list']):
                        discography_article = self.get_article(result.title)
                        break
                if discography_article:
                    break
            
            if not discography_article:
                logger.warning(f"⚠️ No discography found for {artist_name}")
                return []
            
            # Extract album information from content
            albums = self._parse_discography_content(discography_article.content, album_type)
            
            logger.info(f"✅ Found {len(albums)} {album_type} albums for {artist_name}")
            return albums
            
        except Exception as e:
            logger.error(f"❌ Discography extraction error: {e}")
            return []
    
    def _get_article_metadata(self, page_id: int) -> Dict[str, Any]:
        """Get additional metadata for an article."""
        try:
            params = {
                'action': 'query',
                'format': 'json',
                'pageids': page_id,
                'prop': 'info|revisions',
                'inprop': 'protection|talkid|watched|watchers|notificationtimestamp|subjectid|url|readable|preload|displaytitle',
                'rvprop': 'timestamp|user|comment',
                'rvlimit': 1
            }
            
            response = requests.get(self.api_url, params=params)
            response.raise_for_status()
            data = response.json()
            
            metadata = {}
            if 'query' in data and 'pages' in data['query']:
                page_data = list(data['query']['pages'].values())[0]
                
                if 'revisions' in page_data:
                    metadata['last_modified'] = page_data['revisions'][0]['timestamp']
                
                # Check if it's a featured article
                # This would require additional API calls to check featured status
                
            return metadata
            
        except Exception as e:
            logger.warning(f"⚠️ Error getting article metadata: {e}")
            return {}
    
    def _extract_infobox(self, content: str) -> Dict[str, Any]:
        """Extract infobox data from article content."""
        infobox = {}
        
        try:
            # Look for infobox patterns
            infobox_pattern = r'\{\{[Ii]nfobox[^}]*\}\}'
            matches = re.findall(infobox_pattern, content, re.DOTALL)
            
            if matches:
                infobox_text = matches[0]
                # Parse key-value pairs
                lines = infobox_text.split('\n')
                for line in lines:
                    if '=' in line and not line.strip().startswith('{{'):
                        parts = line.split('=', 1)
                        if len(parts) == 2:
                            key = parts[0].strip().replace('|', '')
                            value = parts[1].strip()
                            if key and value:
                                infobox[key] = value
            
        except Exception as e:
            logger.warning(f"⚠️ Error extracting infobox: {e}")
        
        return infobox
    
    def _get_article_categories(self, page_id: int) -> List[str]:
        """Get categories for an article."""
        try:
            params = {
                'action': 'query',
                'format': 'json',
                'pageids': page_id,
                'prop': 'categories',
                'cllimit': 100
            }
            
            response = requests.get(self.api_url, params=params)
            response.raise_for_status()
            data = response.json()
            
            categories = []
            if 'query' in data and 'pages' in data['query']:
                page_data = list(data['query']['pages'].values())[0]
                if 'categories' in page_data:
                    categories = [cat['title'].replace('Category:', '') for cat in page_data['categories']]
            
            return categories
            
        except Exception as e:
            logger.warning(f"⚠️ Error getting categories: {e}")
            return []
    
    def _calculate_search_score(self, item: Dict[str, Any], query: str) -> float:
        """Calculate relevance score for search results."""
        score = 0.0
        query_lower = query.lower()
        title_lower = item['title'].lower()
        snippet_lower = item.get('snippet', '').lower()
        
        # Title match scoring
        if query_lower == title_lower:
            score += 1.0
        elif query_lower in title_lower:
            score += 0.8
        elif any(word in title_lower for word in query_lower.split()):
            score += 0.6
        
        # Snippet match scoring
        if query_lower in snippet_lower:
            score += 0.4
        elif any(word in snippet_lower for word in query_lower.split()):
            score += 0.2
        
        # Size and word count boost
        size = item.get('size', 0)
        if size > 10000:  # Larger articles often more comprehensive
            score += 0.1
        
        return score
    
    def _parse_discography_content(self, content: str, album_type: str) -> List[Dict[str, Any]]:
        """Parse discography content to extract album information."""
        albums = []
        
        try:
            # Look for album sections
            lines = content.split('\n')
            current_section = ""
            
            for line in lines:
                line = line.strip()
                
                # Check for section headers
                if line.startswith('==') and album_type.lower() in line.lower():
                    current_section = album_type
                    continue
                elif line.startswith('==') and album_type.lower() not in line.lower():
                    current_section = ""
                    continue
                
                # If we're in the right section, look for album entries
                if current_section == album_type and line:
                    # Look for patterns like "* ''Album Name'' (Year)"
                    album_match = re.search(r"[*#]\s*['\"]?([^'\"]+)['\"]?\s*\((\d{4})\)", line)
                    if album_match:
                        album_name = album_match.group(1).strip()
                        year = album_match.group(2)
                        
                        albums.append({
                            'title': album_name,
                            'year': int(year),
                            'type': album_type
                        })
        
        except Exception as e:
            logger.warning(f"⚠️ Error parsing discography: {e}")
        
        return albums
    
    def search_mercedes_sosa_albums(self, start_year: int = 2000, end_year: int = 2009) -> List[Dict[str, Any]]:
        """
        Specific method to search for Mercedes Sosa studio albums in a date range.
        This addresses one of the failing GAIA questions.
        """
        try:
            logger.info(f"🎵 Searching Mercedes Sosa studio albums ({start_year}-{end_year})")
            
            # Get Mercedes Sosa discography
            albums = self.extract_discography_info("Mercedes Sosa", "studio")
            
            # Filter by date range
            filtered_albums = [
                album for album in albums 
                if start_year <= album.get('year', 0) <= end_year
            ]
            
            logger.info(f"✅ Found {len(filtered_albums)} Mercedes Sosa studio albums in {start_year}-{end_year}")
            return filtered_albums
            
        except Exception as e:
            logger.error(f"❌ Mercedes Sosa search error: {e}")
            return []
    
    def find_featured_article_by_date(self, target_date: str, topic_keywords: List[str]) -> Optional[str]:
        """
        Find featured article for a specific date matching topic keywords.
        This addresses the dinosaur Featured Article GAIA question.
        """
        try:
            logger.info(f"🌟 Searching featured article for {target_date} with keywords: {topic_keywords}")
            
            featured_articles = self.get_featured_articles(target_date)
            
            for article in featured_articles:
                title = article.get('title', '').lower()
                extract = article.get('extract', '').lower()
                
                # Check if any keywords match
                for keyword in topic_keywords:
                    if keyword.lower() in title or keyword.lower() in extract:
                        logger.info(f"✅ Found matching featured article: {article['title']}")
                        return article['title']
            
            logger.warning(f"⚠️ No featured article found for {target_date} with keywords {topic_keywords}")
            return None
            
        except Exception as e:
            logger.error(f"❌ Featured article search error: {e}")
            return None