# content_state.py

from typing import List, Set, Optional, Dict, Any
from datetime import datetime

class ContentState:
    """Tracks the state and progression of content generation"""
    def __init__(self):
        """Initialize content state tracking variables"""
        # Core content tracking
        self.current_summary: str = ""
        self.section_outlines: List[str] = []
        self.generated_sections: List[str] = []
        self.narrative_threads: List[str] = []
        self.key_points_covered: Set[str] = set()
        self.transition_points: List[str] = []
        self.manual_content: Optional[str] = None
        
        # Additional prompt and context management
        self.additional_prompts: Dict[str, str] = {}  # Store additional prompts per section
        self.context_history: List[Dict[str, Any]] = []  # Track generation context
        self.custom_instructions: Dict[str, str] = {}  # Store custom instructions per section
        
        # Generation state tracking
        self.total_words_generated: int = 0
        self.generation_timestamps: List[float] = []
        self.current_section_index: int = 0
        self.content_status: str = "not_started"  # not_started, in_progress, completed
        self.manual_edits: List[str] = []
        
    def set_manual_content(self, content: str) -> None:
        """Set pre-written content for this section"""
        self.manual_content = content
        self.manual_edits.append(content)
        self.content_status = "in_progress"
        
    def get_manual_content(self) -> Optional[str]:
        """Get pre-written content if it exists"""
        return self.manual_content
    
    def add_section(self, content: str) -> None:
        """Add a new generated section and update tracking"""
        self.generated_sections.append(content)
        self.total_words_generated += len(content.split())
        self.current_section_index += 1
        self.generation_timestamps.append(datetime.now().timestamp())
        
    def add_narrative_thread(self, thread: str) -> None:
        """Add a new narrative thread to track"""
        if thread not in self.narrative_threads:
            self.narrative_threads.append(thread)
        
    def add_key_point(self, point: str) -> None:
        """Add a key point to the covered points set"""
        self.key_points_covered.add(point)
        
    def add_transition(self, transition: str) -> None:
        """Add a transition point between sections"""
        self.transition_points.append(transition)
        
    def add_custom_prompt(self, section_id: str, prompt: str) -> None:
        """Add a custom prompt for a specific section"""
        self.additional_prompts[section_id] = prompt
        self._update_generation_status()
        
    def add_custom_instructions(self, section_id: str, instructions: str) -> None:
        """Add custom instructions for a specific section"""
        self.custom_instructions[section_id] = instructions
        self._update_generation_status()
        
    def get_section_context(self, section_id: str) -> Dict[str, Any]:
        """Get complete context for a section including custom prompts and instructions"""
        context = {
            'manual_content': self.manual_content,
            'additional_prompt': self.additional_prompts.get(section_id, ''),
            'custom_instructions': self.custom_instructions.get(section_id, ''),
            'current_summary': self.current_summary,
            'narrative_threads': self.narrative_threads,
            'key_points_covered': list(self.key_points_covered),
            'generation_status': self.content_status,
            'total_words': self.total_words_generated,
            'section_count': len(self.generated_sections)
        }
        
        if self.context_history:
            context['previous_context'] = self.context_history[-1]
        
        return context
        
    def update_context_history(self, section_id: str, generation_context: Dict[str, Any]) -> None:
        """Update context history with latest generation context"""
        self.context_history.append({
            'section_id': section_id,
            'context': generation_context,
            'timestamp': datetime.now().isoformat()
        })
        self._update_generation_status()
        
    def get_generation_progress(self) -> float:
        """Calculate generation progress as a percentage"""
        if not self.section_outlines:
            return 0.0
        return (self.current_section_index / len(self.section_outlines)) * 100
    
    def _update_generation_status(self) -> None:
        """Update the content generation status"""
        if self.total_words_generated == 0 and not self.manual_content:
            self.content_status = "not_started"
        elif self.get_generation_progress() >= 100:
            self.content_status = "completed"
        else:
            self.content_status = "in_progress"
    
    def reset_state(self) -> None:
        """Reset the content state to initial values"""
        self.__init__()
        
    def get_state_summary(self) -> Dict[str, Any]:
        """Get a summary of the current content state"""
        return {
            'status': self.content_status,
            'progress': self.get_generation_progress(),
            'total_words': self.total_words_generated,
            'sections_generated': len(self.generated_sections),
            'has_manual_content': bool(self.manual_content),
            'narrative_threads_count': len(self.narrative_threads),
            'key_points_count': len(self.key_points_covered),
            'custom_prompts_count': len(self.additional_prompts),
            'last_updated': self.generation_timestamps[-1] if self.generation_timestamps else None
        }