comfy_api_nodes/apis/client.py

"""
API Client Framework for api.comfy.org.

This module provides a flexible framework for making API requests from ComfyUI nodes.
It supports both synchronous and asynchronous API operations with proper type validation.

Key Components:
--------------
1. ApiClient - Handles HTTP requests with authentication and error handling
2. ApiEndpoint - Defines a single HTTP endpoint with its request/response models
3. ApiOperation - Executes a single synchronous API operation

Usage Examples:
--------------

# Example 1: Synchronous API Operation
# ------------------------------------
# For a simple API call that returns the result immediately:

# 1. Create the API client
api_client = ApiClient(
    base_url="https://api.example.com",
    auth_token="your_auth_token_here",
    comfy_api_key="your_comfy_api_key_here",
    timeout=30.0,
    verify_ssl=True
)

# 2. Define the endpoint
user_info_endpoint = ApiEndpoint(
    path="/v1/users/me",
    method=HttpMethod.GET,
    request_model=EmptyRequest,  # No request body needed
    response_model=UserProfile,   # Pydantic model for the response
    query_params=None
)

# 3. Create the request object
request = EmptyRequest()

# 4. Create and execute the operation
operation = ApiOperation(
    endpoint=user_info_endpoint,
    request=request
)
user_profile = operation.execute(client=api_client)  # Returns immediately with the result


# Example 2: Asynchronous API Operation with Polling
# -------------------------------------------------
# For an API that starts a task and requires polling for completion:

# 1. Define the endpoints (initial request and polling)
generate_image_endpoint = ApiEndpoint(
    path="/v1/images/generate",
    method=HttpMethod.POST,
    request_model=ImageGenerationRequest,
    response_model=TaskCreatedResponse,
    query_params=None
)

check_task_endpoint = ApiEndpoint(
    path="/v1/tasks/{task_id}",
    method=HttpMethod.GET,
    request_model=EmptyRequest,
    response_model=ImageGenerationResult,
    query_params=None
)

# 2. Create the request object
request = ImageGenerationRequest(
    prompt="a beautiful sunset over mountains",
    width=1024,
    height=1024,
    num_images=1
)

# 3. Create and execute the polling operation
operation = PollingOperation(
    initial_endpoint=generate_image_endpoint,
    initial_request=request,
    poll_endpoint=check_task_endpoint,
    task_id_field="task_id",
    status_field="status",
    completed_statuses=["completed"],
    failed_statuses=["failed", "error"]
)

# This will make the initial request and then poll until completion
result = operation.execute(client=api_client)  # Returns the final ImageGenerationResult when done
"""

from __future__ import annotations
import logging
import time
import io
import socket
from typing import Dict, Type, Optional, Any, TypeVar, Generic, Callable, Tuple
from enum import Enum
import json
import requests
from urllib.parse import urljoin, urlparse
from pydantic import BaseModel, Field
import uuid # For generating unique operation IDs

from server import PromptServer
from comfy.cli_args import args
from comfy import utils
from . import request_logger

T = TypeVar("T", bound=BaseModel)
R = TypeVar("R", bound=BaseModel)
P = TypeVar("P", bound=BaseModel)  # For poll response

PROGRESS_BAR_MAX = 100


class NetworkError(Exception):
    """Base exception for network-related errors with diagnostic information."""
    pass


class LocalNetworkError(NetworkError):
    """Exception raised when local network connectivity issues are detected."""
    pass


class ApiServerError(NetworkError):
    """Exception raised when the API server is unreachable but internet is working."""
    pass


class EmptyRequest(BaseModel):
    """Base class for empty request bodies.
    For GET requests, fields will be sent as query parameters."""

    pass


class UploadRequest(BaseModel):
    file_name: str = Field(..., description="Filename to upload")
    content_type: str | None = Field(
        None,
        description="Mime type of the file. For example: image/png, image/jpeg, video/mp4, etc.",
    )


class UploadResponse(BaseModel):
    download_url: str = Field(..., description="URL to GET uploaded file")
    upload_url: str = Field(..., description="URL to PUT file to upload")


class HttpMethod(str, Enum):
    GET = "GET"
    POST = "POST"
    PUT = "PUT"
    DELETE = "DELETE"
    PATCH = "PATCH"


