import logging
from typing import Any

from opentelemetry import trace
from opentelemetry.trace import SpanContext, SpanKind, TraceFlags

from ..utils import format_hex_id, get, safe_json_loads, uses_pydantic_base_model
from .base import BaseContext
from .types import (
    FINISH_REASON_MAP,
    VALID_FINISH_REASONS,
    CompletionChoice,
    CompletionChoiceMessage,
    MessagePart,
    OutputMessage,
)


def _extract_parts_from_message(message: dict) -> list[MessagePart]:
    """Extract message parts from a message dict.

    Handles content, function_call, and tool_calls fields consistently.
    Also handles tool role messages with tool_call_response parts.
    """
    parts: list[MessagePart] = []
    role = get(message, "role", "user")

    if role == "tool":
        content = get_content_from_choice_message(message)
        response = safe_json_loads(content, default=content if content is not None else {})

        tool_response_part = MessagePart(
            type="tool_call_response",
            response=response,
        )
        if tool_id := get(message, "id"):
            tool_response_part["id"] = tool_id
        parts.append(tool_response_part)
        return parts

    if content := get_content_from_choice_message(message):
        if isinstance(content, list):
            for part in content:
                parts.append(MessagePart(type="text", content=part))
        else:
            parts.append(MessagePart(type="text", content=content))

    func_args, func_name = get_function_call(message)
    if func_name:
        part = MessagePart(type="tool_call", name=func_name)
        if func_args:
            part["arguments"] = func_args
        parts.append(part)

    for tool_call in get_tool_call_from_choice_message(message):
        tool_func_args, tool_func_name = get_function_call(tool_call, "function")

        if tool_func_name:
            part = MessagePart(type="tool_call", name=tool_func_name)
            if tool_id := get(tool_call, "id"):
                part["id"] = tool_id
            if tool_func_args:
                part["arguments"] = tool_func_args
            parts.append(part)

    return parts or [MessagePart(type="text", content="")]


def convert_completion_choice_to_output_message(
    choice: CompletionChoice,
    finish_reason: str | None = None,
) -> OutputMessage:
    """Convert CompletionChoice to new OutputMessage format with parts."""
    message = get_choice_message(choice)

    if message:
        parts = _extract_parts_from_message(message)
        role = get(message, "role", "assistant")
        message_name = get(message, "name")
    elif text := get(choice, "text", ""):
        parts = [MessagePart(type="text", content=text)]
        role = "assistant"
        message_name = None
    else:
        parts = [MessagePart(type="text", content="")]
        role = "assistant"
        message_name = None

    return OutputMessage(
        role=role,
        parts=parts,
        finish_reason=normalize_finish_reason(finish_reason) or "stop",
        name=message_name,
    )


def create_span_in_existing_trace(trace_id_hex: str, parent_span_id_hex: str):
    """Create a span without using context manager."""
    if not isinstance(trace_id_hex, str):
        trace_id_hex = str(trace_id_hex)
    if not isinstance(parent_span_id_hex, str):
        parent_span_id_hex = str(parent_span_id_hex)

    trace_id = format_hex_id(trace_id_hex)
    parent_span_id = format_hex_id(parent_span_id_hex)

    parent_span_context = SpanContext(
        trace_id=trace_id,
        span_id=parent_span_id,
        is_remote=True,
        trace_flags=TraceFlags(0x01),
    )

    parent_span = trace.NonRecordingSpan(parent_span_context)
    ctx = trace.set_span_in_context(parent_span)

    tracer = trace.get_tracer(__name__)
    span = tracer.start_span("llm_generation", context=ctx, kind=SpanKind.INTERNAL)
    return span


def get_choice_message(
    choice: CompletionChoice | None,
) -> CompletionChoiceMessage | None:
    """Get message dict from choice"""
    if uses_pydantic_base_model(choice):
        choice = choice.model_dump()

    if hasattr(choice, "message") and choice.message:
        if uses_pydantic_base_model(choice.message):
            return choice.message.model_dump()
        return choice.message
    elif hasattr(choice, "delta") and choice.delta:
        return choice.delta
    elif isinstance(choice, dict):
        return choice.get("message") or choice.get("delta") or None
    return None


def get_content_from_choice_message(message: CompletionChoiceMessage | None) -> Any:
    """Get the content from the choice message"""
    content = get(message, "content")

    if content is None:
        return ""
    elif isinstance(content, list):
        text_parts = []
        for part in content:
            text = get(part, "text") if isinstance(part, dict) else str(part)
            if text:
                text_parts.append(text)
        return "\n".join(text_parts)
    elif not isinstance(content, str):
        return str(content)

    return content


def get_current_context(cls: type[BaseContext]) -> BaseContext:
    def _get_current_context():
        """Get the current context from the current span"""
        current_span = trace.get_current_span()

        if (
            current_span
            and current_span.get_span_context().is_valid
            and hasattr(current_span, "_attributes")
        ):
            return cls(current_span)
        else:
            logging.error(
                "No valid span found to create context from",
                extra={"current_span": current_span},
            )

            raise ValueError("Cannot create context from current span")

    return _get_current_context


def get_function_call(
    message: Any | None,
    default_name: str = "function_call",
) -> tuple[str | None, str | None]:
    """Get function call details from message"""
    func_call = get(message, default_name)
    return (
        (get(func_call, "arguments"), get(func_call, "name"))
        if func_call
        else (None, None)
    )


def get_tool_call_from_choice_message(
    message: CompletionChoiceMessage | None,
) -> list[dict]:
    """Get tools calls from message if available"""
    tool_calls = get(message, "tool_calls", [])
    return tool_calls if isinstance(tool_calls, list) else []


def normalize_finish_reason(reason: str | None) -> str | None:
    """Normalize vendor-specific finish reasons to OpenTelemetry spec values."""
    if not reason:
        return None

    normalized = FINISH_REASON_MAP.get(reason, reason)

    if normalized not in VALID_FINISH_REASONS:
        return "other"

    return normalized
