from enum import Enum
from typing import Annotated, Optional, Union

import strawberry
from strawberry.scalars import JSON


@strawberry.enum
class CanonicalParameterName(str, Enum):
    TEMPERATURE = "temperature"
    MAX_COMPLETION_TOKENS = "max_completion_tokens"
    STOP_SEQUENCES = "stop_sequences"
    TOP_P = "top_p"
    RANDOM_SEED = "random_seed"
    RESPONSE_FORMAT = "response_format"
    REASONING_EFFORT = "reasoning_effort"
    ANTHROPIC_EXTENDED_THINKING = "anthropic_extended_thinking"


@strawberry.enum
class InvocationInputField(str, Enum):
    value_int = "value_int"
    value_float = "value_float"
    value_bool = "value_bool"
    value_string = "value_string"
    value_json = "value_json"
    value_string_list = "value_string_list"
    value_boolean = "value_boolean"


@strawberry.interface
class InvocationParameterBase:
    invocation_name: str
    canonical_name: Optional[CanonicalParameterName] = None
    label: str
    required: bool = False


@strawberry.type
class IntInvocationParameter(InvocationParameterBase):
    invocation_input_field: InvocationInputField = InvocationInputField.value_int
    default_value: Optional[int] = None


@strawberry.type
class FloatInvocationParameter(InvocationParameterBase):
    invocation_input_field: InvocationInputField = InvocationInputField.value_float
    default_value: Optional[float] = None


@strawberry.type
class BoundedFloatInvocationParameter(InvocationParameterBase):
    invocation_input_field: InvocationInputField = InvocationInputField.value_float
    default_value: Optional[float] = None
    min_value: float
    max_value: float


@strawberry.type
class StringInvocationParameter(InvocationParameterBase):
    invocation_input_field: InvocationInputField = InvocationInputField.value_string
    default_value: Optional[str] = None


@strawberry.type
class JSONInvocationParameter(InvocationParameterBase):
    invocation_input_field: InvocationInputField = InvocationInputField.value_json
    default_value: Optional[JSON] = None


@strawberry.type
class StringListInvocationParameter(InvocationParameterBase):
    invocation_input_field: InvocationInputField = InvocationInputField.value_string_list
    default_value: Optional[list[str]] = None


@strawberry.type
class BooleanInvocationParameter(InvocationParameterBase):
    invocation_input_field: InvocationInputField = InvocationInputField.value_bool
    default_value: Optional[bool] = None


# Create the union for output types
InvocationParameter = Annotated[
    Union[
        IntInvocationParameter,
        FloatInvocationParameter,
        BoundedFloatInvocationParameter,
        StringInvocationParameter,
        JSONInvocationParameter,
        StringListInvocationParameter,
        BooleanInvocationParameter,
    ],
    strawberry.union("InvocationParameter"),
]