class ApiClient:
    """
    Client for making HTTP requests to an API with authentication, error handling, and retry logic.
    """

    def __init__(
        self,
        base_url: str,
        auth_token: Optional[str] = None,
        comfy_api_key: Optional[str] = None,
        timeout: float = 3600.0,
        verify_ssl: bool = True,
        max_retries: int = 3,
        retry_delay: float = 1.0,
        retry_backoff_factor: float = 2.0,
        retry_status_codes: Optional[Tuple[int, ...]] = None,
    ):
        self.base_url = base_url
        self.auth_token = auth_token
        self.comfy_api_key = comfy_api_key
        self.timeout = timeout
        self.verify_ssl = verify_ssl
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        self.retry_backoff_factor = retry_backoff_factor
        # Default retry status codes: 408 (Request Timeout), 429 (Too Many Requests),
        # 500, 502, 503, 504 (Server Errors)
        self.retry_status_codes = retry_status_codes or (408, 429, 500, 502, 503, 504)

    def _generate_operation_id(self, path: str) -> str:
        """Generates a unique operation ID for logging."""
        return f"{path.strip('/').replace('/', '_')}_{uuid.uuid4().hex[:8]}"

    def _create_json_payload_args(
        self,
        data: Optional[Dict[str, Any]] = None,
        headers: Optional[Dict[str, str]] = None,
    ) -> Dict[str, Any]:
        return {
            "json": data,
            "headers": headers,
        }

    def _create_form_data_args(
        self,
        data: Dict[str, Any],
        files: Dict[str, Any],
        headers: Optional[Dict[str, str]] = None,
        multipart_parser = None,
    ) -> Dict[str, Any]:
        if headers and "Content-Type" in headers:
            del headers["Content-Type"]

        if multipart_parser:
            data = multipart_parser(data)

        return {
            "data": data,
            "files": files,
            "headers": headers,
        }

    def _create_urlencoded_form_data_args(
        self,
        data: Dict[str, Any],
        headers: Optional[Dict[str, str]] = None,
    ) -> Dict[str, Any]:
        headers = headers or {}
        headers["Content-Type"] = "application/x-www-form-urlencoded"

        return {
            "data": data,
            "headers": headers,
        }

    def get_headers(self) -> Dict[str, str]:
        """Get headers for API requests, including authentication if available"""
        headers = {"Content-Type": "application/json", "Accept": "application/json"}

        if self.auth_token:
            headers["Authorization"] = f"Bearer {self.auth_token}"
        elif self.comfy_api_key:
            headers["X-API-KEY"] = self.comfy_api_key

        return headers

    def _check_connectivity(self, target_url: str) -> Dict[str, bool]:
        """
        Check connectivity to determine if network issues are local or server-related.

        Args:
            target_url: URL to check connectivity to

        Returns:
            Dictionary with connectivity status details
        """
        results = {
            "internet_accessible": False,
            "api_accessible": False,
            "is_local_issue": False,
            "is_api_issue": False
        }

        # First check basic internet connectivity using a reliable external site
        try:
            # Use a reliable external domain for checking basic connectivity
            check_response = requests.get("https://www.google.com",
                                         timeout=5.0,
                                         verify=self.verify_ssl)
            if check_response.status_code < 500:
                results["internet_accessible"] = True
        except (requests.RequestException, socket.error):
            results["internet_accessible"] = False
            results["is_local_issue"] = True
            return results

        # Now check API server connectivity
        try:
            # Extract domain from the target URL to do a simpler health check
            parsed_url = urlparse(target_url)
            api_base = f"{parsed_url.scheme}://{parsed_url.netloc}"

            # Try to reach the API domain
            api_response = requests.get(f"{api_base}/health", timeout=5.0, verify=self.verify_ssl)
            if api_response.status_code < 500:
                results["api_accessible"] = True
            else:
                results["api_accessible"] = False
                results["is_api_issue"] = True
        except requests.RequestException:
            results["api_accessible"] = False
            # If we can reach the internet but not the API, it's an API issue
            results["is_api_issue"] = True

        return results

    def request(
        self,
        method: str,
        path: str,
        params: Optional[Dict[str, Any]] = None,
        data: Optional[Dict[str, Any]] = None,
        files: Optional[Dict[str, Any]] = None,
        headers: Optional[Dict[str, str]] = None,
        content_type: str = "application/json",
        multipart_parser: Callable = None,
        retry_count: int = 0,  # Used internally for tracking retries
    ) -> Dict[str, Any]:
        """
        Make an HTTP request to the API with automatic retries for transient errors.

        Args:
            method: HTTP method (GET, POST, etc.)
            path: API endpoint path (will be joined with base_url)
            params: Query parameters
            data: body data
            files: Files to upload
            headers: Additional headers
            content_type: Content type of the request. Defaults to application/json.
            retry_count: Internal parameter for tracking retries, do not set manually

        Returns:
            Parsed JSON response

        Raises:
            LocalNetworkError: If local network connectivity issues are detected
            ApiServerError: If the API server is unreachable but internet is working
            Exception: For other request failures
        """
        url = urljoin(self.base_url, path)
        self.check_auth(self.auth_token, self.comfy_api_key)
        # Combine default headers with any provided headers
        request_headers = self.get_headers()
        if headers:
            request_headers.update(headers)

        # Let requests handle the content type when files are present.
        if files:
            del request_headers["Content-Type"]

        logging.debug(f"[DEBUG] Request Headers: {request_headers}")
        logging.debug(f"[DEBUG] Files: {files}")
        logging.debug(f"[DEBUG] Params: {params}")
        logging.debug(f"[DEBUG] Data: {data}")

        if content_type == "application/x-www-form-urlencoded":
            payload_args = self._create_urlencoded_form_data_args(data, request_headers)
        elif content_type == "multipart/form-data":
            payload_args = self._create_form_data_args(
                data, files, request_headers, multipart_parser
            )
        else:
            payload_args = self._create_json_payload_args(data, request_headers)

        operation_id = self._generate_operation_id(path)
        request_logger.log_request_response(
            operation_id=operation_id,
            request_method=method,
            request_url=url,
            request_headers=request_headers,
            request_params=params,
            request_data=data if content_type == "application/json" else "[form-data or other]"
        )

        try:
            response = requests.request(
                method=method,
                url=url,
                params=params,
                timeout=self.timeout,
                verify=self.verify_ssl,
                **payload_args,
            )

            # Check if we should retry based on status code
            if (response.status_code in self.retry_status_codes and
                retry_count < self.max_retries):

                # Calculate delay with exponential backoff
                delay = self.retry_delay * (self.retry_backoff_factor ** retry_count)

                logging.warning(
                    f"Request failed with status {response.status_code}. "
                    f"Retrying in {delay:.2f}s ({retry_count + 1}/{self.max_retries})"
                )

                time.sleep(delay)
                return self.request(
                    method=method,
                    path=path,
                    params=params,
                    data=data,
                    files=files,
                    headers=headers,
                    content_type=content_type,
                    multipart_parser=multipart_parser,
                    retry_count=retry_count + 1,
                )

            # Raise exception for error status codes
            response.raise_for_status()

            # Log successful response
            response_content_to_log = response.content
            try:
                # Attempt to parse JSON for prettier logging, fallback to raw content
                response_content_to_log = response.json()
            except json.JSONDecodeError:
                pass # Keep as bytes/str if not JSON

            request_logger.log_request_response(
                operation_id=operation_id,
                request_method=method, # Pass request details again for context in log
                request_url=url,
                response_status_code=response.status_code,
                response_headers=dict(response.headers),
                response_content=response_content_to_log
            )

        except requests.ConnectionError as e:
            error_message = f"ConnectionError: {str(e)}"
            request_logger.log_request_response(
                operation_id=operation_id,
                request_method=method,
                request_url=url,
                error_message=error_message
            )
            # Only perform connectivity check if we've exhausted all retries
            if retry_count >= self.max_retries:
                # Check connectivity to determine if it's a local or API issue
                connectivity = self._check_connectivity(self.base_url)

                if connectivity["is_local_issue"]:
                    raise LocalNetworkError(
                        "Unable to connect to the API server due to local network issues. "
                        "Please check your internet connection and try again."
                    ) from e
                elif connectivity["is_api_issue"]:
                    raise ApiServerError(
                        f"The API server at {self.base_url} is currently unreachable. "
                        f"The service may be experiencing issues. Please try again later."
                    ) from e

            # If we haven't exhausted retries yet, retry the request
            if retry_count < self.max_retries:
                delay = self.retry_delay * (self.retry_backoff_factor ** retry_count)
                logging.warning(
                    f"Connection error: {str(e)}. "
                    f"Retrying in {delay:.2f}s ({retry_count + 1}/{self.max_retries})"
                )
                time.sleep(delay)
                return self.request(
                    method=method,
                    path=path,
                    params=params,
                    data=data,
                    files=files,
                    headers=headers,
                    content_type=content_type,
                    multipart_parser=multipart_parser,
                    retry_count=retry_count + 1,
                )

            # If we've exhausted retries and didn't identify the specific issue,
            # raise a generic exception
            final_error_message = (
                f"Unable to connect to the API server after {self.max_retries} attempts. "
                f"Please check your internet connection or try again later."
            )
            request_logger.log_request_response( # Log final failure
                operation_id=operation_id,
                request_method=method, request_url=url,
                error_message=final_error_message
            )
            raise Exception(final_error_message) from e

        except requests.Timeout as e:
            error_message = f"Timeout: {str(e)}"
            request_logger.log_request_response(
                operation_id=operation_id,
                request_method=method, request_url=url,
                error_message=error_message
            )
            # Retry timeouts if we haven't exhausted retries
            if retry_count < self.max_retries:
                delay = self.retry_delay * (self.retry_backoff_factor ** retry_count)
                logging.warning(
                    f"Request timed out. "
                    f"Retrying in {delay:.2f}s ({retry_count + 1}/{self.max_retries})"
                )
                time.sleep(delay)
                return self.request(
                    method=method,
                    path=path,
                    params=params,
                    data=data,
                    files=files,
                    headers=headers,
                    content_type=content_type,
                    multipart_parser=multipart_parser,
                    retry_count=retry_count + 1,
                )
            final_error_message = (
                f"Request timed out after {self.timeout} seconds and {self.max_retries} retry attempts. "
                f"The server might be experiencing high load or the operation is taking longer than expected."
            )
            request_logger.log_request_response( # Log final failure
                operation_id=operation_id,
                request_method=method, request_url=url,
                error_message=final_error_message
            )
            raise Exception(final_error_message) from e

        except requests.HTTPError as e:
            status_code = e.response.status_code if hasattr(e, "response") else None
            original_error_message = f"HTTP Error: {str(e)}"
            error_content_for_log = None
            if hasattr(e, "response") and e.response is not None:
                error_content_for_log = e.response.content
                try:
                    error_content_for_log = e.response.json()
                except json.JSONDecodeError:
                    pass


            # Try to extract detailed error message from JSON response for user display
            # but log the full error content.
            user_display_error_message = original_error_message

            try:
                if hasattr(e, "response") and e.response is not None and e.response.content:
                    error_json = e.response.json()
                    if "error" in error_json and "message" in error_json["error"]:
                        user_display_error_message = f"API Error: {error_json['error']['message']}"
                        if "type" in error_json["error"]:
                            user_display_error_message += f" (Type: {error_json['error']['type']})"
                    elif isinstance(error_json, dict): # Handle cases where error is just a JSON dict
                        user_display_error_message = f"API Error: {json.dumps(error_json)}"
                    else: # Non-dict JSON error
                        user_display_error_message = f"API Error: {str(error_json)}"
            except json.JSONDecodeError:
                # If not JSON, use the raw content if it's not too long, or a summary
                if hasattr(e, "response") and e.response is not None and e.response.content:
                    raw_content = e.response.content.decode(errors='ignore')
                    if len(raw_content) < 200: # Arbitrary limit for display
                        user_display_error_message = f"API Error (raw): {raw_content}"
                    else:
                        user_display_error_message = f"API Error (raw, status {status_code})"

            request_logger.log_request_response(
                operation_id=operation_id,
                request_method=method, request_url=url,
                response_status_code=status_code,
                response_headers=dict(e.response.headers) if hasattr(e, "response") and e.response is not None else None,
                response_content=error_content_for_log,
                error_message=original_error_message # Log the original exception string as error
            )

            logging.debug(f"[DEBUG] API Error: {user_display_error_message} (Status: {status_code})")
            if hasattr(e, "response") and e.response is not None and e.response.content:
                logging.debug(f"[DEBUG] Response content: {e.response.content}")

            # Retry if the status code is in our retry list and we haven't exhausted retries
            if (status_code in self.retry_status_codes and
                retry_count < self.max_retries):

                delay = self.retry_delay * (self.retry_backoff_factor ** retry_count)
                logging.warning(
                    f"HTTP error {status_code}. "
                    f"Retrying in {delay:.2f}s ({retry_count + 1}/{self.max_retries})"
                )
                time.sleep(delay)
                return self.request(
                    method=method,
                    path=path,
                    params=params,
                    data=data,
                    files=files,
                    headers=headers,
                    content_type=content_type,
                    multipart_parser=multipart_parser,
                    retry_count=retry_count + 1,
                )

            # Specific error messages for common status codes for user display
            if status_code == 401:
                user_display_error_message = "Unauthorized: Please login first to use this node."
            elif status_code == 402:
                user_display_error_message = "Payment Required: Please add credits to your account to use this node."
            elif status_code == 409:
                user_display_error_message = "There is a problem with your account. Please contact support@comfy.org."
            elif status_code == 429:
                user_display_error_message = "Rate Limit Exceeded: Please try again later."
            # else, user_display_error_message remains as parsed from response or original HTTPError string

            raise Exception(user_display_error_message) # Raise with the user-friendly message

        # Parse and return JSON response
        if response.content:
            return response.json()
        return {}

    def check_auth(self, auth_token, comfy_api_key):
        """Verify that an auth token is present or comfy_api_key is present"""
        if auth_token is None and comfy_api_key is None:
            raise Exception("Unauthorized: Please login first to use this node.")
        return auth_token or comfy_api_key

    @staticmethod
    def upload_file(
        upload_url: str,
        file: io.BytesIO | str,
        content_type: str | None = None,
        max_retries: int = 3,
        retry_delay: float = 1.0,
        retry_backoff_factor: float = 2.0,
    ):
        """Upload a file to the API with retry logic.

        Args:
            upload_url: The URL to upload to
            file: Either a file path string, BytesIO object, or tuple of (file_path, filename)
            content_type: Optional mime type to set for the upload
            max_retries: Maximum number of retry attempts
            retry_delay: Initial delay between retries in seconds
            retry_backoff_factor: Multiplier for the delay after each retry
        """
        headers = {}
        if content_type:
            headers["Content-Type"] = content_type

        # Prepare the file data
        if isinstance(file, io.BytesIO):
            file.seek(0)  # Ensure we're at the start of the file
            data = file.read()
        elif isinstance(file, str):
            with open(file, "rb") as f:
                data = f.read()
        else:
            raise ValueError("File must be either a BytesIO object or a file path string")

        # Try the upload with retries
        last_exception = None
        operation_id = f"upload_{upload_url.split('/')[-1]}_{uuid.uuid4().hex[:8]}" # Simplified ID for uploads

        # Log initial attempt (without full file data for brevity)
        request_logger.log_request_response(
            operation_id=operation_id,
            request_method="PUT",
            request_url=upload_url,
            request_headers=headers,
            request_data=f"[File data of type {content_type or 'unknown'}, size {len(data)} bytes]"
        )

        for retry_attempt in range(max_retries + 1):
            try:
                response = requests.put(upload_url, data=data, headers=headers)
                response.raise_for_status()
                request_logger.log_request_response(
                    operation_id=operation_id,
                    request_method="PUT", request_url=upload_url, # For context
                    response_status_code=response.status_code,
                    response_headers=dict(response.headers),
                    response_content="File uploaded successfully." # Or response.text if available
                )
                return response

            except (requests.ConnectionError, requests.Timeout, requests.HTTPError) as e:
                last_exception = e
                error_message_for_log = f"{type(e).__name__}: {str(e)}"
                response_content_for_log = None
                status_code_for_log = None
                headers_for_log = None

                if hasattr(e, 'response') and e.response is not None:
                    status_code_for_log = e.response.status_code
                    headers_for_log = dict(e.response.headers)
                    try:
                        response_content_for_log = e.response.json()
                    except json.JSONDecodeError:
                        response_content_for_log = e.response.content


                request_logger.log_request_response(
                    operation_id=operation_id,
                    request_method="PUT", request_url=upload_url,
                    response_status_code=status_code_for_log,
                    response_headers=headers_for_log,
                    response_content=response_content_for_log,
                    error_message=error_message_for_log
                )

                if retry_attempt < max_retries:
                    delay = retry_delay * (retry_backoff_factor ** retry_attempt)
                    logging.warning(
                        f"File upload failed: {str(e)}. "
                        f"Retrying in {delay:.2f}s ({retry_attempt + 1}/{max_retries})"
                    )
                    time.sleep(delay)
                else:
                    break # Max retries reached

        # If we've exhausted all retries, determine the final error type and raise
        final_error_message = f"Failed to upload file after {max_retries + 1} attempts. Error: {str(last_exception)}"
        try:
            # Check basic internet connectivity
            check_response = requests.get("https://www.google.com", timeout=5.0, verify=True) # Assuming verify=True is desired
            if check_response.status_code >= 500: # Google itself has an issue (rare)
                 final_error_message = (f"Failed to upload file. Internet connectivity check to Google failed "
                                       f"(status {check_response.status_code}). Original error: {str(last_exception)}")
                 # Not raising LocalNetworkError here as Google itself might be down.
            # If Google is reachable, the issue is likely with the upload server or a more specific local problem
            # not caught by a simple Google ping (e.g., DNS for the specific upload URL, firewall).
            # The original last_exception is probably most relevant.

        except (requests.RequestException, socket.error) as conn_check_exc:
            # Could not reach Google, likely a local network issue
            final_error_message = (f"Failed to upload file due to network connectivity issues "
                                   f"(cannot reach Google: {str(conn_check_exc)}). "
                                   f"Original upload error: {str(last_exception)}")
            request_logger.log_request_response( # Log final failure reason
                operation_id=operation_id,
                request_method="PUT", request_url=upload_url,
                error_message=final_error_message
            )
            raise LocalNetworkError(final_error_message) from last_exception

        request_logger.log_request_response( # Log final failure reason if not LocalNetworkError
            operation_id=operation_id,
            request_method="PUT", request_url=upload_url,
            error_message=final_error_message
        )
        raise Exception(final_error_message) from last_exception


