import logging
from typing import Any

from ...contexts import CompletionChoice, GenerationUsage, InputMessage, MessagePart
from ...utils import get


def get_finish_reason(response_obj: Any) -> str | None:
    """Extract finish reason from Google GenAI response."""
    try:
        if hasattr(response_obj, "candidates") and response_obj.candidates:
            candidate = response_obj.candidates[0]
            if hasattr(candidate, "finish_reason"):
                finish_reason = str(candidate.finish_reason)
                if "STOP" in finish_reason:
                    return "stop"
                elif "MAX_TOKENS" in finish_reason:
                    return "length"
                elif "SAFETY" in finish_reason or "RECITATION" in finish_reason:
                    return "content_filter"
                else:
                    return "other"
        return None
    except Exception as e:
        logging.warning("Error getting finish reason from Google GenAI", exc_info=e)
        return None


def get_completion_choices(response_obj: Any) -> list[CompletionChoice]:
    """Extract completion choices from Google GenAI response."""
    try:
        choices = []
        if hasattr(response_obj, "candidates") and response_obj.candidates:
            for candidate in response_obj.candidates:
                if hasattr(candidate, "content") and candidate.content:
                    content = candidate.content

                    text_parts = []
                    tool_calls = []

                    if hasattr(content, "parts") and content.parts:
                        for part in content.parts:
                            if hasattr(part, "text") and part.text:
                                text_parts.append(part.text)
                            elif hasattr(part, "function_call"):
                                func_call = part.function_call
                                tool_call = {
                                    "id": getattr(
                                        func_call, "id", f"call_{hash(func_call.name)}"
                                    ),
                                    "type": "function",
                                    "function": {
                                        "name": func_call.name,
                                        "arguments": str(func_call.args),
                                    },
                                }
                                tool_calls.append(tool_call)

                    text_content = "".join(text_parts) if text_parts else None

                    role = getattr(content, "role", "model")
                    if role == "model":
                        role = "assistant"

                    choice = CompletionChoice(
                        message={
                            "role": role,
                            "content": text_content,
                            "function_call": None,
                            "name": None,
                            "tool_calls": tool_calls if tool_calls else None,
                        },
                        delta=None,
                        text=text_content,
                    )
                    choices.append(choice)

        return choices
    except Exception as e:
        logging.warning(
            "Error getting completion choices from Google GenAI", exc_info=e
        )
        return []


def get_usage_from_response(response_obj: Any) -> GenerationUsage:
    """Extract usage information from Google GenAI response."""
    try:
        if hasattr(response_obj, "usage_metadata"):
            usage = response_obj.usage_metadata
            prompt_tokens = getattr(usage, "prompt_token_count", None)
            completion_tokens = getattr(usage, "response_token_count", None)
            total_tokens = getattr(usage, "total_token_count", None)

            return {
                "input_tokens": prompt_tokens,
                "output_tokens": completion_tokens,
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
                "total_tokens": total_tokens,
                "max_tokens": None,
                "temperature": None,
                "top_p": None,
                "stream": None,
                "tool_calls": None,
            }
        return {}
    except Exception as e:
        logging.warning("Error getting usage from Google GenAI response", exc_info=e)
        return {}


def get_request_parameters(kwargs: dict) -> GenerationUsage:
    """Extract request parameters from kwargs."""
    try:
        config = kwargs.get("config")
        max_tokens = None
        temperature = None
        top_p = None

        if config:
            if isinstance(config, dict):
                max_tokens = config.get("max_output_tokens")
                temperature = config.get("temperature")
                top_p = config.get("top_p")
            else:
                if hasattr(config, "max_output_tokens"):
                    max_tokens = config.max_output_tokens
                if hasattr(config, "temperature"):
                    temperature = config.temperature
                if hasattr(config, "top_p"):
                    top_p = config.top_p

        return {
            "max_tokens": max_tokens,
            "temperature": temperature,
            "top_p": top_p,
            "stream": False,
            "input_tokens": None,
            "output_tokens": None,
            "prompt_tokens": None,
            "completion_tokens": None,
            "total_tokens": None,
            "tool_calls": None,
        }
    except Exception as e:
        logging.warning(
            "Error getting request parameters from Google GenAI", exc_info=e
        )
        return {}


