src/state.py

from typing import Optional, List, Dict, Any, Union
from pydantic import BaseModel
import os
import re
from src.tools.assets_utils import sanitize_filename


class CodeSegment(BaseModel):
    code_snippet: str
    explanation: str
    audio_path: Optional[str] = None
    audio_duration: Optional[float] = None
    start_char_idx: Optional[int] = None
    end_char_idx: Optional[int] = None


class VideoGenerationState(BaseModel):
    # User-provided inputs
    topic: Optional[str] = None
    course_name: Optional[str] = None
    subtitle: Optional[str] = None
    programming_language: Optional[str] = None
    slide_colour: Optional[str] = "blue"
    # Allow target_language to be either a single language string or list of languages
    target_language: Optional[Union[str, List[str]]] = None
    tts_gender: Optional[str] = "male"
    visualization_plan: Optional[Dict[str, Any]] = None
    tts_voice: Optional[str] = "kore"
    toggle_hinglish: Optional[bool] = False
    video_type: Optional[str] = "base_video"  # or "personalised_video"
    user_code_input: Optional[str] = None

    # Memory & Context
    session_id: Optional[str] = None
    user_profile: Optional[Dict[str, Any]] = None  # pulled from DB or vector store

    # Intermediate content
    presentation_txt: Optional[str] = None
    slides: Optional[List[str]] = None
    rendered_slides: Optional[List[str]] = None  # images or paths
    slide_contents: Optional[List[Dict[str, Any]]] = None
    code_blocks: Optional[List[str]] = None
    segmented_code_data: Optional[List[CodeSegment]] = None
    # Video paths - ALL path fields must be Dict[str, str] keyed by language
    web_output_screenshot_path: Optional[str] = None
    rendered_web_videos: Optional[List[str]] = None
    slide_videos: Optional[Dict[str, str]] = None
    visualization_videos: Optional[List[str]] = None  # Added this field
    slide_video_path: Optional[Dict[str, str]] = None
    final_avatar_video: Optional[Dict[str, str]] = None
    web_code_file_paths: Optional[List[str]] = None
    final_avatar_video_path: Optional[Dict[str, str]] = None
    code_video_path: Optional[Dict[str, str]] = None

    ppt_viz_path: Optional[str] = None
    final_video_path: Optional[Dict[str, str]] = None

    master_audio_path: Optional[str] = None
    # Multi-language & audio flags
    generate_multi_language: Optional[bool] = False
    languages: Optional[List[str]] = None
    upload_audio_to_s3: Optional[bool] = False
    multi_audio_mode: Optional[str] = "per_language_videos"  # or 'single_video_multi_audio_tracks'
    audio_s3_urls: Optional[Dict[str, str]] = None
    temp_frame_dir: Optional[str] = None
    actual_execution_output: Optional[str] = None
    upload_link: Optional[Union[str, Dict[str, Optional[str]]]] = None
    # Error handling
    error: Optional[str] = None
    subtitle_content: Optional[Union[str, Dict[str, Optional[str]]]] = None
    

    # Additional fields for input node
    http_headers: Optional[Dict[str, Any]] = None
    optional_params: Optional[Dict[str, Any]] = None
    context_metadata: Optional[Dict[str, Any]] = None
    memory_retrieval_status: Optional[str] = None
    fallback_indicator: Optional[bool] = None
    status: str = "pending"
    output_type: str = "terminal"

    # visualisation options
    enable_maths: Optional[bool] = False
    enable_highlighting: Optional[bool] = False
    enable_images: Optional[bool] = True
    enable_visualizations: Optional[bool] = True

    def get_base_dirs(self):
        """Get base directories for outputs."""

        # Create topic-specific directory (sanitize topic for filesystem safety)
        current_dir = os.path.dirname(os.path.abspath(__file__))
        safe_topic = sanitize_filename(self.topic)
        base_dir = os.path.join(current_dir, "output", safe_topic, "visualisation")
        os.makedirs(base_dir, exist_ok=True)



        return {
            "base_dir": base_dir,
            "web_code_dir": os.path.join(base_dir, "web_code"),
            "media_dir": os.path.join(base_dir, "media"),
            "audio_dir": os.path.join(base_dir, "audio"),
            "oration_dir": os.path.join(base_dir, "oration"),
            "combined_dir": os.path.join(base_dir, "combined"),
            "temp_dir": os.path.join(base_dir, "temp")
        }


    def _safe_title(self, title):
        """Convert title to safe filename."""
        safe = re.sub(r'[^\w\s-]', '', str(title)).strip().replace(' ', '_').lower()
        return re.sub(r'_+', '_', safe)


    def set_error(self, message: str):
        """Set error message and update status."""
        self.error = message
        self.status = 'failed'