class ApiEndpoint(Generic[T, R]):
    """Defines an API endpoint with its request and response types"""

    def __init__(
        self,
        path: str,
        method: HttpMethod,
        request_model: Type[T],
        response_model: Type[R],
        query_params: Optional[Dict[str, Any]] = None,
    ):
        """Initialize an API endpoint definition.

        Args:
            path: The URL path for this endpoint, can include placeholders like {id}
            method: The HTTP method to use (GET, POST, etc.)
            request_model: Pydantic model class that defines the structure and validation rules for API requests to this endpoint
            response_model: Pydantic model class that defines the structure and validation rules for API responses from this endpoint
            query_params: Optional dictionary of query parameters to include in the request
        """
        self.path = path
        self.method = method
        self.request_model = request_model
        self.response_model = response_model
        self.query_params = query_params or {}


class SynchronousOperation(Generic[T, R]):
    """
    Represents a single synchronous API operation.
    """

    def __init__(
        self,
        endpoint: ApiEndpoint[T, R],
        request: T,
        files: Optional[Dict[str, Any]] = None,
        api_base: str | None = None,
        auth_token: Optional[str] = None,
        comfy_api_key: Optional[str] = None,
        auth_kwargs: Optional[Dict[str,str]] = None,
        timeout: float = 604800.0,
        verify_ssl: bool = True,
        content_type: str = "application/json",
        multipart_parser: Callable = None,
        max_retries: int = 3,
        retry_delay: float = 1.0,
        retry_backoff_factor: float = 2.0,
    ):
        self.endpoint = endpoint
        self.request = request
        self.response = None
        self.error = None
        self.api_base: str = api_base or args.comfy_api_base
        self.auth_token = auth_token
        self.comfy_api_key = comfy_api_key
        if auth_kwargs is not None:
            self.auth_token = auth_kwargs.get("auth_token", self.auth_token)
            self.comfy_api_key = auth_kwargs.get("comfy_api_key", self.comfy_api_key)
        self.timeout = timeout
        self.verify_ssl = verify_ssl
        self.files = files
        self.content_type = content_type
        self.multipart_parser = multipart_parser
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        self.retry_backoff_factor = retry_backoff_factor

    def execute(self, client: Optional[ApiClient] = None) -> R:
        """Execute the API operation using the provided client or create one with retry support"""
        try:
            # Create client if not provided
            if client is None:
                client = ApiClient(
                    base_url=self.api_base,
                    auth_token=self.auth_token,
                    comfy_api_key=self.comfy_api_key,
                    timeout=self.timeout,
                    verify_ssl=self.verify_ssl,
                    max_retries=self.max_retries,
                    retry_delay=self.retry_delay,
                    retry_backoff_factor=self.retry_backoff_factor,
                )

            # Convert request model to dict, but use None for EmptyRequest
            request_dict = (
                None
                if isinstance(self.request, EmptyRequest)
                else self.request.model_dump(exclude_none=True)
            )
            if request_dict:
                for key, value in request_dict.items():
                    if isinstance(value, Enum):
                        request_dict[key] = value.value

            # Debug log for request
            logging.debug(
                f"[DEBUG] API Request: {self.endpoint.method.value} {self.endpoint.path}"
            )
            logging.debug(f"[DEBUG] Request Data: {json.dumps(request_dict, indent=2)}")
            logging.debug(f"[DEBUG] Query Params: {self.endpoint.query_params}")

            # Make the request with built-in retry
            resp = client.request(
                method=self.endpoint.method.value,
                path=self.endpoint.path,
                data=request_dict,
                params=self.endpoint.query_params,
                files=self.files,
                content_type=self.content_type,
                multipart_parser=self.multipart_parser
            )

            # Debug log for response
            logging.debug("=" * 50)
            logging.debug("[DEBUG] RESPONSE DETAILS:")
            logging.debug("[DEBUG] Status Code: 200 (Success)")
            logging.debug(f"[DEBUG] Response Body: {json.dumps(resp, indent=2)}")
            logging.debug("=" * 50)

            # Parse and return the response
            return self._parse_response(resp)

        except LocalNetworkError as e:
            # Propagate specific network error types
            logging.error(f"[ERROR] Local network error: {str(e)}")
            raise

        except ApiServerError as e:
            # Propagate API server errors
            logging.error(f"[ERROR] API server error: {str(e)}")
            raise

        except Exception as e:
            logging.error(f"[ERROR] API Exception: {str(e)}")
            raise Exception(str(e))

    def _parse_response(self, resp):
        """Parse response data - can be overridden by subclasses"""
        # The response is already the complete object, don't extract just the "data" field
        # as that would lose the outer structure (created timestamp, etc.)

        # Parse response using the provided model
        self.response = self.endpoint.response_model.model_validate(resp)
        logging.debug(f"[DEBUG] Parsed Response: {self.response}")
        return self.response


