from typing import Protocol, Optional, Tuple, TYPE_CHECKING
from .trace_types import TraceParams
from .experiment_types import ExperimentParams
from .experiment_types import Experiment
from .trace_types import Trace
from .generation_types import GenerationParams, Generation
from .log_types import LogParams, Log

if TYPE_CHECKING:
    from ...utils.dtos import CreateExperimentResult

class IMonitorSDK(Protocol):
    """Interface for interacting with Basalt monitoring.

    The MonitorSDK provides methods to create and manage traces, generations, and logs
    for monitoring and tracking AI application flows.

    Examples:
        ```python
        # Example 1: Creating a trace
        trace = basalt.monitor.create_trace('user-session', {
            'input': 'User started a new session',
            'metadata': {'userId': '123', 'sessionType': 'web'}
        })

        # Example 2: Creating a generation within a trace
        generation = basalt.monitor.create_generation({
            'name': 'text-completion',
            'prompt': {'slug': 'text-completion-prompt', 'version': '1.0.0'},
            'input': 'Tell me a joke',
            'trace': trace
        })

        # Example 3: Creating a span for a processing step
        span = basalt.monitor.create_log({
            'type': 'span',
            'name': 'data-processing',
            'trace': trace,
            'metadata': {'processingType': 'text-analysis'}
        })
        ```
    """

    def create_trace(self, slug: str, params: TraceParams = {}) -> Trace:
        """Creates a new trace to monitor a complete user interaction or process flow.

        Args:
            slug: The unique identifier of the feature to which the trace belongs.
            params: Optional parameters for the trace.
                - output: Final output data for the trace.
                - start_time: When the trace started (defaults to now if not provided).
                - end_time: When the trace ended.
                - user: User information associated with this trace.
                - organization: Organization information associated with this trace.
                - metadata: Additional contextual information for the trace.

        Examples:
            ```python
            # Create a basic trace
            basic_trace = basalt.monitor.create_trace('user-query')

            # Create a trace with parameters
            detailed_trace = basalt.monitor.create_trace('document-processing', {
                'input': 'Raw document text',
                'start_time': datetime.now(),
                'user': {'id': 'user-123', 'name': 'John Doe'},
                'metadata': {'documentId': 'doc-456', 'documentType': 'invoice'}
            })
            ```

        Returns:
            A Trace object that can be used to track the process flow.
        """
        ...

    def create_generation(self, params: GenerationParams) -> Generation:
        """Creates a new generation to track an AI model generation within a trace.

        Args:
            params: Parameters for the generation.
                - name: Name of the generation (required).
                - trace: The parent trace this generation belongs to (required).
                - prompt: Information about the prompt used for generation.
                    - slug: Prompt identifier.
                    - version: Prompt version.
                    - tag: Prompt tag.
                - input: The input provided to the model.
                - variables: Variables used in the prompt template.
                - output: The output generated by the model.
                - start_time: When the generation started.
                - end_time: When the generation completed.
                - metadata: Additional contextual information.
                - parent: Optional parent log if this generation is part of a log.

        Examples:
            ```python
            # Create a generation with a prompt reference
            generation = basalt.monitor.create_generation({
                'name': 'answer-generation',
                'trace': trace,
                'prompt': {'slug': 'qa-prompt', 'version': '2.1.0'},
                'input': 'What is the capital of France?',
                'variables': {'style': 'concise', 'language': 'en'},
                'metadata': {'modelVersion': 'gpt-4'}
            })

            # Create a generation without a prompt reference
            simple_generation = basalt.monitor.create_generation({
                'name': 'text-completion',
                'trace': trace,
                'input': 'Complete this sentence: The sky is',
                'output': 'The sky is blue and vast.'
            })
            ```

        Returns:
            A Generation object that can be used to track the AI generation.
        """
        ...

    def create_log(self, params: LogParams) -> Log:
        """Creates a new log to track a specific operation or step within a trace.

        Args:
            params: Parameters for the log.
                - name: Name of the log (required).
                - trace: The parent trace this log belongs to (required).
                - input: The input data for this operation.
                - start_time: When the log started.
                - end_time: When the log completed.
                - metadata: Additional contextual information.
                - parent: Optional parent log if this is a nested log.

        Examples:
            ```python
            # Create a basic log
            basic_log = basalt.monitor.create_log({
                'name': 'data-fetching',
                'trace': trace
            })

            # Create a detailed log
            detailed_log = basalt.monitor.create_log({
                'name': 'user-validation',
                'trace': trace,
                'input': 'user credentials',
                'metadata': {'validationRules': ['password-strength', 'email-format']},
                'parent': parent_log  # Another log this is nested under
            })
            ```

        Returns:
            A Log object that can be used to track the operation.
        """
        ...

    async def create_experiment(self, feature_slug: str, params: ExperimentParams) -> 'CreateExperimentResult':
        """Creates a new experiment to bundle multiple traces together in.

        You can pass this experiment to the create_trace method to add the generated traces to the experiment.
        It's used mostly for local experimentations, to compare the performance between different versions of a workflow.

        Args:
            feature_slug: The unique identifier of the feature to which the experiment belongs.
            params: Parameters for the experiment.
                - name: Name of the experiment (required).

        Examples:
            ```python
            experiment = await basalt.monitor.create_experiment('user-query', {'name': 'my-experiment'})

            # Create a trace and add it to the experiment
            trace = basalt.monitor.create_trace('user-query', {'experiment': experiment})
            ```

        Returns:
            A tuple containing (Optional[Exception], Optional[Experiment]). The Experiment object can be used to track the AI generation.
        """
        ...

    def create_experiment_sync(self, feature_slug: str, params: ExperimentParams) -> 'CreateExperimentResult':
        """Synchronously creates a new experiment to bundle multiple traces together in.

        You can pass this experiment to the create_trace method to add the generated traces to the experiment.
        It's used mostly for local experimentations, to compare the performance between different versions of a workflow.

        Args:
            feature_slug: The unique identifier of the feature to which the experiment belongs.
            params: Parameters for the experiment.
                - name: Name of the experiment (required).

        Examples:
            ```python
            experiment = basalt.monitor.create_experiment_sync('user-query', {'name': 'my-experiment'})

            # Create a trace and add it to the experiment
            trace = basalt.monitor.create_trace('user-query', {'experiment': experiment})
            ```

        Returns:
            A tuple containing (Optional[Exception], Optional[Experiment]). The Experiment object can be used to track the AI generation.
        """
        ...
