"""
Task Planning Pydantic Models
"""
from __future__ import annotations

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


class AgentInfo(BaseModel):
    """Information about an agent"""
    id: str
    name: str
    model_id: str
    description: Optional[str] = None

    @field_validator('description', mode='before')
    @classmethod
    def empty_str_to_none(cls, v):
        """Convert empty string to None for optional fields"""
        if v == '':
            return None
        return v

    @field_validator('model_id', mode='before')
    @classmethod
    def default_model(cls, v):
        """Provide default model if empty"""
        if not v or v == '':
            return 'claude-sonnet-4'
        return v


class TeamInfo(BaseModel):
    """Information about a team"""
    id: str
    name: str
    agents: List[Dict] = []
    description: Optional[str] = None

    @field_validator('description', mode='before')
    @classmethod
    def empty_str_to_none(cls, v):
        """Convert empty string to None for optional fields"""
        if v == '':
            return None
        return v


class EnvironmentInfo(BaseModel):
    """Information about an execution environment"""
    id: str
    name: str
    type: Optional[str] = "production"
    status: Optional[str] = "active"


class WorkerQueueInfo(BaseModel):
    """Information about a worker queue"""
    id: str
    name: str
    environment_id: Optional[str] = None
    active_workers: int = 0
    status: Optional[str] = "active"


class TaskPlanRequest(BaseModel):
    """Request to plan a task"""
    description: str = Field(..., description="Task description")
    priority: Literal['low', 'medium', 'high', 'critical'] = Field('medium', description="Task priority")
    project_id: Optional[str] = Field(None, description="Associated project ID")
    agents: List[AgentInfo] = Field([], description="Available agents")
    teams: List[TeamInfo] = Field([], description="Available teams")
    environments: List[EnvironmentInfo] = Field([], description="Available execution environments")
    worker_queues: List[WorkerQueueInfo] = Field([], description="Available worker queues")
    refinement_feedback: Optional[str] = Field(None, description="User feedback for plan refinement")
    conversation_context: Optional[str] = Field(None, description="Conversation history for context")
    previous_plan: Optional[Dict] = Field(None, description="Previous plan for refinement")
    iteration: int = Field(1, description="Planning iteration number")
    planning_strategy: Optional[Literal['claude_code_sdk', 'agno']] = Field('claude_code_sdk', description="Planning strategy to use (claude_code_sdk or agno)")


class ComplexityInfo(BaseModel):
    """Task complexity assessment"""
    story_points: int = Field(..., ge=1, le=21, description="Story points (1-21)")
    confidence: Literal['low', 'medium', 'high'] = Field(..., description="Confidence level")
    reasoning: str = Field(..., description="Reasoning for complexity assessment")


class AgentModelInfo(BaseModel):
    """Information about the model an agent will use"""
    model_id: str  # e.g., "claude-sonnet-4", "gpt-4o"
    estimated_input_tokens: int
    estimated_output_tokens: int
    cost_per_1k_input_tokens: float
    cost_per_1k_output_tokens: float
    total_model_cost: float


class ToolUsageInfo(BaseModel):
    """Expected tool usage for an agent"""
    tool_name: str  # e.g., "aws_s3", "kubectl", "bash"
    estimated_calls: int
    cost_per_call: float
    total_tool_cost: float


class TaskItem(BaseModel):
    """Detailed task breakdown item with dependencies and testing strategy"""
    id: int = Field(..., description="Unique task ID for tracking dependencies")
    title: str = Field(..., description="Short, clear task title")
    description: str = Field(..., description="Brief overview of what needs to be done")
    details: str = Field(..., description="Step-by-step implementation details, code snippets, and specific instructions")
    test_strategy: str = Field(..., description="How to verify this task was completed correctly")
    priority: Literal["high", "medium", "low"] = Field(default="medium", description="Task priority level")
    dependencies: List[int] = Field(default_factory=list, description="List of task IDs that must be completed before this one")
    status: Literal["pending", "in_progress", "done"] = Field(default="pending", description="Current task status")
    subtasks: List["TaskItem"] = Field(default_factory=list, description="Optional nested subtasks")
    # Optional context-driven fields - planner decides what to use based on agent capabilities
    skills_to_use: Optional[List[str]] = Field(default=None, description="Optional: Specific skills from agent's skillset to use (e.g., ['aws_s3', 'kubectl'])")
    env_vars_to_use: Optional[List[str]] = Field(default=None, description="Optional: Environment variables from execution_environment to use (e.g., ['AWS_REGION', 'KUBECONFIG'])")
    secrets_to_use: Optional[List[str]] = Field(default=None, description="Optional: Secrets/credentials from execution_environment to use (e.g., ['AWS_ACCESS_KEY_ID', 'GITHUB_TOKEN'])")
    knowledge_references: Optional[List[str]] = Field(default=None, description="Optional: References to organizational knowledge used for this task")


