"""
A2A Task Handler
-----------------
Handler for A2A protocol task requests.

Uses the thin client pattern - all verification happens server-side via /v1/stamp.

The A2A protocol defines a task lifecycle with states:
- pending: Task received but not started
- running: Task in progress
- completed: Task finished successfully
- failed: Task failed
- cancelled: Task was cancelled
"""

from dataclasses import dataclass, field
from typing import Dict, Any, Optional, Callable, Awaitable
from enum import Enum
import time
import uuid

from lockstock_core import LockStockClient
from lockstock_core.types import map_tool_to_capability


class A2ATaskState(str, Enum):
    """A2A task lifecycle states."""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


@dataclass
class A2ATask:
    """An A2A protocol task."""
    task_id: str
    agent_id: str
    capability: str
    input_data: Dict[str, Any]
    state: A2ATaskState = A2ATaskState.PENDING
    output_data: Optional[Dict[str, Any]] = None
    error: Optional[str] = None
    created_at: float = field(default_factory=time.time)
    updated_at: float = field(default_factory=time.time)
    lockstock_hash: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        return {
            "task_id": self.task_id,
            "agent_id": self.agent_id,
            "capability": self.capability,
            "state": self.state.value,
            "input_data": self.input_data,
            "output_data": self.output_data,
            "error": self.error,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
            "lockstock_verification": {
                "hash": self.lockstock_hash
            } if self.lockstock_hash else None
        }


class A2ATaskHandler:
    """
    Handler for A2A protocol task requests.

    Manages the task lifecycle while integrating with LockStock
    for authorization and audit.
    """

    # Type for task executor functions
    TaskExecutor = Callable[[str, Dict[str, Any]], Awaitable[Dict[str, Any]]]

    def __init__(
        self,
        agent_id: str,
        api_key: str,
        endpoint: str = "https://lockstock-api-i9kp.onrender.com"
    ):
        """
        Initialize A2A task handler.

        Args:
            agent_id: The agent's unique identifier
            api_key: API key for authentication (from dashboard)
            endpoint: LockStock API endpoint
        """
        self.client = LockStockClient(
            agent_id=agent_id,
            api_key=api_key,
            endpoint=endpoint
        )
        self.agent_id = agent_id
        self._tasks: Dict[str, A2ATask] = {}
        self._executors: Dict[str, "A2ATaskHandler.TaskExecutor"] = {}

    def register_executor(
        self,
        capability: str,
        executor: "A2ATaskHandler.TaskExecutor"
    ):
        """
        Register an executor for a capability.

        Args:
            capability: The capability this executor handles
            executor: Async function that executes tasks
        """
        self._executors[capability.upper()] = executor

    async def create_task(
        self,
        capability: str,
        input_data: Dict[str, Any],
        task_id: Optional[str] = None
    ) -> A2ATask:
        """
        Create a new A2A task.

        Args:
            capability: Required capability
            input_data: Task input data
            task_id: Optional task ID (generated if not provided)

        Returns:
            Created task
        """
        task_id = task_id or f"task_{uuid.uuid4().hex[:16]}"

        task = A2ATask(
            task_id=task_id,
            agent_id=self.agent_id,
            capability=capability.upper(),
            input_data=input_data
        )

        self._tasks[task_id] = task

        # Stamp task creation (server handles audit logging)
        await self.client.stamp("a2a_task_create")

        return task

    async def execute_task(self, task_id: str) -> A2ATask:
        """
        Execute a pending task.

        Args:
            task_id: Task to execute

        Returns:
            Updated task with result
        """
        task = self._tasks.get(task_id)
        if not task:
            raise ValueError(f"Task not found: {task_id}")

        if task.state != A2ATaskState.PENDING:
            raise ValueError(f"Task is not pending: {task.state}")

        # Stamp capability with LockStock
        result = await self.client.stamp(task=task.capability)

        if not result.authorized:
            task.state = A2ATaskState.FAILED
            task.error = f"LockStock authorization denied: {result.reason}"
            task.updated_at = time.time()
            return task

        task.lockstock_hash = result.state_hash

        # Find executor
        executor = self._executors.get(task.capability)
        if not executor:
            task.state = A2ATaskState.FAILED
            task.error = f"No executor registered for capability: {task.capability}"
            task.updated_at = time.time()
            return task

        # Execute
        task.state = A2ATaskState.RUNNING
        task.updated_at = time.time()

        try:
            output = await executor(task.capability, task.input_data)
            task.state = A2ATaskState.COMPLETED
            task.output_data = output

        except Exception as e:
            task.state = A2ATaskState.FAILED
            task.error = str(e)

        task.updated_at = time.time()
        return task

    async def get_task_status(self, task_id: str) -> Optional[A2ATask]:
        """
        Get task status.

        Args:
            task_id: Task identifier

        Returns:
            Task if found, None otherwise
        """
        return self._tasks.get(task_id)

    async def cancel_task(self, task_id: str) -> Optional[A2ATask]:
        """
        Cancel a pending or running task.

        Args:
            task_id: Task to cancel

        Returns:
            Updated task if found
        """
        task = self._tasks.get(task_id)
        if not task:
            return None

        if task.state in [A2ATaskState.COMPLETED, A2ATaskState.FAILED]:
            return task  # Already finished

        task.state = A2ATaskState.CANCELLED
        task.updated_at = time.time()

        # Stamp cancellation (server handles audit logging)
        await self.client.stamp("a2a_task_cancel")

        return task

    async def list_tasks(
        self,
        state: Optional[A2ATaskState] = None,
        limit: int = 100
    ) -> list:
        """
        List tasks.

        Args:
            state: Filter by state (optional)
            limit: Maximum tasks to return

        Returns:
            List of task summaries
        """
        tasks = list(self._tasks.values())

        if state:
            tasks = [t for t in tasks if t.state == state]

        # Sort by created_at descending
        tasks.sort(key=lambda t: t.created_at, reverse=True)

        return [t.to_dict() for t in tasks[:limit]]

    async def close(self):
        """Close the underlying client."""
        await self.client.close()