def format_messages_for_tracing(contents: Any) -> list[InputMessage]:
    """Format Google GenAI contents for tracing."""
    try:
        formatted_messages: list[InputMessage] = []

        def _build_tool_call_part(part: dict) -> MessagePart:
            func_call = get(part, "function_call", {})
            func_args = get(func_call, "args", {})
            func_name = get(func_call, "name", "unknown")

            tool_call_part = MessagePart(
                type="tool_call",
                name=func_name,
                arguments=str(func_args),
            )

            tool_call_part["id"] = get(func_call, "id", None)
            return tool_call_part

        def _get_response_part(item: Any) -> list[InputMessage]:
            parts = get(item, "parts", [])
            role = get(item, "role", "user")

            if role == "model":
                role = "assistant"

            return _process_parts(parts, role)

        def _get_tool_response_part(part: Any) -> MessagePart:
            func_response = get(part, "function_response", {})
            response_data = get(func_response, "response", {})

            tool_response_part = MessagePart(
                type="tool_call_response",
                response=response_data,
            )

            tool_response_part["id"] = get(func_response, "id", None)
            return tool_response_part

        def has_parts(item: Any) -> bool:
            return isinstance(item, dict) or (
                hasattr(item, "role") and hasattr(item, "parts")
            )

        def _process_parts(parts: list, role: str) -> list[InputMessage]:
            messages: list[InputMessage] = []
            text_parts: list[MessagePart] = []
            tool_call_parts: list[MessagePart] = []
            tool_response_parts: list[MessagePart] = []

            for part in parts:
                if isinstance(part, str):
                    text_parts.append(MessagePart(type="text", content=part))
                elif isinstance(part, dict):
                    if "text" in part:
                        text_parts.append(
                            MessagePart(type="text", content=part["text"])
                        )
                    elif "function_call" in part:
                        tool_call_parts.append(_build_tool_call_part(part))
                    elif "function_response" in part:
                        tool_response_parts.append(_get_tool_response_part(part))
                elif hasattr(part, "text") and part.text:
                    text_parts.append(MessagePart(type="text", content=part.text))
                elif hasattr(part, "function_call") and part.function_call:
                    tool_call_parts.append(_build_tool_call_part(part))
                elif hasattr(part, "function_response") and part.function_response:
                    tool_response_parts.append(_get_tool_response_part(part))

            if text_parts or tool_call_parts:
                combined_parts = text_parts + tool_call_parts
                messages.append(InputMessage(role=role, parts=combined_parts))
            elif not tool_response_parts:
                messages.append(
                    InputMessage(
                        role=role, parts=[MessagePart(type="text", content="")]
                    )
                )

            if tool_response_parts:
                for tool_response_part in tool_response_parts:
                    messages.append(
                        InputMessage(role="tool", parts=[tool_response_part])
                    )

            return messages

        # Recursively process the contents using helpers above
        # Leave text-note here as visual block between internals
        if isinstance(contents, str):
            formatted_messages.append(
                InputMessage(
                    role="user", parts=[MessagePart(type="text", content=contents)]
                )
            )
        elif isinstance(contents, list):
            for item in contents:
                if isinstance(item, str):
                    formatted_messages.append(
                        InputMessage(
                            role="user", parts=[MessagePart(type="text", content=item)]
                        )
                    )
                elif has_parts(item):
                    formatted_messages.extend(_get_response_part(item))
        elif has_parts(contents):
            formatted_messages.extend(_get_response_part(contents))

        return formatted_messages
    except Exception as e:
        logging.warning("Error formatting messages for tracing", exc_info=e)
        return []


def get_stream_text(chunk: Any) -> str:
    """Extract text from streaming chunk."""
    try:
        if hasattr(chunk, "text") and chunk.text:
            return chunk.text

        if hasattr(chunk, "candidates") and chunk.candidates:
            candidate = chunk.candidates[0]
            if hasattr(candidate, "content") and candidate.content:
                content = candidate.content
                if hasattr(content, "parts") and content.parts:
                    text_parts = []
                    for part in content.parts:
                        if hasattr(part, "text") and part.text:
                            text_parts.append(part.text)
                    return "".join(text_parts)

        return ""
    except Exception:
        return ""
