"""Asynchronous resource implementation for long-running operations."""

import asyncio
import enum
import time
import uuid
from typing import Any, Callable, Dict, Optional

import pydantic
from pydantic import AnyUrl, Field

from mcp.server.fastmcp.resources.base import Resource
from mcp.server.fastmcp.utilities.logging import get_logger

logger = get_logger(__name__)


class AsyncStatus(str, enum.Enum):
    """Status of an asynchronous operation."""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELED = "canceled"


class AsyncResource(Resource):
    """A resource representing an asynchronous operation.
    
    This resource type is used to track long-running operations that are executed
    asynchronously. It provides methods for updating progress, completing with a result,
    failing with an error, and canceling the operation.
    """
    
    status: AsyncStatus = Field(
        default=AsyncStatus.PENDING,
        description="Current status of the asynchronous operation"
    )
    progress: float = Field(
        default=0,
        description="Current progress value (0-100 or raw count)"
    )
    total: Optional[float] = Field(
        default=None,
        description="Total expected progress value, if known"
    )
    error: Optional[str] = Field(
        default=None,
        description="Error message if the operation failed"
    )
    created_at: float = Field(
        default_factory=time.time,
        description="Timestamp when the resource was created"
    )
    started_at: Optional[float] = Field(
        default=None,
        description="Timestamp when the operation started running"
    )
    completed_at: Optional[float] = Field(
        default=None,
        description="Timestamp when the operation completed, failed, or was canceled"
    )
    
    # Fields not included in serialization
    _task: Optional[asyncio.Task] = pydantic.PrivateAttr(default=None)
    _mcp_server = pydantic.PrivateAttr(default=None)
    
    async def read(self) -> str:
        """Read the current state of the resource as JSON.
        
        Returns the current status and progress information.
        """
        # Convert the resource to a dictionary, excluding private fields
        data = self.model_dump(exclude={"_task", "_mcp_server"})
        
        # Return status info as JSON
        import json
        return json.dumps(data, indent=2)
    
    async def start(self, task: asyncio.Task) -> None:
        """Mark the resource as running and store the task.
        
        Args:
            task: The asyncio task that is executing the operation
        """
        self._task = task
        self.status = AsyncStatus.RUNNING
        self.started_at = time.time()
        await self._notify_changed()
        
        logger.debug(
            "Started async operation",
            extra={
                "uri": self.uri,
            }
        )
    
    async def update_progress(self, progress: float, total: Optional[float] = None) -> None:
        """Update the progress information.
        
        Args:
            progress: Current progress value
            total: Total expected progress value, if known
        """
        self.progress = progress
        if total is not None:
            self.total = total
        await self._notify_changed()
        
        logger.debug(
            "Updated async operation progress",
            extra={
                "uri": self.uri,
                "progress": self.progress,
                "total": self.total,
            }
        )
    
    async def complete(self, result: Any = None) -> None:
        """Mark the resource as completed.
        
        Args:
            result: The result of the operation (not stored)
        """
        self.status = AsyncStatus.COMPLETED
        self.completed_at = time.time()
        
        # If we have a total, set progress to match it
        if self.total is not None:
            self.progress = self.total
            
        await self._notify_changed()
        
        logger.info(
            "Completed async operation",
            extra={
                "uri": self.uri,
                "duration": self.completed_at - (self.started_at or self.created_at),
            }
        )
    
    async def fail(self, error: str) -> None:
        """Mark the resource as failed and store the error.
        
        Args:
            error: Error message describing why the operation failed
        """
        self.status = AsyncStatus.FAILED
        self.error = error
        self.completed_at = time.time()
        await self._notify_changed()
        
        logger.error(
            "Failed async operation",
            extra={
                "uri": self.uri,
                "error": error,
                "duration": self.completed_at - (self.started_at or self.created_at),
            }
        )
    
    async def cancel(self) -> None:
        """Cancel the operation if it's still running."""
        if self.status in (AsyncStatus.PENDING, AsyncStatus.RUNNING) and self._task:
            self._task.cancel()
            self.status = AsyncStatus.CANCELED
            self.completed_at = time.time()
            await self._notify_changed()
            
            logger.info(
                "Canceled async operation",
                extra={
                    "uri": self.uri,
                    "duration": self.completed_at - (self.started_at or self.created_at),
                }
            )
    
    async def _notify_changed(self) -> None:
        """Notify subscribers that the resource has changed."""
        if self._mcp_server:
            # This will be implemented in the MCP server to notify clients
            # of resource changes via the notification protocol
            await self._mcp_server.notify_resource_changed(self.uri)

