"""
Asteroid Agents Python SDK - High-Level Client Interface

Provides a clean, easy-to-use interface for interacting with the Asteroid Agents API,
similar to the TypeScript SDK.

This module provides a high-level client that wraps the generated OpenAPI client
without modifying any generated files.
"""

import time
import os
from typing import Dict, Any, Optional, List, Union, Tuple
from .openapi_client import (
    Configuration,
    ApiClient,
    SDKApi,
    ExecutionApi,
    ExecutionStatusResponse,
    ExecutionResultResponse,
    BrowserSessionRecordingResponse,
    UploadExecutionFiles200Response,
    Status,
    StructuredAgentExecutionRequest
)
from .openapi_client.exceptions import ApiException


class AsteroidClient:
    """
    High-level client for the Asteroid Agents API.
    
    This class provides a convenient interface for executing agents and managing
    their execution lifecycle, similar to the TypeScript SDK.
    """
    
    def __init__(self, api_key: str, base_url: Optional[str] = None):
        """
        Create an API client with the provided API key.
        
        Args:
            api_key: Your API key for authentication
            base_url: Optional base URL (defaults to https://odyssey.asteroid.ai/api/v1)
        
        Example:
            client = AsteroidClient('your-api-key')
        """
        if api_key is None:
            raise TypeError("API key cannot be None")
        
        # Configure the API client
        config = Configuration(
            host=base_url or "https://odyssey.asteroid.ai/api/v1",
            api_key={'ApiKeyAuth': api_key}
        )
        
        self.api_client = ApiClient(config)
        self.sdk_api = SDKApi(self.api_client)
        self.execution_api = ExecutionApi(self.api_client)
        
    def execute_agent(self, agent_id: str, execution_data: Dict[str, Any]) -> str: 
        """
        Execute an agent with the provided parameters.
        
        Args:
            agent_id: The ID of the agent to execute
            agent_profile_id: The ID of the agent profile
            execution_data: The execution parameters
            
        Returns:
            The execution ID
            
        Raises:
            Exception: If the execution request fails
            
        Example:
            execution_id = client.execute_structured_agent('my-agent-id', 'agent-profile-id', {'input': 'some dynamic value'})
        """
        req = StructuredAgentExecutionRequest(dynamic_data=execution_data)
        try:
            response = self.sdk_api.execute_agent_structured(agent_id, req)
            return response.execution_id
        except ApiException as e:
            raise Exception(f"Failed to execute agent: {e}")
    
    def get_execution_status(self, execution_id: str) -> ExecutionStatusResponse:
        """
        Get the current status for an execution.
        
        Args:
            execution_id: The execution identifier
            
        Returns:
            The execution status details
            
        Raises:
            Exception: If the status request fails
            
        Example:
            status = client.get_execution_status(execution_id)
            print(status.status)
        """
        try:
            return self.sdk_api.get_execution_status(execution_id)
        except ApiException as e:
            raise Exception(f"Failed to get execution status: {e}")
    
    def get_execution_result(self, execution_id: str) -> Dict[str, Any]:
        """
        Get the final result of an execution.
        
        Args:
            execution_id: The execution identifier
            
        Returns:
            The result object of the execution
            
        Raises:
            Exception: If the result request fails or execution failed
            
        Example:
            result = client.get_execution_result(execution_id)
            print(result)
        """
        try:
            response = self.sdk_api.get_execution_result(execution_id)
            
            if response.error:
                raise Exception(response.error)
            
            return response.result or {}
        except ApiException as e:
            raise Exception(f"Failed to get execution result: {e}")
    
    def wait_for_execution_result(
        self, 
        execution_id: str, 
        interval: float = 1.0, 
        timeout: float = 3600.0
    ) -> Dict[str, Any]:
        """
        Wait for an execution to reach a terminal state and return the result.
        
        Continuously polls the execution status until it's either "completed", 
        "cancelled", or "failed".
        
        Args:
            execution_id: The execution identifier
            interval: Polling interval in seconds (default is 1.0)
            timeout: Maximum wait time in seconds (default is 3600 - 1 hour)
            
        Returns:
            The execution result if completed
            
        Raises:
            Exception: If the execution ends as "cancelled" or "failed", or times out
            
        Example:
            result = client.wait_for_execution_result(execution_id, interval=2.0)
        """
        start_time = time.time()
        
        while True:
            elapsed_time = time.time() - start_time
            if elapsed_time >= timeout:
                raise Exception(f"Execution {execution_id} timed out after {timeout}s")
            
            status_response = self.get_execution_status(execution_id)
            current_status = status_response.status
            
            if current_status == Status.COMPLETED:
                return self.get_execution_result(execution_id)
            elif current_status in [Status.FAILED, Status.CANCELLED]:
                reason = f" - {status_response.reason}" if status_response.reason else ""
                raise Exception(f"Execution {execution_id} ended with status: {current_status.value}{reason}")
            
            # Wait for the specified interval before polling again
            time.sleep(interval)
    
    def upload_execution_files(
        self, 
        execution_id: str, 
        files: List[Union[bytes, str, Tuple[str, bytes]]],
        default_filename: str = "file.txt"
    ) -> UploadExecutionFiles200Response:
        """
        Upload files to an execution.
        
        Args:
            execution_id: The execution identifier
            files: List of files to upload. Each file can be:
                   - bytes: Raw file content (will use default_filename)
                   - str: File path as string (will read file and use filename)
                   - Tuple[str, bytes]: (filename, file_content) tuple
            default_filename: Default filename to use when file is provided as bytes
            
        Returns:
            The upload response containing message and file IDs
            
        Raises:
            Exception: If the upload request fails
            
        Example:
            # Upload with file content (file should be in your current working directory)
            with open('hello.txt', 'r') as f:
                file_content = f.read()
            
            response = client.upload_execution_files(execution_id, [file_content.encode()])
            print(f"Uploaded files: {response.file_ids}")
            
            # Upload with filename and content
            files = [('hello.txt', file_content.encode())]
            response = client.upload_execution_files(execution_id, files)
            
            # Or create content directly
            hello_content = "Hello World!".encode()
            response = client.upload_execution_files(execution_id, [hello_content])
        """
        try:
            # Process files to ensure proper format
            processed_files = []
            for file_item in files:
                if isinstance(file_item, tuple):
                    # Already in (filename, content) format
                    filename, content = file_item
                    if isinstance(content, str):
                        content = content.encode()
                    processed_files.append((filename, content))
                elif isinstance(file_item, str):
                    # Check if string is a file path that exists, otherwise treat as content
                    if os.path.isfile(file_item):
                        # File path - read the file
                        filename = os.path.basename(file_item)
                        with open(file_item, 'rb') as f:
                            content = f.read()
                        processed_files.append((filename, content))
                    else:
                        # String content - encode and use default filename
                        content = file_item.encode()
                        processed_files.append((default_filename, content))
                elif isinstance(file_item, bytes):
                    # Raw bytes - use default filename
                    processed_files.append((default_filename, file_item))
                else:
                    # Other types - convert to string content and encode
                    content = str(file_item).encode()
                    processed_files.append((default_filename, content))
            
            response = self.execution_api.upload_execution_files(execution_id, files=processed_files)
            return response
        except ApiException as e:
            raise Exception(f"Failed to upload execution files: {e}")
    
    def get_browser_session_recording(self, execution_id: str) -> str:
        """
        Get the browser session recording URL for a completed execution.
        
        Args:
            execution_id: The execution identifier
            
        Returns:
            The URL of the browser session recording
            
        Raises:
            Exception: If the recording request fails
            
        Example:
            recording_url = client.get_browser_session_recording(execution_id)
            print(f"Recording available at: {recording_url}")
        """
        try:
            response = self.sdk_api.get_browser_session_recording(execution_id)
            return response.recording_url
        except ApiException as e:
            raise Exception(f"Failed to get browser session recording: {e}")
    
    def __enter__(self):
        """Context manager entry."""
        return self
    
    def __exit__(self, exc_type, exc_value, traceback):
        """Context manager exit."""
        pass


