from __future__ import annotations

import json
import logging
from typing import Any

from ..utils import get
from .span import SpanContext
from .types import (
    CompletionChoice,
    ContextType,
    GenerationUsage,
    InputMessage,
    ObservationType,
)
from .utils import (
    convert_completion_choice_to_output_message,
    normalize_finish_reason,
)


class GenerationContext(SpanContext):
    """Context for generation spans (LLM operations)."""

    observation_type: ObservationType = "observation"
    type: ContextType = "generation"

    def __init__(self, *args, **kwargs):
        """Initialize the generation context with a special attribute."""
        super().__init__(*args, **kwargs)
        self.operation(self.observation_type)

    def completions(
        self,
        completion_choices: list[CompletionChoice],
    ) -> GenerationContext:
        """Trace output messages using gen_ai.output.messages specification.

        Converts CompletionChoice format to new OutputMessage format with parts.
        Sets gen_ai.output.messages as a JSON array of all output messages.
        """
        try:
            if not completion_choices:
                return self

            output_messages = [
                convert_completion_choice_to_output_message(
                    choice,
                    get(choice, "finish_reason"),
                )
                for choice in completion_choices
                if choice is not None
            ]

            if output_messages:
                self._span.set_attribute(
                    "gen_ai.output.messages", json.dumps(output_messages)
                )

            return self
        except Exception as e:
            logging.warning("Completions could not be traced", exc_info=e)
            return self

    def input(self, content: Any) -> GenerationContext:
        """Set the input for the generation."""
        self._span.set_attribute("gen_ai.prompt", str(content))
        return self

    def messages(self, messages: list[InputMessage]) -> GenerationContext:
        """Trace input messages using gen_ai.input.messages specification.

        Sets gen_ai.input.messages as a JSON array of all input messages.
        """
        try:
            if not messages:
                return self

            input_messages = [
                message for message in messages if message is not None
            ]

            if input_messages:
                self._span.set_attribute(
                    "gen_ai.input.messages", json.dumps(input_messages)
                )

            return self
        except Exception as e:
            logging.warning("Messages could not be traced", exc_info=e)
            return self

    def model(self, model_name: str) -> GenerationContext:
        """Set the desired model for generation."""
        self._span.set_attribute("gen_ai.request.model", model_name)
        return self

    def model_used(self, model_used: str) -> GenerationContext:
        """Set the actual model used for generation."""
        self._span.set_attribute("gen_ai.response.model", model_used)
        return self

    def operation(self, operation_name: str) -> GenerationContext:
        """Set the operation name for the generation."""
        self._span.set_attribute("gen_ai.operation.name", operation_name)
        return self

    def output(self, content: Any) -> GenerationContext:
        """Set the output of the generation."""
        self._span.set_attribute("gen_ai.response.completion", str(content))
        return self

    def provider(self, provider_name: str) -> GenerationContext:
        """Set the provider name for the generation."""
        self._span.set_attribute("gen_ai.provider.name", provider_name)
        return self

    def reason(self, reason: str) -> GenerationContext:
        """Set the reason for the generation."""
        normalized = normalize_finish_reason(reason)
        if normalized:
            self._span.set_attribute("gen_ai.response.finish_reasons", [normalized])
        return self

    def request_params(self, usage: GenerationUsage) -> GenerationContext:
        """Set request parameters from usage dict."""
        if usage.get("max_tokens") is not None:
            self._span.set_attribute("gen_ai.request.max_tokens", usage["max_tokens"])
        if usage.get("temperature") is not None:
            self._span.set_attribute("gen_ai.request.temperature", usage["temperature"])
        if usage.get("top_p") is not None:
            self._span.set_attribute("gen_ai.request.top_p", usage["top_p"])
        if usage.get("top_k") is not None:
            self._span.set_attribute("gen_ai.request.top_k", usage["top_k"])
        return self

    def response_id(self, response_id: str) -> GenerationContext:
        """Set the response ID for the generation."""
        self._span.set_attribute("gen_ai.response.id", response_id)
        return self

    def usage(
        self,
        usage: GenerationUsage,
    ) -> GenerationContext:
        """Set token usage metrics."""
        self.set_attributes(
            prefix="gen_ai.usage",
            **usage,
        )

        return self

    def vendor(self, vendor: str) -> GenerationContext:
        """Set the vendor for the generation."""
        self._span.set_attribute("gen_ai.request.vendor", vendor)
        return self
