from pydantic import BaseModel, Field
from typing import Optional, Dict, Any, List, Literal

# Schema definitions for Jasper's internal state management
class Task(BaseModel):
    id: str = Field(..., description="Unique identifier for the task")
    description: str = Field(..., description="Description of the task")
    tool_name: Optional[str] = Field(default=None, description="Name of the tool to be used for the task")
    tool_args: Optional[Dict[str, Any]] = Field(default=None, description="Arguments for the tool")
    status: Literal["pending", "in_progress", "completed", "failed"] = Field(default="pending", description="Current status of the task")
    error: Optional[str] = Field(default=None, description="Error message if the task failed")

# --- Confidence Breakdown ---
# Provides a detailed view of the confidence score components
class ConfidenceBreakdown(BaseModel):
    data_coverage: float      # % of required data fetched
    data_quality: float       # provider reliability
    inference_strength: float # logic depth
    overall: float

class validationresult(BaseModel):
    is_valid: bool = Field(..., description="Indicates if the state is valid")
    issues: List[str] = Field(default_factory=list, description="List of issues found during validation")
    confidence: float = Field(default=0.0, description="Confidence score of the validation result")
    breakdown: Optional[ConfidenceBreakdown] = Field(default=None, description="Detailed confidence breakdown")


class Jasperstate(BaseModel):
    query: str = Field(..., description="The original user query")

    plan: List[Task] = Field(default_factory=list, description="List of tasks in the plan")
    current_task_index: int = Field(default=0, description="Index of the current task being executed")

    task_results: Dict[str, Dict] = Field(default_factory=dict, description="Results of executed tasks, keyed by task ID")

    validation: Optional[validationresult] = Field(default=None, description="Validation result of the current state")

    retries: int = Field(default=0, description="Number of retries attempted for failed tasks")
    max_retries: int = Field(default=3, description="Maximum number of retries allowed for failed tasks")

    status: Literal["Planning", "Executing", "Validating", "Synthesizing", "Completed", "Failed"] = Field(default="Planning", description="Current status of the agent")

    final_answer: Optional[str] = Field(default=None, description="The final answer generated by Jasper")
    error: Optional[str] = Field(default=None, description="Error message if the run failed unexpectedly")
    error_source: Optional[str] = Field(default=None, description="Source of the error: 'data_provider', 'llm_service', 'llm_auth', 'llm_timeout', 'llm_unknown', or 'query'")