class TaskStatus(str, Enum):
    """Enum for task status values"""

    COMPLETED = "completed"
    FAILED = "failed"
    PENDING = "pending"


class PollingOperation(Generic[T, R]):
    """
    Represents an asynchronous API operation that requires polling for completion.
    """

    def __init__(
        self,
        poll_endpoint: ApiEndpoint[EmptyRequest, R],
        completed_statuses: list,
        failed_statuses: list,
        status_extractor: Callable[[R], str],
        progress_extractor: Callable[[R], float] = None,
        result_url_extractor: Callable[[R], str] = None,
        request: Optional[T] = None,
        api_base: str | None = None,
        auth_token: Optional[str] = None,
        comfy_api_key: Optional[str] = None,
        auth_kwargs: Optional[Dict[str,str]] = None,
        poll_interval: float = 5.0,
        max_poll_attempts: int = 120,  # Default max polling attempts (10 minutes with 5s interval)
        max_retries: int = 3,  # Max retries per individual API call
        retry_delay: float = 1.0,
        retry_backoff_factor: float = 2.0,
        estimated_duration: Optional[float] = None,
        node_id: Optional[str] = None,
    ):
        self.poll_endpoint = poll_endpoint
        self.request = request
        self.api_base: str = api_base or args.comfy_api_base
        self.auth_token = auth_token
        self.comfy_api_key = comfy_api_key
        if auth_kwargs is not None:
            self.auth_token = auth_kwargs.get("auth_token", self.auth_token)
            self.comfy_api_key = auth_kwargs.get("comfy_api_key", self.comfy_api_key)
        self.poll_interval = poll_interval
        self.max_poll_attempts = max_poll_attempts
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        self.retry_backoff_factor = retry_backoff_factor
        self.estimated_duration = estimated_duration

        # Polling configuration
        self.status_extractor = status_extractor or (
            lambda x: getattr(x, "status", None)
        )
        self.progress_extractor = progress_extractor
        self.result_url_extractor = result_url_extractor
        self.node_id = node_id
        self.completed_statuses = completed_statuses
        self.failed_statuses = failed_statuses

        # For storing response data
        self.final_response = None
        self.error = None

    def execute(self, client: Optional[ApiClient] = None) -> R:
        """Execute the polling operation using the provided client. If failed, raise an exception."""
        try:
            if client is None:
                client = ApiClient(
                    base_url=self.api_base,
                    auth_token=self.auth_token,
                    comfy_api_key=self.comfy_api_key,
                    max_retries=self.max_retries,
                    retry_delay=self.retry_delay,
                    retry_backoff_factor=self.retry_backoff_factor,
                )
            return self._poll_until_complete(client)
        except LocalNetworkError as e:
            # Provide clear message for local network issues
            raise Exception(
                f"Polling failed due to local network issues. Please check your internet connection. "
                f"Details: {str(e)}"
            ) from e
        except ApiServerError as e:
            # Provide clear message for API server issues
            raise Exception(
                f"Polling failed due to API server issues. The service may be experiencing problems. "
                f"Please try again later. Details: {str(e)}"
            ) from e
        except Exception as e:
            raise Exception(f"Error during polling: {str(e)}")

    def _display_text_on_node(self, text: str):
        """Sends text to the client which will be displayed on the node in the UI"""
        if not self.node_id:
            return

        PromptServer.instance.send_progress_text(text, self.node_id)

    def _display_time_progress_on_node(self, time_completed: int):
        if not self.node_id:
            return

        if self.estimated_duration is not None:
            estimated_time_remaining = max(
                0, int(self.estimated_duration) - int(time_completed)
            )
            message = f"Task in progress: {time_completed:.0f}s (~{estimated_time_remaining:.0f}s remaining)"
        else:
            message = f"Task in progress: {time_completed:.0f}s"
        self._display_text_on_node(message)

    def _check_task_status(self, response: R) -> TaskStatus:
        """Check task status using the status extractor function"""
        try:
            status = self.status_extractor(response)
            if status in self.completed_statuses:
                return TaskStatus.COMPLETED
            elif status in self.failed_statuses:
                return TaskStatus.FAILED
            return TaskStatus.PENDING
        except Exception as e:
            logging.error(f"Error extracting status: {e}")
            return TaskStatus.PENDING

    def _poll_until_complete(self, client: ApiClient) -> R:
        """Poll until the task is complete"""
        poll_count = 0
        consecutive_errors = 0
        max_consecutive_errors = min(5, self.max_retries * 2)  # Limit consecutive errors

        if self.progress_extractor:
            progress = utils.ProgressBar(PROGRESS_BAR_MAX)

        while poll_count < self.max_poll_attempts:
            try:
                poll_count += 1
                logging.debug(f"[DEBUG] Polling attempt #{poll_count}")

                request_dict = (
                    self.request.model_dump(exclude_none=True)
                    if self.request is not None
                    else None
                )

                if poll_count == 1:
                    logging.debug(
                        f"[DEBUG] Poll Request: {self.poll_endpoint.method.value} {self.poll_endpoint.path}"
                    )
                    logging.debug(
                        f"[DEBUG] Poll Request Data: {json.dumps(request_dict, indent=2) if request_dict else 'None'}"
                    )

                # Query task status
                resp = client.request(
                    method=self.poll_endpoint.method.value,
                    path=self.poll_endpoint.path,
                    params=self.poll_endpoint.query_params,
                    data=request_dict,
                )

                # Successfully got a response, reset consecutive error count
                consecutive_errors = 0

                # Parse response
                response_obj = self.poll_endpoint.response_model.model_validate(resp)

                # Check if task is complete
                status = self._check_task_status(response_obj)
                logging.debug(f"[DEBUG] Task Status: {status}")

                # If progress extractor is provided, extract progress
                if self.progress_extractor:
                    new_progress = self.progress_extractor(response_obj)
                    if new_progress is not None:
                        progress.update_absolute(new_progress, total=PROGRESS_BAR_MAX)

                if status == TaskStatus.COMPLETED:
                    message = "Task completed successfully"
                    if self.result_url_extractor:
                        result_url = self.result_url_extractor(response_obj)
                        if result_url:
                            message = f"Result URL: {result_url}"
                    else:
                        message = "Task completed successfully!"
                    logging.debug(f"[DEBUG] {message}")
                    self._display_text_on_node(message)
                    self.final_response = response_obj
                    if self.progress_extractor:
                        progress.update(100)
                    return self.final_response
                elif status == TaskStatus.FAILED:
                    message = f"Task failed: {json.dumps(resp)}"
                    logging.error(f"[DEBUG] {message}")
                    raise Exception(message)
                else:
                    logging.debug("[DEBUG] Task still pending, continuing to poll...")

                # Wait before polling again
                logging.debug(
                    f"[DEBUG] Waiting {self.poll_interval} seconds before next poll"
                )
                for i in range(int(self.poll_interval)):
                    time_completed = (poll_count * self.poll_interval) + i
                    self._display_time_progress_on_node(time_completed)
                    time.sleep(1)

            except (LocalNetworkError, ApiServerError) as e:
                # For network-related errors, increment error count and potentially abort
                consecutive_errors += 1
                if consecutive_errors >= max_consecutive_errors:
                    raise Exception(
                        f"Polling aborted after {consecutive_errors} consecutive network errors: {str(e)}"
                    ) from e

                # Log the error but continue polling
                logging.warning(
                    f"Network error during polling (attempt {poll_count}/{self.max_poll_attempts}): {str(e)}. "
                    f"Will retry in {self.poll_interval} seconds."
                )
                time.sleep(self.poll_interval)

            except Exception as e:
                # For other errors, increment count and potentially abort
                consecutive_errors += 1
                if consecutive_errors >= max_consecutive_errors or status == TaskStatus.FAILED:
                    raise Exception(
                        f"Polling aborted after {consecutive_errors} consecutive errors: {str(e)}"
                    ) from e

                logging.error(f"[DEBUG] Polling error: {str(e)}")
                logging.warning(
                    f"Error during polling (attempt {poll_count}/{self.max_poll_attempts}): {str(e)}. "
                    f"Will retry in {self.poll_interval} seconds."
                )
                time.sleep(self.poll_interval)

        # If we've exhausted all polling attempts
        raise Exception(
            f"Polling timed out after {poll_count} attempts ({poll_count * self.poll_interval} seconds). "
            f"The operation may still be running on the server but is taking longer than expected."
        )