from typing import Optional

import strawberry
from strawberry import UNSET
from strawberry.relay.types import GlobalID
from strawberry.scalars import JSON

from phoenix.server.api.helpers.prompts.models import (
    PromptTemplateFormat,
)
from phoenix.server.api.input_types.GenerativeCredentialInput import GenerativeCredentialInput
from phoenix.server.api.types.Identifier import Identifier

from .ChatCompletionMessageInput import ChatCompletionMessageInput
from .GenerativeModelInput import GenerativeModelInput
from .InvocationParameters import InvocationParameterInput
from .PlaygroundEvaluatorInput import PlaygroundEvaluatorInput
from .PromptTemplateOptions import PromptTemplateOptions


@strawberry.input
class ChatCompletionInput:
    messages: list[ChatCompletionMessageInput]
    model: GenerativeModelInput
    credentials: Optional[list[GenerativeCredentialInput]] = UNSET
    invocation_parameters: list[InvocationParameterInput] = strawberry.field(default_factory=list)
    tools: Optional[list[JSON]] = UNSET
    template: Optional[PromptTemplateOptions] = UNSET
    prompt_name: Optional[Identifier] = None
    repetitions: int
    evaluators: list[PlaygroundEvaluatorInput] = strawberry.field(default_factory=list)


@strawberry.input
class ChatCompletionOverDatasetInput:
    messages: list[ChatCompletionMessageInput]
    model: GenerativeModelInput
    credentials: Optional[list[GenerativeCredentialInput]] = UNSET
    invocation_parameters: list[InvocationParameterInput] = strawberry.field(default_factory=list)
    tools: Optional[list[JSON]] = UNSET
    template_format: PromptTemplateFormat = PromptTemplateFormat.MUSTACHE
    repetitions: int
    dataset_id: GlobalID
    dataset_version_id: Optional[GlobalID] = None
    split_ids: Optional[list[GlobalID]] = None
    experiment_name: Optional[str] = None
    experiment_description: Optional[str] = None
    experiment_metadata: Optional[JSON] = strawberry.field(default_factory=dict)
    prompt_name: Optional[Identifier] = None
    evaluators: list[PlaygroundEvaluatorInput] = strawberry.field(default_factory=list)
    appended_messages_path: Optional[str] = strawberry.field(
        default=None,
        description="Dot-notation path to messages in dataset example input to append to prompt",
    )
    template_variables_path: Optional[str] = strawberry.field(
        default="input",
        description="Dot-notation path prefix for template variables. Default 'input' means "
        "{{query}} resolves to input.query. Empty string means full paths like "
        "{{input.query}} or {{reference.answer}} are required.",
    )
    tracing_enabled: bool = True
