#!/usr/bin/env python3
"""
AI Relevance Filter for News Digest Bot
Uses AI to score articles for relevance, importance, and quality.
"""

import json
import logging
import os
import requests
from datetime import datetime
from typing import List, Dict, Any, Optional
import time

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class AIFilter:
    def __init__(self, api_key: str = None, model: str = "gpt-3.5-turbo"):
        """Initialize AI filter with API configuration."""
        self.api_key = api_key or os.getenv('OPENAI_API_KEY')
        self.model = model
        self.base_url = "https://api.openai.com/v1/chat/completions"
        
        if not self.api_key:
            logger.warning("No OpenAI API key found. AI filtering will be disabled.")
            self.enabled = False
        else:
            self.enabled = True
    
    def score_article(self, article: Dict[str, Any]) -> Dict[str, Any]:
        """Score a single article for relevance and importance."""
        if not self.enabled:
            # Fallback scoring without AI
            return self.fallback_score(article)
        
        try:
            prompt = self.create_scoring_prompt(article)
            response = self.call_openai_api(prompt)
            
            if response:
                score_data = self.parse_ai_response(response)
                score_data['scored_by'] = 'ai'
                return score_data
            else:
                return self.fallback_score(article)
                
        except Exception as e:
            logger.error(f"Error scoring article '{article.get('title', 'Unknown')}': {e}")
            return self.fallback_score(article)
    
    def create_scoring_prompt(self, article: Dict[str, Any]) -> str:
        """Create a prompt for AI scoring of the article."""
        title = article.get('title', '')
        description = article.get('description', '')
        source_name = article.get('source', {}).get('name', '')
        category = article.get('source', {}).get('category', '')
        
        prompt = f"""
Please score this AI news article on the following criteria (0-10 scale):

ARTICLE:
Title: {title}
Description: {description}
Source: {source_name} ({category})

SCORING CRITERIA:
1. Relevance (0-10): How relevant is this to AI developments, research, or industry?
2. Importance (0-10): How significant is this news for the AI field?
3. Novelty (0-10): How new or unique is this information?
4. Quality (0-10): How well-written and informative is the content?

ADDITIONAL FACTORS:
- Breaking news or major announcements should score higher
- Technical research breakthroughs are highly valued
- Industry partnerships, funding, and product launches are important
- Opinion pieces and general commentary score lower
- Duplicate or rehashed content scores lower

Please respond with ONLY a JSON object in this format:
{{
    "relevance": 8,
    "importance": 7,
    "novelty": 6,
    "quality": 8,
    "overall": 7.25,
    "category": "research|product|funding|partnership|opinion|other",
    "reasoning": "Brief explanation of the scoring"
}}"""

        return prompt
    
    def call_openai_api(self, prompt: str) -> Optional[str]:
        """Make API call to OpenAI for scoring."""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        data = {
            "model": self.model,
            "messages": [
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            "max_tokens": 300,
            "temperature": 0.3
        }
        
        try:
            response = requests.post(self.base_url, headers=headers, json=data, timeout=30)
            response.raise_for_status()
            
            result = response.json()
            content = result['choices'][0]['message']['content'].strip()
            return content
            
        except requests.exceptions.RequestException as e:
            logger.error(f"API request error: {e}")
            return None
        except (KeyError, IndexError) as e:
            logger.error(f"API response parsing error: {e}")
            return None
    
    def parse_ai_response(self, response: str) -> Dict[str, Any]:
        """Parse AI response into scoring data."""
        try:
            # Try to find JSON in the response
            import re
            json_match = re.search(r'\{.*\}', response, re.DOTALL)
            if json_match:
                json_str = json_match.group(0)
                score_data = json.loads(json_str)
                
                # Validate required fields
                required_fields = ['relevance', 'importance', 'novelty', 'quality', 'overall']
                for field in required_fields:
                    if field not in score_data:
                        raise ValueError(f"Missing required field: {field}")
                
                return score_data
            else:
                raise ValueError("No JSON found in response")
                
        except (json.JSONDecodeError, ValueError) as e:
            logger.error(f"Error parsing AI response: {e}")
            logger.debug(f"Response was: {response}")
            return self.fallback_score({})
    
    def fallback_score(self, article: Dict[str, Any]) -> Dict[str, Any]:
        """Provide fallback scoring when AI is unavailable."""
        # Simple heuristic-based scoring
        source_weight = article.get('source', {}).get('weight', 0.5)
        category = article.get('source', {}).get('category', 'general')
        
        # Category-based scoring
        category_scores = {
            'research': 8.0,
            'industry': 7.0,
            'business': 6.0,
            'startup': 6.5,
            'technical': 7.5,
            'consumer': 5.5,
            'ethics': 6.0,
            'safety': 7.0,
            'opensource': 6.5,
            'future': 5.0,
            'general': 5.5
        }
        
        base_score = category_scores.get(category, 5.5)
        final_score = (base_score * 0.7) + (source_weight * 10 * 0.3)
        
        return {
            'relevance': final_score,
            'importance': final_score * 0.9,
            'novelty': final_score * 0.8,
            'quality': source_weight * 10,
            'overall': final_score,
            'category': 'unknown',
            'reasoning': 'Heuristic scoring (AI unavailable)',
            'scored_by': 'heuristic'
        }
    
    def filter_articles(self, articles: List[Dict[str, Any]], 
                       min_score: float = 6.0, 
                       max_articles: int = 50) -> List[Dict[str, Any]]:
        """Filter and score multiple articles."""
        logger.info(f"Filtering {len(articles)} articles...")
        
        scored_articles = []
        
        for i, article in enumerate(articles):
            try:
                logger.info(f"Scoring article {i+1}/{len(articles)}: {article.get('title', 'Unknown')[:60]}...")
                
                score_data = self.score_article(article)
                
                # Add scoring to article
                article['ai_score'] = score_data
                
                # Only include articles above minimum score
                if score_data.get('overall', 0) >= min_score:
                    scored_articles.append(article)
                
                # Add delay to respect rate limits
                if self.enabled:
                    time.sleep(1)  # 1 second delay between API calls
                
            except Exception as e:
                logger.error(f"Error processing article {i+1}: {e}")
                continue
        
        # Sort by overall score (descending)
        scored_articles.sort(key=lambda x: x.get('ai_score', {}).get('overall', 0), reverse=True)
        
        # Limit to maximum number of articles
        if len(scored_articles) > max_articles:
            scored_articles = scored_articles[:max_articles]
        
        logger.info(f"Filtered to {len(scored_articles)} high-quality articles")
        return scored_articles
    
    def get_top_stories(self, articles: List[Dict[str, Any]], 
                       count: int = 5) -> List[Dict[str, Any]]:
        """Get the top N stories from scored articles."""
        if not articles:
            return []
        
        # Ensure articles are scored
        for article in articles:
            if 'ai_score' not in article:
                article['ai_score'] = self.score_article(article)
        
        # Sort by score and return top N
        sorted_articles = sorted(articles, 
                               key=lambda x: x.get('ai_score', {}).get('overall', 0), 
                               reverse=True)
        
        return sorted_articles[:count]

def main():
    """Main function for command-line usage."""
    import argparse
    
    parser = argparse.ArgumentParser(description='Filter AI news articles by relevance')
    parser.add_argument('input', help='Input JSON file with articles')
    parser.add_argument('--output', help='Output file path')
    parser.add_argument('--min-score', type=float, default=6.0, help='Minimum score threshold')
    parser.add_argument('--max-articles', type=int, default=50, help='Maximum number of articles')
    parser.add_argument('--top-stories', type=int, help='Extract only top N stories')
    parser.add_argument('--api-key', help='OpenAI API key (or use OPENAI_API_KEY env var)')
    parser.add_argument('--model', default='gpt-3.5-turbo', help='OpenAI model to use')
    
    args = parser.parse_args()
    
    # Load articles
    try:
        with open(args.input, 'r') as f:
            data = json.load(f)
            articles = data.get('articles', [])
    except FileNotFoundError:
        print(f"Error: Input file not found: {args.input}")
        return
    except json.JSONDecodeError as e:
        print(f"Error parsing input file: {e}")
        return
    
    # Initialize filter
    ai_filter = AIFilter(api_key=args.api_key, model=args.model)
    
    # Filter articles
    if args.top_stories:
        filtered_articles = ai_filter.get_top_stories(articles, args.top_stories)
    else:
        filtered_articles = ai_filter.filter_articles(articles, args.min_score, args.max_articles)
    
    # Prepare output
    output_data = {
        'filtered_at': datetime.now().isoformat(),
        'original_count': len(articles),
        'filtered_count': len(filtered_articles),
        'filter_settings': {
            'min_score': args.min_score if not args.top_stories else None,
            'max_articles': args.max_articles if not args.top_stories else None,
            'top_stories': args.top_stories,
            'model': args.model,
            'ai_enabled': ai_filter.enabled
        },
        'articles': filtered_articles
    }
    
    # Save output
    output_file = args.output or f"filtered_articles_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    
    try:
        with open(output_file, 'w') as f:
            json.dump(output_data, f, indent=2)
        
        print(f"Filtered {len(articles)} articles to {len(filtered_articles)} high-quality articles")
        print(f"Saved to: {output_file}")
        
    except Exception as e:
        print(f"Error saving output: {e}")

if __name__ == "__main__":
    main()