src/gemini_client.py

"""
Google Gemini API client for text analysis and content generation
"""
import os
import json
import logging
from typing import Dict, List, Optional, Any
import google.generativeai as genai
from config import Config

logger = logging.getLogger(__name__)

class GeminiClient:
    """Client for Google Gemini API integration"""
    
    def __init__(self):
        """Initialize Gemini client"""
        self.api_key = Config.GEMINI_API_KEY
        if not self.api_key:
            logger.warning("Gemini API key not found. Some features may be limited.")
            self.client = None
        else:
            try:
                genai.configure(api_key=self.api_key)
                self.model = genai.GenerativeModel(Config.GEMINI_MODEL)
                self.client = self.model
                logger.info("Gemini client initialized successfully")
            except Exception as e:
                logger.error(f"Failed to initialize Gemini client: {e}")
                self.client = None
    
    def analyze_content(self, text_analysis: Dict) -> Dict:
        """
        Analyze content using Gemini API
        
        Args:
            text_analysis: Text analysis from TextProcessor
            
        Returns:
            Enhanced analysis with AI insights
        """
        if not self.client:
            return self._fallback_analysis(text_analysis)
        
        try:
            prompt = self._create_analysis_prompt(text_analysis)
            response = self.client.generate_content(
                prompt,
                generation_config=genai.types.GenerationConfig(
                    temperature=Config.GEMINI_TEMPERATURE,
                    max_output_tokens=Config.GEMINI_MAX_TOKENS,
                )
            )
            
            ai_analysis = self._parse_gemini_response(response.text)
            
            # Combine with original analysis
            enhanced_analysis = {
                **text_analysis,
                'ai_insights': ai_analysis,
                'suggested_template': ai_analysis.get('template', 'Modern'),
                'visual_elements': ai_analysis.get('visual_elements', []),
                'content_structure': ai_analysis.get('structure', {}),
                'design_recommendations': ai_analysis.get('design', {})
            }
            
            return enhanced_analysis
            
        except Exception as e:
            logger.error(f"Gemini analysis failed: {e}")
            return self._fallback_analysis(text_analysis)
    
    def generate_title(self, text: str) -> str:
        """Generate an engaging title for the content"""
        if not self.client:
            return self._extract_title_fallback(text)
        
        try:
            prompt = f"""
            Generate a compelling, concise title for this content that would work well in an infographic:
            
            Content: {text[:500]}...
            
            Requirements:
            - Maximum 8 words
            - Engaging and clear
            - Suitable for visual presentation
            - Return only the title, nothing else
            """
            
            response = self.client.generate_content(prompt)
            title = response.text.strip().strip('"').strip("'")
            return title if len(title.split()) <= 8 else self._extract_title_fallback(text)
            
        except Exception as e:
            logger.error(f"Title generation failed: {e}")
            return self._extract_title_fallback(text)
    
    def suggest_visual_elements(self, content: Dict) -> List[Dict]:
        """Suggest visual elements for the infographic"""
        if not self.client:
            return self._fallback_visual_elements(content)
        
        try:
            prompt = f"""
            Based on this content analysis, suggest 5-8 visual elements for an infographic:
            
            Keywords: {', '.join(content.get('keywords', [])[:10])}
            Key Points: {content.get('key_points', [])[:3]}
            Content Type: {content.get('sentiment', 'neutral')}
            Has Data: {content.get('data_elements', {}).get('has_data', False)}
            
            For each visual element, provide:
            - type (icon, chart, illustration, diagram)
            - description
            - placement_suggestion (header, body, footer)
            - importance (1-10)
            
            Return as JSON array with these fields.
            """
            
            response = self.client.generate_content(prompt)
            visual_elements = self._parse_json_response(response.text)
            
            return visual_elements if isinstance(visual_elements, list) else self._fallback_visual_elements(content)
            
        except Exception as e:
            logger.error(f"Visual element suggestion failed: {e}")
            return self._fallback_visual_elements(content)
    
    def optimize_content_for_visual(self, sections: List[Dict]) -> List[Dict]:
        """Optimize content sections for visual presentation"""
        if not self.client:
            return self._fallback_optimize_content(sections)
        
        try:
            content_summary = "\n".join([f"Section {s['id']}: {s['content'][:100]}..." for s in sections[:5]])
            
            prompt = f"""
            Optimize these content sections for infographic presentation:
            
            {content_summary}
            
            For each section, provide:
            - condensed_text (max 15 words)
            - visual_treatment (bullet, number, highlight, quote)
            - priority (1-10)
            - color_suggestion (primary, secondary, accent)
            
            Return as JSON array maintaining the same section IDs.
            """
            
            response = self.client.generate_content(prompt)
            optimized = self._parse_json_response(response.text)
            
            # Merge with original sections
            if isinstance(optimized, list):
                for i, section in enumerate(sections):
                    if i < len(optimized):
                        section.update(optimized[i])
            
            return sections
            
        except Exception as e:
            logger.error(f"Content optimization failed: {e}")
            return self._fallback_optimize_content(sections)
    
    def _create_analysis_prompt(self, analysis: Dict) -> str:
        """Create comprehensive analysis prompt for Gemini"""
        return f"""
        Analyze this content for infographic design and provide recommendations:
        
        Content Statistics:
        - Word count: {analysis['statistics']['word_count']}
        - Sentiment: {analysis['sentiment']}
        - Key points: {len(analysis['key_points'])}
        - Has data elements: {analysis['data_elements']['has_data']}
        
        Content Structure:
        - Paragraphs: {analysis['structure']['paragraph_count']}
        - Has headers: {analysis['structure']['has_headers']}
        - Has lists: {analysis['structure']['has_lists']}
        - Has numbers: {analysis['structure']['has_numbers']}
        
        Top Keywords: {', '.join(analysis['keywords'][:8])}
        
        Based on this analysis, provide:
        1. Best template style (Modern, Corporate, Creative, Minimalist, Academic)
        2. Recommended layout (Vertical, Horizontal, Grid, Flow)
        3. Color scheme suggestions (primary mood)
        4. Visual element recommendations
        5. Content hierarchy suggestions
        
        Return response as JSON with these keys:
        - template
        - layout  
        - color_mood
        - visual_elements (array)
        - structure (object)
        - design (object)
        """
    
    def _parse_gemini_response(self, response_text: str) -> Dict:
        """Parse Gemini response into structured data"""
        try:
            # Try to extract JSON from response
            json_match = response_text.find('{')
            if json_match != -1:
                json_end = response_text.rfind('}') + 1
                json_str = response_text[json_match:json_end]
                return json.loads(json_str)
        except:
            pass
        
        # Fallback parsing
        return {
            'template': self._extract_template_from_text(response_text),
            'layout': self._extract_layout_from_text(response_text),
            'color_mood': self._extract_color_mood_from_text(response_text),
            'visual_elements': [],
            'structure': {'hierarchy': 'standard'},
            'design': {'emphasis': 'balanced'}
        }
    
    def _parse_json_response(self, response_text: str) -> Any:
        """Parse JSON from Gemini response"""
        try:
            json_match = response_text.find('[')
            if json_match != -1:
                json_end = response_text.rfind(']') + 1
                json_str = response_text[json_match:json_end]
                return json.loads(json_str)
            
            json_match = response_text.find('{')
            if json_match != -1:
                json_end = response_text.rfind('}') + 1
                json_str = response_text[json_match:json_end]
                return json.loads(json_str)
        except:
            pass
        
        return []
    
    def _extract_template_from_text(self, text: str) -> str:
        """Extract template suggestion from text"""
        templates = ['Modern', 'Corporate', 'Creative', 'Minimalist', 'Academic']
        for template in templates:
            if template.lower() in text.lower():
                return template
        return 'Modern'
    
    def _extract_layout_from_text(self, text: str) -> str:
        """Extract layout suggestion from text"""
        layouts = ['Vertical', 'Horizontal', 'Grid', 'Flow']
        for layout in layouts:
            if layout.lower() in text.lower():
                return layout
        return 'Vertical'
    
    def _extract_color_mood_from_text(self, text: str) -> str:
        """Extract color mood from text"""
        if any(word in text.lower() for word in ['professional', 'business', 'corporate']):
            return 'professional'
        elif any(word in text.lower() for word in ['creative', 'vibrant', 'colorful']):
            return 'creative'
        elif any(word in text.lower() for word in ['minimal', 'clean', 'simple']):
            return 'minimal'
        else:
            return 'balanced'
    
    def _extract_title_fallback(self, text: str) -> str:
        """Fallback title extraction"""
        first_line = text.split('\n')[0].strip()
        if len(first_line.split()) <= 8:
            return first_line
        
        # Extract first sentence
        first_sentence = text.split('.')[0].strip()
        if len(first_sentence.split()) <= 8:
            return first_sentence
        
        # Create title from keywords
        words = text.split()[:8]
        return ' '.join(words).title()
    
    def _fallback_analysis(self, text_analysis: Dict) -> Dict:
        """Fallback analysis when Gemini is not available"""
        structure = text_analysis.get('structure', {})
        sentiment = text_analysis.get('sentiment', 'neutral')
        
        # Rule-based template suggestion
        template = 'Modern'
        if sentiment == 'positive' and structure.get('has_data'):
            template = 'Corporate'
        elif len(text_analysis.get('keywords', [])) > 10:
            template = 'Creative'
        elif structure.get('paragraph_count', 0) > 5:
            template = 'Academic'
        
        return {
            **text_analysis,
            'ai_insights': {
                'template': template,
                'layout': structure.get('suggested_layout', 'Vertical'),
                'color_mood': 'balanced'
            },
            'suggested_template': template,
            'visual_elements': self._fallback_visual_elements(text_analysis),
            'content_structure': {'hierarchy': 'standard'},
            'design_recommendations': {'emphasis': 'balanced'}
        }
    
    def _fallback_visual_elements(self, content: Dict) -> List[Dict]:
        """Fallback visual elements"""
        elements = []
        
        if content.get('data_elements', {}).get('has_data'):
            elements.append({
                'type': 'chart',
                'description': 'Data visualization',
                'placement': 'body',
                'importance': 8
            })
        
        keywords = content.get('keywords', [])[:3]
        for i, keyword in enumerate(keywords):
            elements.append({
                'type': 'icon',
                'description': f'Icon for {keyword}',
                'placement': 'body',
                'importance': 6 - i
            })
        
        return elements
    
    def _fallback_optimize_content(self, sections: List[Dict]) -> List[Dict]:
        """Fallback content optimization"""
        for section in sections:
            content = section.get('content', '')
            words = content.split()[:15]
            section['condensed_text'] = ' '.join(words)
            section['visual_treatment'] = 'bullet' if len(words) < 10 else 'highlight'
            section['priority'] = section.get('priority', 5)
            section['color_suggestion'] = 'primary'
        
        return sections