src/template_manager.py

"""
Template management for different infographic styles
"""
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass
from config import Config, TEMPLATE_COLORS
import logging

logger = logging.getLogger(__name__)

@dataclass
class TemplateStyle:
    """Template style configuration"""
    name: str
    colors: Dict[str, str]
    fonts: Dict[str, Tuple[str, int, str]]
    spacing: Dict[str, int]
    layout_rules: Dict[str, any]
    visual_style: Dict[str, any]

class TemplateManager:
    """Manage and apply infographic templates"""
    
    def __init__(self):
        """Initialize template manager with predefined styles"""
        self.templates = self._initialize_templates()
        logger.info(f"Template manager initialized with {len(self.templates)} templates")
    
    def get_template(self, template_name: str, custom_color: Optional[str] = None) -> TemplateStyle:
        """
        Get template configuration
        
        Args:
            template_name: Name of the template
            custom_color: Optional custom primary color
            
        Returns:
            TemplateStyle configuration
        """
        template = self.templates.get(template_name, self.templates['Modern'])
        
        # Apply custom color if provided
        if custom_color:
            template = self._apply_custom_color(template, custom_color)
        
        return template
    
    def apply_template_to_content(self, content: Dict, template_name: str, 
                                 custom_color: Optional[str] = None, 
                                 layout_type: str = "Vertical") -> Dict:
        """
        Apply template styling to content
        
        Args:
            content: Analyzed content from TextProcessor and GeminiClient
            template_name: Selected template name
            custom_color: Optional custom color
            layout_type: Layout arrangement
            
        Returns:
            Styled content ready for image generation
        """
        template = self.get_template(template_name, custom_color)
        
        styled_content = {
            'template': template,
            'layout_type': layout_type,
            'title': self._style_title(content, template),
            'sections': self._style_sections(content, template, layout_type),
            'visual_elements': self._style_visual_elements(content, template),
            'design_specs': self._create_design_specs(template, layout_type),
            'color_scheme': template.colors,
            'typography': template.fonts
        }
        
        return styled_content
    
    def _initialize_templates(self) -> Dict[str, TemplateStyle]:
        """Initialize predefined templates"""
        templates = {}
        
        # Modern Template
        templates['Modern'] = TemplateStyle(
            name='Modern',
            colors=TEMPLATE_COLORS['Modern'],
            fonts={
                'title': ('Arial', 32, 'bold'),
                'heading': ('Arial', 24, 'bold'),
                'body': ('Arial', 16, 'normal'),
                'caption': ('Arial', 14, 'normal')
            },
            spacing={
                'title_margin': 40,
                'section_margin': 30,
                'line_spacing': 8,
                'padding': 60
            },
            layout_rules={
                'max_sections': 6,
                'title_position': 'top',
                'alignment': 'left',
                'use_icons': True,
                'use_dividers': True
            },
            visual_style={
                'border_radius': 12,
                'shadow': True,
                'gradient': False,
                'icon_style': 'filled',
                'emphasis': 'color'
            }
        )
        
        # Corporate Template
        templates['Corporate'] = TemplateStyle(
            name='Corporate',
            colors=TEMPLATE_COLORS['Corporate'],
            fonts={
                'title': ('Arial', 30, 'bold'),
                'heading': ('Arial', 22, 'bold'),
                'body': ('Arial', 15, 'normal'),
                'caption': ('Arial', 13, 'normal')
            },
            spacing={
                'title_margin': 35,
                'section_margin': 25,
                'line_spacing': 6,
                'padding': 50
            },
            layout_rules={
                'max_sections': 8,
                'title_position': 'top',
                'alignment': 'left',
                'use_icons': False,
                'use_dividers': True
            },
            visual_style={
                'border_radius': 4,
                'shadow': False,
                'gradient': False,
                'icon_style': 'outline',
                'emphasis': 'typography'
            }
        )
        
        # Creative Template (continued)
        templates['Creative'] = TemplateStyle(
            name='Creative',
            colors=TEMPLATE_COLORS['Creative'],
            fonts={
                'title': ('Arial', 36, 'bold'),
                'heading': ('Arial', 26, 'bold'),
                'body': ('Arial', 17, 'normal'),
                'caption': ('Arial', 15, 'italic')
            },
            spacing={
                'title_margin': 50,
                'section_margin': 35,
                'line_spacing': 10,
                'padding': 70
            },
            layout_rules={
                'max_sections': 5,
                'title_position': 'center',
                'alignment': 'center',
                'use_icons': True,
                'use_dividers': False
            },
            visual_style={
                'border_radius': 20,
                'shadow': True,
                'gradient': True,
                'icon_style': 'colorful',
                'emphasis': 'visual'
            }
        )
        
        # Minimalist Template
        templates['Minimalist'] = TemplateStyle(
            name='Minimalist',
            colors=TEMPLATE_COLORS['Minimalist'],
            fonts={
                'title': ('Arial', 28, 'normal'),
                'heading': ('Arial', 20, 'bold'),
                'body': ('Arial', 14, 'normal'),
                'caption': ('Arial', 12, 'normal')
            },
            spacing={
                'title_margin': 60,
                'section_margin': 40,
                'line_spacing': 12,
                'padding': 80
            },
            layout_rules={
                'max_sections': 4,
                'title_position': 'top',
                'alignment': 'left',
                'use_icons': False,
                'use_dividers': True
            },
            visual_style={
                'border_radius': 0,
                'shadow': False,
                'gradient': False,
                'icon_style': 'minimal',
                'emphasis': 'whitespace'
            }
        )
        
        # Academic Template
        templates['Academic'] = TemplateStyle(
            name='Academic',
            colors=TEMPLATE_COLORS['Academic'],
            fonts={
                'title': ('Arial', 28, 'bold'),
                'heading': ('Arial', 20, 'bold'),
                'body': ('Arial', 14, 'normal'),
                'caption': ('Arial', 12, 'italic')
            },
            spacing={
                'title_margin': 30,
                'section_margin': 20,
                'line_spacing': 6,
                'padding': 40
            },
            layout_rules={
                'max_sections': 10,
                'title_position': 'top',
                'alignment': 'justify',
                'use_icons': False,
                'use_dividers': True
            },
            visual_style={
                'border_radius': 0,
                'shadow': False,
                'gradient': False,
                'icon_style': 'academic',
                'emphasis': 'content'
            }
        )
        
        return templates
    
    def _apply_custom_color(self, template: TemplateStyle, custom_color: str) -> TemplateStyle:
        """Apply custom primary color to template"""
        new_colors = template.colors.copy()
        new_colors['primary'] = custom_color
        
        # Create a new template with custom color
        return TemplateStyle(
            name=template.name,
            colors=new_colors,
            fonts=template.fonts,
            spacing=template.spacing,
            layout_rules=template.layout_rules,
            visual_style=template.visual_style
        )
    
    def _style_title(self, content: Dict, template: TemplateStyle) -> Dict:
        """Style the title section"""
        # Try to get AI-generated title first
        title_text = ""
        if 'ai_insights' in content and hasattr(content['ai_insights'], 'get'):
            title_text = content['ai_insights'].get('title', '')
        
        # Fallback to first line or generated title
        if not title_text:
            first_line = content.get('original_text', '').split('\n')[0].strip()
            if len(first_line.split()) <= 10:
                title_text = first_line
            else:
                title_text = ' '.join(content.get('original_text', '').split()[:8])
        
        return {
            'text': title_text,
            'font': template.fonts['title'],
            'color': template.colors['primary'],
            'alignment': template.layout_rules['title_position'],
            'margin': template.spacing['title_margin']
        }
    
    def _style_sections(self, content: Dict, template: TemplateStyle, layout_type: str) -> List[Dict]:
        """Style content sections"""
        sections = content.get('sections', [])
        max_sections = template.layout_rules['max_sections']
        
        # Limit sections based on template
        sections = sections[:max_sections]
        
        styled_sections = []
        for i, section in enumerate(sections):
            styled_section = {
                'id': section.get('id', i + 1),
                'content': section.get('content', ''),
                'condensed_text': section.get('condensed_text', section.get('content', '')[:100]),
                'type': section.get('type', 'body'),
                'priority': section.get('priority', 5),
                'visual_treatment': section.get('visual_treatment', 'bullet'),
                'color_suggestion': section.get('color_suggestion', 'primary'),
                'styling': self._get_section_styling(section, template, i)
            }
            styled_sections.append(styled_section)
        
        return styled_sections
    
    def _get_section_styling(self, section: Dict, template: TemplateStyle, index: int) -> Dict:
        """Get styling for individual section"""
        section_type = section.get('type', 'body')
        
        if section_type == 'header':
            font = template.fonts['heading']
            color = template.colors['primary']
        elif section_type == 'key_point':
            font = template.fonts['body']
            color = template.colors['accent']
        elif section_type == 'data':
            font = template.fonts['heading']
            color = template.colors['secondary']
        else:
            font = template.fonts['body']
            color = template.colors['text']
        
        return {
            'font': font,
            'color': color,
            'background_color': self._get_alternating_background(index, template),
            'margin': template.spacing['section_margin'],
            'padding': 20,
            'border_radius': template.visual_style['border_radius'],
            'alignment': template.layout_rules['alignment']
        }
    
    def _get_alternating_background(self, index: int, template: TemplateStyle) -> str:
        """Get alternating background colors"""
        if template.name == 'Minimalist':
            return 'transparent'
        elif index % 2 == 0:
            return template.colors['background']
        else:
            return self._lighten_color(template.colors['background'], 0.05)
    
    def _lighten_color(self, hex_color: str, factor: float) -> str:
        """Lighten a hex color"""
        # Simple color lightening (fallback)
        return hex_color  # For now, return original color
    
    def _style_visual_elements(self, content: Dict, template: TemplateStyle) -> List[Dict]:
        """Style visual elements like icons and charts"""
        visual_elements = content.get('visual_elements', [])
        
        styled_elements = []
        for element in visual_elements:
            styled_element = {
                'type': element.get('type', 'icon'),
                'description': element.get('description', ''),
                'placement': element.get('placement', 'body'),
                'importance': element.get('importance', 5),
                'styling': {
                    'color': template.colors['accent'],
                    'size': self._get_element_size(element.get('importance', 5)),
                    'style': template.visual_style.get('icon_style', 'filled')
                }
            }
            styled_elements.append(styled_element)
        
        return styled_elements
    
    def _get_element_size(self, importance: int) -> int:
        """Get element size based on importance"""
        size_map = {
            1: 16, 2: 20, 3: 24, 4: 28, 5: 32,
            6: 36, 7: 40, 8: 44, 9: 48, 10: 52
        }
        return size_map.get(importance, 32)
    
    def _create_design_specs(self, template: TemplateStyle, layout_type: str) -> Dict:
        """Create comprehensive design specifications"""
        return {
            'canvas_size': self._get_canvas_size(layout_type),
            'margins': {
                'top': template.spacing['padding'],
                'bottom': template.spacing['padding'],
                'left': template.spacing['padding'],
                'right': template.spacing['padding']
            },
            'grid': self._get_grid_specs(layout_type),
            'typography': {
                'line_height': template.spacing['line_spacing'],
                'paragraph_spacing': template.spacing['section_margin']
            },
            'colors': template.colors,
            'visual_effects': {
                'shadows': template.visual_style['shadow'],
                'gradients': template.visual_style['gradient'],
                'border_radius': template.visual_style['border_radius']
            },
            'layout_rules': template.layout_rules
        }
    
    def _get_canvas_size(self, layout_type: str) -> Tuple[int, int]:
        """Get canvas size based on layout"""
        sizes = {
            'Vertical': (Config.DEFAULT_WIDTH, Config.DEFAULT_HEIGHT),
            'Horizontal': (Config.DEFAULT_HEIGHT, Config.DEFAULT_WIDTH),
            'Grid': (Config.DEFAULT_WIDTH, Config.DEFAULT_WIDTH),
            'Flow': (Config.DEFAULT_WIDTH, int(Config.DEFAULT_HEIGHT * 0.75))
        }
        return sizes.get(layout_type, (Config.DEFAULT_WIDTH, Config.DEFAULT_HEIGHT))
    
    def _get_grid_specs(self, layout_type: str) -> Dict:
        """Get grid specifications for layout"""
        grid_specs = {
            'Vertical': {'columns': 1, 'rows': 'auto', 'gap': 30},
            'Horizontal': {'columns': 2, 'rows': 'auto', 'gap': 40},
            'Grid': {'columns': 2, 'rows': 3, 'gap': 25},
            'Flow': {'columns': 'auto', 'rows': 'auto', 'gap': 20}
        }
        return grid_specs.get(layout_type, grid_specs['Vertical'])
    
    def get_available_templates(self) -> List[str]:
        """Get list of available template names"""
        return list(self.templates.keys())
    
    def get_template_preview_data(self, template_name: str) -> Dict:
        """Get template preview information"""
        template = self.templates.get(template_name)
        if not template:
            return {}
        
        return {
            'name': template.name,
            'colors': template.colors,
            'description': self._get_template_description(template.name),
            'best_for': self._get_template_use_cases(template.name)
        }
    
    def _get_template_description(self, template_name: str) -> str:
        """Get template description"""
        descriptions = {
            'Modern': 'Clean, contemporary design with vibrant colors and modern typography',
            'Corporate': 'Professional business style with conservative colors and clean lines',
            'Creative': 'Artistic and expressive with bold colors and creative elements',
            'Minimalist': 'Simple, elegant design focusing on whitespace and essential elements',
            'Academic': 'Scholarly presentation style optimized for educational content'
        }
        return descriptions.get(template_name, 'Professional infographic template')
    
    def _get_template_use_cases(self, template_name: str) -> List[str]:
        """Get template use cases"""
        use_cases = {
            'Modern': ['Social Media', 'Marketing', 'Tech Content', 'Startups'],
            'Corporate': ['Business Reports', 'Annual Reports', 'Corporate Communications', 'Finance'],
            'Creative': ['Creative Industries', 'Art Projects', 'Entertainment', 'Personal Branding'],
            'Minimalist': ['Luxury Brands', 'Architecture', 'Design Portfolios', 'Premium Products'],
            'Academic': ['Research Papers', 'Educational Content', 'Scientific Reports', 'Publications']
        }
        return use_cases.get(template_name, ['General Purpose'])