"""
Shared Data Models for Task Manager MCP Servers

This module defines the core data models used across both MCP servers.
These models ensure consistency and type safety throughout the system.
"""

from typing import Dict, List, Optional, Any, Literal, Union
from datetime import datetime
from pydantic import BaseModel, Field
from enum import Enum


class TaskStatus(str, Enum):
    """Task status enumeration"""
    TODO = "todo"
    IN_PROGRESS = "in_progress"
    BLOCKED = "blocked"
    COMPLETED = "completed"
    CANCELLED = "cancelled"


class TaskPriority(int, Enum):
    """Task priority levels (1=highest, 5=lowest)"""
    CRITICAL = 1
    HIGH = 2
    MEDIUM = 3
    LOW = 4
    MINIMAL = 5


class ResourceType(str, Enum):
    """Resource type enumeration"""
    HUMAN = "human"
    AGENT = "agent"


class Task(BaseModel):
    """Core task model"""
    id: str = Field(..., description="Unique task identifier")
    name: str = Field(..., description="Task name")
    description: str = Field(default="", description="Detailed task description")
    status: TaskStatus = Field(default=TaskStatus.TODO, description="Current task status")
    priority: TaskPriority = Field(default=TaskPriority.MEDIUM, description="Task priority level")
    assignee: Optional[str] = Field(default=None, description="Assigned resource ID")
    tags: List[str] = Field(default_factory=list, description="Task tags")
    due_date: Optional[datetime] = Field(default=None, description="Task due date")
    time_cost: float = Field(default=1.0, description="Estimated time cost in hours")
    importance: int = Field(default=3, ge=1, le=5, description="Task importance (1-5)")
    benefit: int = Field(default=5, ge=1, le=10, description="Expected benefit (1-10)")
    marginal_cost: int = Field(default=3, ge=1, le=5, description="Marginal cost (1-5)")
    dependencies: List[str] = Field(default_factory=list, description="List of dependent task IDs")
    subtasks: List[str] = Field(default_factory=list, description="List of subtask IDs")
    created_at: datetime = Field(default_factory=datetime.now, description="Creation timestamp")
    updated_at: datetime = Field(default_factory=datetime.now, description="Last update timestamp")
    project_id: Optional[str] = Field(default=None, description="Associated project ID")
    
    # Risk and progress tracking fields
    risk_level: int = Field(default=3, ge=1, le=5, description="Risk level (1-5)")
    progress: float = Field(default=0.0, ge=0.0, le=1.0, description="Progress percentage (0-1)")
    estimated_completion: Optional[datetime] = Field(default=None, description="Estimated completion date")
    actual_time_spent: float = Field(default=0.0, description="Actual time spent in hours")


class Project(BaseModel):
    """Project model"""
    id: str = Field(..., description="Unique project identifier")
    name: str = Field(..., description="Project name")
    description: str = Field(default="", description="Project description")
    requirements_doc: str = Field(default="", description="Project requirements document")
    created_at: datetime = Field(default_factory=datetime.now, description="Creation timestamp")
    updated_at: datetime = Field(default_factory=datetime.now, description="Last update timestamp")
    status: str = Field(default="active", description="Project status")


class HumanResource(BaseModel):
    """Human resource model"""
    id: str = Field(..., description="Unique resource identifier")
    name: str = Field(..., description="Resource name")
    type: Literal[ResourceType.HUMAN] = Field(default=ResourceType.HUMAN, description="Resource type")
    skills: List[str] = Field(default_factory=list, description="List of skills")
    load: float = Field(default=0.0, ge=0.0, le=1.0, description="Current workload (0-1)")
    cost_per_hour: float = Field(default=50.0, description="Cost per hour")
    status: str = Field(default="available", description="Current status")
    availability: Dict[str, Any] = Field(default_factory=dict, description="Availability schedule")


class AgentResource(BaseModel):
    """AI Agent resource model"""
    id: str = Field(..., description="Unique resource identifier")
    name: str = Field(..., description="Resource name")
    type: Literal[ResourceType.AGENT] = Field(default=ResourceType.AGENT, description="Resource type")
    skills: List[str] = Field(default_factory=list, description="List of skills")
    load: float = Field(default=0.0, ge=0.0, description="Current task load")
    cost_per_task: float = Field(default=0.1, description="Cost per task")
    system_prompt: str = Field(default="", description="Agent system prompt")
    tools: List[Dict[str, str]] = Field(default_factory=list, description="Available tools")
    status: str = Field(default="available", description="Current status")


# Union type for resources
Resource = Union[HumanResource, AgentResource]


class RiskAssessment(BaseModel):
    """Risk assessment model"""
    task_id: str = Field(..., description="Associated task ID")
    risk_level: int = Field(..., ge=1, le=5, description="Risk level (1-5)")
    risk_factors: List[str] = Field(default_factory=list, description="Identified risk factors")
    mitigation_strategies: List[str] = Field(default_factory=list, description="Suggested mitigation strategies")
    probability: float = Field(default=0.5, ge=0.0, le=1.0, description="Risk probability")
    impact: int = Field(default=3, ge=1, le=5, description="Risk impact level")
    created_at: datetime = Field(default_factory=datetime.now, description="Assessment timestamp")


class ProgressUpdate(BaseModel):
    """Progress update model"""
    task_id: str = Field(..., description="Associated task ID")
    progress: float = Field(..., ge=0.0, le=1.0, description="Progress percentage")
    time_spent: float = Field(default=0.0, description="Time spent since last update")
    notes: str = Field(default="", description="Progress notes")
    updated_by: str = Field(..., description="Who updated the progress")
    timestamp: datetime = Field(default_factory=datetime.now, description="Update timestamp")


class TaskRecommendation(BaseModel):
    """Task recommendation model"""
    task_id: str = Field(..., description="Recommended task ID")
    score: float = Field(..., description="Recommendation score")
    reasoning: str = Field(..., description="Reasoning for recommendation")
    context: Dict[str, Any] = Field(default_factory=dict, description="Additional context")
    timestamp: datetime = Field(default_factory=datetime.now, description="Recommendation timestamp")


class Reminder(BaseModel):
    """Reminder model"""
    id: str = Field(..., description="Unique reminder identifier")
    task_id: str = Field(..., description="Associated task ID")
    reminder_type: str = Field(..., description="Type of reminder (start, deadline, progress)")
    message: str = Field(..., description="Reminder message")
    scheduled_time: datetime = Field(..., description="When to send the reminder")
    sent: bool = Field(default=False, description="Whether reminder was sent")
    created_at: datetime = Field(default_factory=datetime.now, description="Creation timestamp")