from typing import Dict, Optional, Union, Any, TypedDict
from dataclasses import dataclass, field

from .base_log_types import BaseLog, BaseLogParams, LogType, Input, Output

class _PromptReferenceRequired(TypedDict):
    """Required fields for PromptReference."""
    slug: str

class PromptReference(_PromptReferenceRequired, TypedDict, total=False):
    """Reference to a prompt template.

    This class represents a reference to a prompt template used in AI model generations.

    Attributes:
        slug (str): Unique identifier for the prompt template (required).
        version (str): Version of the prompt template (optional).
        tag (str): Tag for the prompt template (optional).

    Example:
        ```python
        # Basic prompt reference
        prompt = PromptReference(slug="qa-prompt", version="2.1.0")
        ```
    """
    version: Optional[str]
    tag: Optional[str]

class GenerationParams(BaseLogParams, total=False):
    """Parameters for creating a new generation.

    This class defines the parameters that can be used to create a new generation,
    either with or without a prompt reference.

    Attributes:
        prompt (Optional[PromptReference]): Reference to the prompt template used.
        input (Optional[str]): The input provided to the model.
        output (Optional[str]): The output generated by the model.
        variables (Optional[Dict[str, Any]]): Variables used in the prompt template.
        options (Optional[Dict[str, Any]]): Additional options for the generation.

    Example:
        ```python
        # Create generation parameters with a prompt reference
        params = GenerationParams(
            name="answer-generation",
            prompt=PromptReference(slug="qa-prompt", version="2.1.0"),
            input="What is the capital of France?",
            variables={"style": "concise", "language": "en"}
        )

        # Create generation parameters without a prompt reference
        params = GenerationParams(
            name="text-completion",
            input="Complete this sentence: The sky is",
            output="The sky is blue and vast."
        )
        ```
    """
    prompt: Optional[PromptReference]
    variables: Optional[Dict[str, Any]]
    options: Optional[Dict[str, Any]]
    input_tokens: Optional[int]
    output_tokens: Optional[int]
    cost: Optional[float]

class UpdateGenerationParams(GenerationParams, total=False):
    """Parameters for updating a generation."""
    name: Optional[str]

@dataclass
class Generation(BaseLog):
    """Generation class representing an AI model generation within a trace.

    This class tracks interactions with AI models, including inputs, outputs,
    and prompt information used for the generation.

    Attributes:
        prompt (Optional[PromptReference]): Reference to the prompt template used.
        input (Optional[str]): The input provided to the model.
        output (Optional[str]): The output generated by the model.
        variables (Optional[Dict[str, Any]]): Variables used in the prompt template.
        type (str): The type of log, defaults to LogType.GENERATION.
        input_tokens (Optional[int]): Number of tokens used for the input.
        output_tokens (Optional[int]): Number of tokens used for the output.
        cost (Optional[float]): Cost of the generation.

    Example:
        ```python
        # Create a generation with a prompt reference
        generation = trace.create_generation(
            name="answer-generation",
            prompt=PromptReference(slug="qa-prompt", version="2.1.0"),
            input="What is the capital of France?"
        )

        # Start the generation
        generation.start()

        # End the generation with output
        generation.end("The capital of France is Paris.")

        # Update generation metadata
        generation.update(metadata={
            "model_version": "gpt-4",
            "tokens_used": 42
        })
        ```
    """
    prompt: Optional[PromptReference] = None
    variables: Optional[Dict[str, Any]] = None
    type: LogType = field(default=LogType.GENERATION)
    input_tokens: Optional[int] = None
    output_tokens: Optional[int] = None
    cost: Optional[float] = None

    def start(self, input: Optional[Input] = None) -> 'Generation':
        """Marks the generation as started and sets the input if provided.

        Args:
            input (Optional[str]): Optional input data to associate with the generation.

        Returns:
            Generation: The generation instance for method chaining.

        Example:
            ```python
            # Start a generation without input
            generation.start()

            # Start a generation with input
            generation.start("What is the capital of France?")
            ```
        """
        ...

    def end(self,
            output: Optional[Output] = None,
            input_tokens: Optional[int] = None,
            output_tokens: Optional[int] = None,
            cost: Optional[float] = None
    ) -> 'Generation':
        """Marks the generation as ended and sets the output if provided.

        Args:
            output (Optional[Union[str, Dict[str, Any]]]): Optional output data from the model.
                Can be either a string or a dictionary containing output parameters.
            input_tokens (Optional[int]): Optional number of tokens used for the input.
            output_tokens (Optional[int]): Optional number of tokens used for the output.
            cost (Optional[float]): Cost of the generation.

        Returns:
            Generation: The generation instance for method chaining.

        Example:
            ```python
            # End a generation without output
            generation.end()

            # End a generation with output as string
            generation.end("The capital of France is Paris.")

            # End a generation with output params
            generation.end(
                output="The capital of France is Paris.",
                input_tokens=10,
                output_tokens=10,
                cost=0.01
            })
            ```
        """
        ...

    def update(self, params: 'UpdateGenerationParams') -> 'Generation':
        """Updates the log with new parameters.

        Args:
            **params: The parameters to update.

        Returns:
            The log instance for method chaining.
        """
        ...
