dom/template_manager.py

import json
import random
from pathlib import Path
from datetime import datetime
from typing import List, Dict, Any
import logging

logger = logging.getLogger(__name__)

class TemplateManager:
    """Manages prompt templates with placeholders and random selection."""
    
    def __init__(self, templates_dir: str = "templates", timezone_name: str = "UTC"):
        self.templates_dir = Path(templates_dir)
        self.templates_dir.mkdir(exist_ok=True)
        self.templates = []
        self._timezone = None
        self._timezone_name = timezone_name
        self._init_timezone()
        self._load_templates()
    
    def _init_timezone(self):
        """Initialize timezone."""
        import pytz
        
        try:
            self._timezone = pytz.timezone(self._timezone_name)
        except Exception as e:
            logger.warning(f"Unknown timezone '{self._timezone_name}', falling back to UTC: {e}")
            self._timezone = pytz.UTC
            self._timezone_name = "UTC"
    
    def _get_current_time(self):
        """Get current time in configured timezone."""
        return datetime.now(self._timezone)
    
    def _load_templates(self):
        """Load templates from JSON files in the templates directory."""
        if not self.templates_dir.exists():
            self._create_default_templates()
        
        for template_file in self.templates_dir.glob("*.json"):
            try:
                with open(template_file, 'r') as f:
                    template_data = json.load(f)
                    if isinstance(template_data, list):
                        self.templates.extend(template_data)
                    else:
                        self.templates.append(template_data)
            except Exception as e:
                logger.error(f"Error loading template {template_file}: {e}")
        
        if not self.templates:
            self._create_default_templates()
    
    def _create_default_templates(self):
        """Create default templates if none exist."""
        default_templates = [
            {
                "name": "Daily Reflection",
                "title": "Daily Reflection",
                "prompt": "What are your thoughts about {date} at {time}?",
                "description": "A simple daily reflection prompt"
            },
            {
                "name": "Mindfulness Check",
                "title": "Mindfulness Check",
                "prompt": "How are you feeling right now on this {day} at {time}?",
                "description": "A mindfulness check-in prompt"
            },
            {
                "name": "Gratitude Prompt",
                "title": "Gratitude Prompt",
                "prompt": "What are three things you're grateful for on {date} during this {time_of_day}?",
                "description": "A gratitude reflection prompt"
            }
        ]
        
        template_file = self.templates_dir / "default.json"
        with open(template_file, 'w') as f:
            json.dump(default_templates, f, indent=2)
        
        self.templates.extend(default_templates)
    
    def _get_time_of_day(self, hour: int) -> str:
        """Return time of day based on hour."""
        if 5 <= hour < 12:
            return "morning"
        elif 12 <= hour < 17:
            return "afternoon"
        elif 17 <= hour < 21:
            return "evening"
        else:
            return "night"
    
    def _fill_placeholders(self, template: str) -> str:
        """Fill placeholders in template with current values."""
        now = self._get_current_time()
        placeholders = {
            "date": now.strftime("%Y-%m-%d"),
            "time": now.strftime("%H:%M"),
            "day": now.strftime("%A"),
            "time_of_day": self._get_time_of_day(now.hour),
            "weekday": now.strftime("%A"),
            "month": now.strftime("%B"),
            "day_of_month": str(now.day)
        }
        
        for placeholder, value in placeholders.items():
            template = template.replace(f"{{{placeholder}}}", value)
        
        return template
    
    def get_random_prompt(self) -> dict:
        """Get a random prompt template with placeholders filled.
        
        Returns:
            dict: Contains 'prompt' and 'title' keys with filled placeholders
        """
        if not self.templates:
            logger.warning("No templates available")
            return {
                "prompt": "What are your thoughts right now?",
                "title": "Random Thought"
            }
        
        template = random.choice(self.templates)
        prompt = template.get("prompt", "")
        title = template.get("title", template.get("name", "Notification"))
        
        return {
            "prompt": self._fill_placeholders(prompt),
            "title": self._fill_placeholders(title)
        }
    
    def add_template(self, name: str, prompt: str, description: str = "", title: str = None):
        """Add a new template to the system."""
        if title is None:
            title = name
        
        new_template = {
            "name": name,
            "title": title,
            "prompt": prompt,
            "description": description
        }
        self.templates.append(new_template)
        
        # Save to file
        template_file = self.templates_dir / "custom.json"
        custom_templates = []
        if template_file.exists():
            with open(template_file, 'r') as f:
                custom_templates = json.load(f)
        
        custom_templates.append(new_template)
        
        with open(template_file, 'w') as f:
            json.dump(custom_templates, f, indent=2)