import logging
from typing import Any

from ...contexts import CompletionChoice, GenerationUsage


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

                    choice = CompletionChoice(
                        message={
                            "role": getattr(content, "role", "model"),
                            "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[dict]:
    """Format Google GenAI contents for tracing."""
    try:
        formatted_messages = []

        if isinstance(contents, str):
            formatted_messages.append({
                "role": "user",
                "content": contents,
                "function_call": None,
                "name": None,
                "tool_calls": None,
            })
        elif isinstance(contents, list):
            for item in contents:
                if isinstance(item, str):
                    formatted_messages.append({
                        "role": "user",
                        "content": item,
                        "function_call": None,
                        "name": None,
                        "tool_calls": None,
                    })
                elif isinstance(item, dict):
                    role = item.get("role", "user")
                    parts = item.get("parts", [])

                    text_content = []
                    for part in parts:
                        if isinstance(part, str):
                            text_content.append(part)
                        elif isinstance(part, dict):
                            if "text" in part:
                                text_content.append(part["text"])
                            elif "function_call" in part:
                                text_content.append(f"[Function Call: {part['function_call'].get('name', 'unknown')}]")
                            elif "function_response" in part:
                                text_content.append("[Function Response]")

                    formatted_messages.append({
                        "role": role,
                        "content": "\n".join(text_content) if text_content else "",
                        "function_call": None,
                        "name": None,
                        "tool_calls": None,
                    })
                elif hasattr(item, "role") and hasattr(item, "parts"):
                    role = getattr(item, "role", "user")
                    parts = getattr(item, "parts", [])

                    text_content = []
                    for part in parts:
                        if hasattr(part, "text"):
                            text_content.append(part.text)
                        elif hasattr(part, "function_call"):
                            func_call = part.function_call
                            text_content.append(f"[Function Call: {func_call.name}]")
                        elif hasattr(part, "function_response"):
                            text_content.append("[Function Response]")

                    formatted_messages.append({
                        "role": role,
                        "content": "\n".join(text_content) if text_content else "",
                        "function_call": None,
                        "name": None,
                        "tool_calls": None,
                    })
        elif hasattr(contents, "role") and hasattr(contents, "parts"):
            role = getattr(contents, "role", "user")
            parts = getattr(contents, "parts", [])

            text_content = []
            for part in parts:
                if hasattr(part, "text"):
                    text_content.append(part.text)
                elif hasattr(part, "function_call"):
                    func_call = part.function_call
                    text_content.append(f"[Function Call: {func_call.name}]")
                elif hasattr(part, "function_response"):
                    text_content.append("[Function Response]")

            formatted_messages.append({
                "role": role,
                "content": "\n".join(text_content) if text_content else "",
                "function_call": None,
                "name": None,
                "tool_calls": None,
            })

        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 ""
