"""
Task — the core work unit in Nexus.

Design notes:
  - Typed input/output via Pydantic schemas
  - Lifecycle mirrors A2A Task states: submitted → working → completed/failed
  - Constraints (budget, latency, quality) drive adaptive topology decisions
  - Each task records its full execution trace (messages, token usage, timing)
"""

from __future__ import annotations

from datetime import datetime
from enum import StrEnum
from typing import Any

from pydantic import BaseModel, Field

from nexus.core.message import Conversation, Message
from nexus.types import (
    AgentCapability,
    ContentPart,
    TaskConstraints,
    TokenUsage,
    new_id,
)


class TaskStatus(StrEnum):
    """
    Task lifecycle states, compatible with A2A protocol.

    State machine:
        SUBMITTED → WORKING → COMPLETED
                           → FAILED
                           → CANCELED
                  → INPUT_REQUIRED (paused, needs human input)
                  → REJECTED (agent refused the task)
    """

    SUBMITTED = "submitted"
    WORKING = "working"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELED = "canceled"
    INPUT_REQUIRED = "input_required"
    REJECTED = "rejected"

    @property
    def is_terminal(self) -> bool:
        return self in (
            TaskStatus.COMPLETED,
            TaskStatus.FAILED,
            TaskStatus.CANCELED,
            TaskStatus.REJECTED,
        )

    @property
    def is_active(self) -> bool:
        return self in (TaskStatus.WORKING, TaskStatus.INPUT_REQUIRED)


class TaskResult(BaseModel):
    """The output produced by completing a task."""

    content: str = ""
    parts: list[ContentPart] = Field(default_factory=list)
    data: dict[str, Any] = Field(default_factory=dict)
    quality_score: float | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)


class Task(BaseModel):
    """
    A unit of work to be executed by one or more agents.

    Tasks are the fundamental building block of workflows. They carry:
    - A description of what needs to be done
    - Typed input/output
    - Constraints that guide agent selection and execution
    - Full execution history for observability
    """

    id: str = Field(default_factory=new_id)
    name: str = ""
    description: str = ""
    task_type: str = "general"

    # Input
    input_data: dict[str, Any] = Field(default_factory=dict)
    input_messages: list[Message] = Field(default_factory=list)

    # Constraints for adaptive topology
    constraints: TaskConstraints = Field(default_factory=TaskConstraints)
    required_capabilities: AgentCapability | None = None

    # Execution state
    status: TaskStatus = TaskStatus.SUBMITTED
    assigned_agent_id: str | None = None
    result: TaskResult | None = None
    error: str | None = None

    # Execution trace
    conversation: Conversation = Field(default_factory=Conversation)
    token_usage: TokenUsage = Field(default_factory=TokenUsage)

    # Dependencies (for TaskGraph)
    depends_on: list[str] = Field(default_factory=list)  # Task IDs

    # Timing
    created_at: datetime = Field(default_factory=datetime.utcnow)
    started_at: datetime | None = None
    completed_at: datetime | None = None

    # Metadata
    metadata: dict[str, Any] = Field(default_factory=dict)
    retry_count: int = 0
    max_retries: int = 3

    # ---- Lifecycle methods ----

    def reset_for_retry(self) -> None:
        """Reset a failed task so it can be retried."""
        if self.status == TaskStatus.FAILED:
            self.status = TaskStatus.SUBMITTED
            self.error = None
            self.completed_at = None
            self.retry_count += 1

    def start(self, agent_id: str) -> None:
        """Mark the task as being worked on by an agent."""
        if self.status == TaskStatus.FAILED:
            self.reset_for_retry()
        if self.status.is_terminal:
            raise ValueError(f"Cannot start task in terminal state: {self.status}")
        self.status = TaskStatus.WORKING
        self.assigned_agent_id = agent_id
        self.started_at = datetime.utcnow()

    def complete(self, result: TaskResult) -> None:
        """Mark the task as successfully completed."""
        self.status = TaskStatus.COMPLETED
        self.result = result
        self.completed_at = datetime.utcnow()

    def fail(self, error: str) -> None:
        """Mark the task as failed."""
        self.status = TaskStatus.FAILED
        self.error = error
        self.completed_at = datetime.utcnow()

    def cancel(self) -> None:
        """Cancel the task."""
        if self.status.is_terminal:
            raise ValueError(f"Cannot cancel task in terminal state: {self.status}")
        self.status = TaskStatus.CANCELED
        self.completed_at = datetime.utcnow()

    def request_input(self, prompt: str) -> None:
        """Pause the task and request human input."""
        self.status = TaskStatus.INPUT_REQUIRED
        self.conversation.append(
            Message.agent(f"[INPUT_REQUIRED] {prompt}", sender_id=self.assigned_agent_id)
        )

    def can_retry(self) -> bool:
        """Check if the task can be retried."""
        return self.retry_count < self.max_retries

    @property
    def elapsed_seconds(self) -> float | None:
        """Wall-clock time from start to completion."""
        if not self.started_at:
            return None
        end = self.completed_at or datetime.utcnow()
        return (end - self.started_at).total_seconds()

    @property
    def is_ready(self) -> bool:
        """Check if all dependencies are satisfied (for TaskGraph scheduling)."""
        # This is checked externally by the orchestrator using the task registry
        return self.status == TaskStatus.SUBMITTED


class TaskGraph(BaseModel):
    """
    A directed acyclic graph of tasks with dependency edges.

    Used by the TopologyEngine to determine execution order and parallelism.
    """

    id: str = Field(default_factory=new_id)
    tasks: dict[str, Task] = Field(default_factory=dict)
    metadata: dict[str, Any] = Field(default_factory=dict)

    def add_task(self, task: Task) -> None:
        """Add a task to the graph."""
        self.tasks[task.id] = task

    def add_dependency(self, task_id: str, depends_on_id: str) -> None:
        """Declare that task_id depends on depends_on_id."""
        if task_id not in self.tasks:
            raise KeyError(f"Task {task_id} not found")
        if depends_on_id not in self.tasks:
            raise KeyError(f"Task {depends_on_id} not found")
        if depends_on_id not in self.tasks[task_id].depends_on:
            self.tasks[task_id].depends_on.append(depends_on_id)

    def get_ready_tasks(self) -> list[Task]:
        """Return tasks whose dependencies are all completed (can run in parallel)."""
        ready = []
        for task in self.tasks.values():
            if task.status != TaskStatus.SUBMITTED:
                continue
            deps_satisfied = all(
                self.tasks[dep_id].status == TaskStatus.COMPLETED
                for dep_id in task.depends_on
                if dep_id in self.tasks
            )
            if deps_satisfied:
                ready.append(task)
        return ready

    def topological_sort(self) -> list[Task]:
        """Return tasks in a valid execution order."""
        visited: set[str] = set()
        result: list[Task] = []

        def visit(task_id: str) -> None:
            if task_id in visited:
                return
            visited.add(task_id)
            task = self.tasks[task_id]
            for dep_id in task.depends_on:
                if dep_id in self.tasks:
                    visit(dep_id)
            result.append(task)

        for task_id in self.tasks:
            visit(task_id)
        return result

    @property
    def is_complete(self) -> bool:
        """Check if all tasks are in a terminal state."""
        return all(t.status.is_terminal for t in self.tasks.values())

    @property
    def has_failures(self) -> bool:
        return any(t.status == TaskStatus.FAILED for t in self.tasks.values())