# Convenience functions that mirror the TypeScript SDK pattern
def create_client(api_key: str, base_url: Optional[str] = None) -> AsteroidClient:
    """
    Create an API client with a provided API key.
    
    This is a convenience function that creates an AsteroidClient instance.
    
    Args:
        api_key: Your API key
        base_url: Optional base URL
        
    Returns:
        A configured AsteroidClient instance
        
    Example:
        client = create_client('your-api-key')
    """
    return AsteroidClient(api_key, base_url)

def execute_agent(client: AsteroidClient, agent_id: str, agent_profile_id: str, execution_data: Dict[str, Any]) -> str:
    """
    Execute an agent with the provided parameters.
    
    Args:
        client: The AsteroidClient instance
        agent_id: The ID of the agent to execute
        agent_profile_id: The ID of the agent profile
        execution_data: The execution parameters
        
    Returns:
        The execution ID
        
    Example:
        execution_id = execute_agent(client, 'my-agent-id', {'input': 'some dynamic value'})
    """
    return client.execute_agent(agent_id, agent_profile_id, execution_data)



def get_execution_status(client: AsteroidClient, execution_id: str) -> ExecutionStatusResponse:
    """
    Get the current status for an execution.
    
    Args:
        client: The AsteroidClient instance
        execution_id: The execution identifier
        
    Returns:
        The execution status details
        
    Example:
        status = get_execution_status(client, execution_id)
        print(status.status)
    """
    return client.get_execution_status(execution_id)


