"""
Sandai Operator SDK - Type definitions

Core types for building batch processing operators.
"""

from typing import Optional, List, Dict, Generic, TypeVar, Generator, Any
from pydantic import BaseModel
from pathlib import Path


class InputArtifact(BaseModel):
    """Input artifact definition"""
    uri: str
    slot: Optional[str] = None
    name: Optional[str] = None
    meta: Optional[Dict[str, object]] = None


class OutputArtifact(BaseModel):
    """Output artifact definition"""
    uri: str
    slot: Optional[str] = None
    name: Optional[str] = None
    meta: Optional[Dict[str, object]] = None


OptionsType = TypeVar("OptionsType", bound=BaseModel)
ResultsType = TypeVar("ResultsType", bound=BaseModel)


class ProcessingContext:
    """
    Processing context for task-level workdir management
    
    Provides isolated working directories for each task to store temporary files
    during processing.
    """
    
    def __init__(self, temp_dir: Path):
        self._temp_dir = Path(temp_dir)
        self._task_workdirs: Dict[str, Path] = {}
    
    def get_task_workdir(self, task_id: str) -> str:
        """
        Get the working directory path for a task
        
        Args:
            task_id: Unique task identifier
            
        Returns:
            Absolute path to the task's working directory
        """
        if task_id not in self._task_workdirs:
            workdir = self._temp_dir / task_id / "workdir"
            workdir.mkdir(parents=True, exist_ok=True)
            self._task_workdirs[task_id] = workdir
        
        return str(self._task_workdirs[task_id])
    



class TaskInput(BaseModel, Generic[OptionsType]):
    """Task input definition"""
    task_id: str
    artifacts: Optional[List[InputArtifact]]
    options: OptionsType
    meta: Optional[Dict] = None
    deadline: Optional[float] = None  # UTC timestamp when task should be abandoned
    cancel_key: Optional[str] = None  # Redis key to check for task cancellation


class TaskError(BaseModel):
    """Task error information"""
    error_type: str  # Error type
    error_message: str  # Error message
    error_details: Optional[Dict[str, Any]] = None  # Error details
    traceback: Optional[str] = None  # Error traceback


class TaskOutput(BaseModel, Generic[ResultsType]):
    """Task output definition"""
    task_id: str
    artifacts: Optional[List[OutputArtifact]] = None
    results: Optional[ResultsType] = None
    error: Optional[TaskError] = None  # Error information
    status: str = "success"  # Task status: success, error, partial


class EOFMarker(BaseModel):
    """EOF marker for job mode to indicate end of input stream"""
    marker_type: str = "eof"
    message: str = "End of input stream"


# Rebuild models to resolve forward references
TaskError.model_rebuild()
TaskOutput.model_rebuild()