class TeamBreakdownItem(BaseModel):
    """Breakdown of work for a specific team/agent"""
    team_id: Optional[str] = None
    team_name: str
    agent_id: Optional[str] = None
    agent_name: Optional[str] = None
    responsibilities: List[str]
    estimated_time_hours: float
    model_info: Optional[AgentModelInfo] = None
    expected_tools: List[ToolUsageInfo] = []
    agent_cost: float = 0.0  # Total cost for this agent (model + tools)
    # New: Detailed task breakdown with dependencies
    tasks: List[TaskItem] = Field(default_factory=list, description="Ordered list of tasks with dependencies and test strategies")


class RecommendedExecution(BaseModel):
    """AI recommendation for which entity should execute the task"""
    entity_type: Literal['agent', 'team']
    entity_id: str
    entity_name: str
    reasoning: str
    recommended_environment_id: Optional[str] = None
    recommended_environment_name: Optional[str] = None
    recommended_worker_queue_id: Optional[str] = None
    recommended_worker_queue_name: Optional[str] = None
    execution_reasoning: Optional[str] = None


class LLMCostBreakdown(BaseModel):
    """Detailed LLM cost breakdown by model"""
    model_id: str
    estimated_input_tokens: int
    estimated_output_tokens: int
    cost_per_1k_input_tokens: float
    cost_per_1k_output_tokens: float
    total_cost: float


class ToolCostBreakdown(BaseModel):
    """Tool execution cost breakdown"""
    category: str  # e.g., "AWS APIs", "Database Queries", "External APIs"
    tools: List[ToolUsageInfo]
    category_total: float


class RuntimeCostBreakdown(BaseModel):
    """Runtime and compute costs"""
    worker_execution_hours: float
    cost_per_hour: float
    total_cost: float


class CostBreakdownItem(BaseModel):
    """Individual cost breakdown item (legacy, kept for backwards compatibility)"""
    item: str
    cost: float


class HumanResourceCost(BaseModel):
    """Human resource cost breakdown by role"""
    role: str  # e.g., "Senior DevOps Engineer", "Security Engineer"
    hourly_rate: float  # e.g., 150.00
    estimated_hours: float  # e.g., 8.0
    total_cost: float  # e.g., 1200.00


class CostEstimate(BaseModel):
    """Enhanced cost estimation for the task"""
    estimated_cost_usd: float
    # Legacy breakdown (keep for backwards compatibility)
    breakdown: List[CostBreakdownItem] = []
    # New detailed breakdowns
    llm_costs: List[LLMCostBreakdown] = []
    tool_costs: List[ToolCostBreakdown] = []
    runtime_cost: Optional[RuntimeCostBreakdown] = None


class RealizedSavings(BaseModel):
    """Realized savings by using Kubiya orchestration platform"""
    # Without Kubiya (manual execution)
    without_kubiya_cost: float  # Total cost if done manually
    without_kubiya_hours: float  # Total time if done manually
    without_kubiya_resources: List[HumanResourceCost]  # Resource breakdown

    # With Kubiya (AI orchestration)
    with_kubiya_cost: float  # AI execution cost
    with_kubiya_hours: float  # AI execution time

    # Realized Savings
    money_saved: float  # Dollars saved
    time_saved_hours: float  # Hours saved
    time_saved_percentage: int  # Percentage of time saved

    # Summary
    savings_summary: str  # Compelling savings narrative


class TaskPlanResponse(BaseModel):
    """AI-generated task plan"""
    title: str
    summary: str
    complexity: ComplexityInfo
    team_breakdown: List[TeamBreakdownItem]
    recommended_execution: RecommendedExecution
    cost_estimate: CostEstimate
    realized_savings: RealizedSavings
    risks: List[str] = []
    prerequisites: List[str] = []
    success_criteria: List[str] = []
    # Optional fields for when AI needs clarification
    has_questions: bool = False
    questions: Optional[str] = None


# Rebuild models to support forward references (for TaskItem.subtasks)
TaskItem.model_rebuild()