def get_execution_result(client: AsteroidClient, execution_id: str) -> Dict[str, Any]:
    """
    Get the final result of an execution.
    
    Args:
        client: The AsteroidClient instance
        execution_id: The execution identifier
        
    Returns:
        The result object of the execution
        
    Example:
        result = get_execution_result(client, execution_id)
        print(result)
    """
    return client.get_execution_result(execution_id)


def wait_for_execution_result(
    client: AsteroidClient, 
    execution_id: str, 
    interval: float = 1.0, 
    timeout: float = 3600.0
) -> Dict[str, Any]:
    """
    Wait for an execution to reach a terminal state and return the result.
    
    Args:
        client: The AsteroidClient instance
        execution_id: The execution identifier
        interval: Polling interval in seconds (default is 1.0)
        timeout: Maximum wait time in seconds (default is 3600 - 1 hour)
        
    Returns:
        The execution result if completed
        
    Example:
        result = wait_for_execution_result(client, execution_id, interval=2.0)
    """
    return client.wait_for_execution_result(execution_id, interval, timeout)


def upload_execution_files(
    client: AsteroidClient, 
    execution_id: str, 
    files: List[Union[bytes, str, Tuple[str, bytes]]],
    default_filename: str = "file.txt"
) -> UploadExecutionFiles200Response:
    """
    Upload files to an execution.
    
    Args:
        client: The AsteroidClient instance
        execution_id: The execution identifier
        files: List of files to upload
        default_filename: Default filename to use when file is provided as bytes
        
    Returns:
        The upload response containing message and file IDs
        
    Example:
        # Create a simple text file with "Hello World!" content
        hello_content = "Hello World!".encode()
        response = upload_execution_files(client, execution_id, [hello_content])
        print(f"Uploaded files: {response.file_ids}")
        
        # Or specify filename with content
        files = [('hello.txt', "Hello World!".encode())]
        response = upload_execution_files(client, execution_id, files)
    """
    return client.upload_execution_files(execution_id, files, default_filename)


def get_browser_session_recording(client: AsteroidClient, execution_id: str) -> str:
    """
    Get the browser session recording URL for a completed execution.
    
    Args:
        client: The AsteroidClient instance
        execution_id: The execution identifier
        
    Returns:
        The URL of the browser session recording
        
    Example:
        recording_url = get_browser_session_recording(client, execution_id)
        print(f"Recording available at: {recording_url}")
    """
    return client.get_browser_session_recording(execution_id)


# Re-export common types for convenience
__all__ = [
    'AsteroidClient',
    'create_client',
    'execute_agent',
    'get_execution_status',
    'get_execution_result',
    'wait_for_execution_result',
    'upload_execution_files',
    'get_browser_session_recording'
] 