from typing import Dict, List, Optional
from dataclasses import dataclass, field
from datetime import datetime

@dataclass
class ProcessMetrics:
    """Metrics for a running process and its children."""
    cpu_percent: float = 0.0
    memory_rss: int = 0  # Resident Set Size in bytes
    memory_vms: int = 0  # Virtual Memory Size in bytes
    memory_percent: float = 0.0
    num_threads: int = 0
    num_children: int = 0
    children_cpu_percent: float = 0.0
    children_memory_rss: int = 0
    last_updated: str = ""

@dataclass
class TaskStatus:
    """Status of an individual task (script)."""
    script_path: str
    script_name: str
    state: str = "PENDING"  # PENDING, WAITING, RUNNING, DONE, ERROR, CRASHED, TIMED_OUT
    start_time: Optional[str] = None
    end_time: Optional[str] = None
    duration_seconds: float = 0.0
    exit_code: Optional[int] = None
    error_message: Optional[str] = None
    pane_id: Optional[str] = None
    process_metrics: ProcessMetrics = field(default_factory=ProcessMetrics)

@dataclass
class DirectoryStatus:
    """Status of a directory containing tasks."""
    directory_num: int
    directory_path: str
    state: str = "PENDING"  # PENDING, RUNNING, DONE, ERROR, TIMED_OUT
    timeout: int = 600
    start_time: Optional[str] = None
    end_time: Optional[str] = None
    duration_seconds: float = 0.0
    tasks: List[TaskStatus] = field(default_factory=list)
    window_name: Optional[str] = None

@dataclass
class DAGStructure:
    """Complete DAG structure for the task run."""
    run_name: str
    run_id: str
    state: str = "INITIALIZING"  # INITIALIZING, RUNNING, COMPLETED, FAILED
    start_time: str = ""
    end_time: Optional[str] = None
    duration_seconds: float = 0.0
    total_directories: int = 0
    completed_directories: int = 0
    failed_directories: int = 0
    directories: List[DirectoryStatus] = field(default_factory=list)
    last_updated: str = ""

class MetaData:
    """Class to hold metadata for a task directory."""
    def __init__(self, timeout: int = 600):  # Default timeout to 10 minutes (600 seconds)
        self.timeout = timeout
        # Add more attributes here in the future