class Tool(BaseModel):
    # ... existing fields ...
    async_supported: bool = Field(
        False, description="Whether this tool supports asynchronous execution"
    )

def tool(
    self,
    name: str | None = None,
    description: str | None = None,
    annotations: ToolAnnotations | None = None,
    async_supported: bool = False,
) -> Callable[[AnyFunction], AnyFunction]:
    # ... existing code ...
    def decorator(fn: AnyFunction) -> AnyFunction:
        tool = self.add_tool(
            fn, 
            name=name, 
            description=description, 
            annotations=annotations,
        )
        
        # Store async support information
        tool.async_supported = async_supported
        
        return fn
    return decorator

def create_async_resource(
    self,
    name: str | None = None,
    description: str | None = None
) -> AsyncResource:
    """Create a new async resource."""
    resource_uri = f"resource://tasks/{uuid.uuid4()}"
    resource = AsyncResource(
        AnyUrl(resource_uri),
        name=name,
        description=description
    )
    
    # Set the MCP server reference if available
    if self._mcp_server:
        resource._mcp_server = self._mcp_server
        
    self.add_resource(resource)
    return resource

class CallToolResult(Result):
  """The server's response to a tools/call request from the client."""
  content: list[TextContent | ImageContent | DataContent | EmbeddedResource | AsyncResource] | None = None

async def call_tool(
  self,
  name: str,
  arguments: dict[str, Any],
  context: Context[ServerSessionT, LifespanContextT] | None = None,
) -> Any:
  """Call a tool by name with arguments."""
  tool = self.get_tool(name)
  if not tool:
      raise ToolError(f"Unknown tool: {name}")
      
  # Check if the tool supports async execution
  if tool.async_supported:
      # Create an async resource
      resource = self._resource_manager.create_async_resource(
          name=tool.name,
          description=tool.description
      )
      
      # Set the resource in the context
      if context:
          context.resource = resource
      
      # Create a task to run the tool
      async def run_tool_async():
          try:
              # Run the tool
              result = await tool.run(arguments, context=context)
              
              # Complete the resource with the result
              await resource.complete(result)
              
              return result
          except Exception as e:
              # Mark the resource as failed
              await resource.fail(str(e))
              raise
              
      # Create and start the task
      task = asyncio.create_task(run_tool_async())
      
      # Start the resource
      await resource.start(task)
      
      # Return the resource URI
      return {
          "type": "resource",
          "uri": str(resource.uri),
          "status": resource.status.value,
      }
  else:
      # Run the tool synchronously
      return await tool.run(arguments, context=context)

{
  "type": "resource",
  "uri": "resource://tasks/1234-5678-90ab-cdef",
  "status": "pending"
}

{
  "uri": "resource://tasks/1234-5678-90ab-cdef",
  "name": "Process Data",
  "description": "Processing data items",
  "status": "running",
  "progress": 3,
  "total": 10,
  "result": null,
  "error": null,
  "created_at": 1625097600.123,
  "started_at": 1625097601.456,
  "completed_at": null
}

@app.tool(async_supported=True)
async def process_data(data: list[str], ctx: Context) -> Dict[str, Any]:
  ......
 

const response = await client.callTool("process_data", {
  data: ["item1", "item2", "item3"]
});

// Check if the response is a resource URI
if (response.type === "resource") {
  // Get the resource URI
  const resourceUri = response.uri;
  
  // Subscribe to the resource
  await client.subscribeResource(resourceUri);
  
  // Set up notification handler
  client.on('notifications/resource/changed', async (params) => {
    if (params.uri === resourceUri) {
      // Get the updated resource
      const resource = await client.readResource(resourceUri);
      
      // Update UI based on resource status
      updateStatus(resource.status);
      
      if (resource.status === "running" && resource.total) {
        // Update progress bar
        updateProgressBar(resource.progress, resource.total);
      }
      
      if (resource.status === "completed") {
        // Display result
        displayResult(resource.result);
        
        // Unsubscribe from updates
        await client.unsubscribeResource(resourceUri);
      }
    }
  });
} else {
  // The tool ran synchronously
  displayResult(response);
}