from datetime import datetime
from typing import Any, Dict, List, Literal, Optional

from pydantic import BaseModel, ConfigDict, Field
from pydantic.alias_generators import to_camel

from persona_models.common import create_uuid

MissionStatus = Literal[
    "draft", "planned", "active", "review", "paused", "completed", "failed", "cancelled"
]
MissionStepStatus = Literal["pending", "in_progress", "completed", "failed", "skipped"]


class WorkerAgent(BaseModel):
    agent_id: str
    scope: str
    name: Optional[str] = None

    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)


class MissionStep(BaseModel):
    id: str = Field(default_factory=create_uuid)
    name: str
    description: str | None = None
    status: MissionStepStatus = "pending"
    worker_agent_id: str | None = None
    parent_step_id: str | None = None
    dependencies: List[str] = Field(default_factory=list)
    output: str | None = None
    error: str | None = None
    metadata: Dict[str, Any] = Field(default_factory=dict)
    created_at: datetime = Field(default_factory=datetime.now)
    updated_at: datetime = Field(default_factory=datetime.now)
    started_at: datetime | None = None
    completed_at: datetime | None = None

    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)


class Mission(BaseModel):
    id: str = Field(default_factory=create_uuid)
    project_id: str
    name: str
    goal: str
    status: MissionStatus = "draft"
    planner_agent_id: str
    worker_agents: List[WorkerAgent] = Field(default_factory=list)
    steps: List[MissionStep] = Field(default_factory=list)
    current_step_id: str | None = None
    version: int = 1
    user_instructions: List[str] = Field(default_factory=list)
    created_at: datetime = Field(default_factory=datetime.now)
    updated_at: datetime = Field(default_factory=datetime.now)

    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)

    def get_step(self, step_id: str) -> Optional[MissionStep]:
        for step in self.steps:
            if step.id == step_id or step.name == step_id:
                return step
        return None

    def get_root_steps(self) -> List[MissionStep]:
        return [s for s in self.steps if s.parent_step_id is None]

    def get_ready_steps(self) -> List[MissionStep]:
        completed_ids = {s.id for s in self.steps if s.status == "completed"}
        return [
            s
            for s in self.steps
            if s.status == "pending"
            and all(dep in completed_ids for dep in s.dependencies)
        ]


class CreateMissionRequest(BaseModel):
    name: str
    goal: str
    planner_agent_id: str
    worker_agents: List[WorkerAgent] = Field(default_factory=list)

    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)
