from __future__ import annotations

import base64
import functools
import logging
from collections.abc import Mapping
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Optional,
    TypeVar,
    Union,
)

from typing_extensions import TypedDict

from langsmith import client as ls_client
from langsmith import run_helpers
from langsmith._internal._beta_decorator import warn_beta
from langsmith._internal._orjson import dumps as _dumps
from langsmith.schemas import InputTokenDetails, OutputTokenDetails, UsageMetadata

if TYPE_CHECKING:
    from google import genai  # type: ignore[import-untyped, attr-defined]

C = TypeVar("C", bound=Union["genai.Client", Any])
logger = logging.getLogger(__name__)


def _strip_none(d: dict) -> dict:
    """Remove `None` values from dictionary."""
    return {k: v for k, v in d.items() if v is not None}


def _convert_config_for_tracing(kwargs: dict) -> None:
    """Convert `GenerateContentConfig` to `dict` for LangSmith compatibility."""
    if "config" in kwargs and not isinstance(kwargs["config"], dict):
        kwargs["config"] = vars(kwargs["config"])


def _to_dict(obj: Any) -> Any:
    """Serialize a Pydantic/model object to dict (or return as-is for dict/str)."""
    if isinstance(obj, (dict, str)):
        return obj
    if hasattr(obj, "model_dump"):
        return obj.model_dump()
    if hasattr(obj, "to_dict"):
        return obj.to_dict()
    return obj


def _process_gemini_inputs(inputs: dict) -> dict:
    r"""Process Gemini inputs to normalize them for LangSmith tracing.

    Example:
        ```txt
        {"contents": "Hello", "model": "gemini-pro"}
        → {"messages": [{"role": "user", "content": "Hello"}], "model": "gemini-pro"}
        {"contents": [{"role": "user", "parts": [{"text": "What is AI?"}]}], "model": "gemini-pro"}
        → {"messages": [{"role": "user", "content": "What is AI?"}], "model": "gemini-pro"}
        ```
    """  # noqa: E501
    # If contents is not present or not in list format, return as-is
    contents = inputs.get("contents")
    if not contents:
        return inputs

    # Handle string input (simple case)
    if isinstance(contents, str):
        return {
            "messages": [{"role": "user", "content": contents}],
            "model": inputs.get("model"),
            **({k: v for k, v in inputs.items() if k not in ("contents", "model")}),
        }

    # Handle list of content objects (multimodal case)
    if isinstance(contents, list):
        # Check if it's a simple list of strings
        if all(isinstance(item, str) for item in contents):
            # Each string becomes a separate user message (matches Gemini's behavior)
            return {
                "messages": [{"role": "user", "content": item} for item in contents],
                "model": inputs.get("model"),
                **({k: v for k, v in inputs.items() if k not in ("contents", "model")}),
            }
        # Handle complex multimodal case (dict or types.Content / types.Part objects)
        messages = []
        for content in contents:
            content = _to_dict(content)
            if not isinstance(content, dict):
                continue
            role = content.get("role", "user")
            raw_parts = content.get("parts", [])

            text_parts: list[str] = []
            content_parts: list[dict[str, Any]] = []

            for part in raw_parts:
                part = _to_dict(part)
                if isinstance(part, str):
                    text_parts.append(part)
                    content_parts.append({"type": "text", "text": part})
                    continue
                if not isinstance(part, dict):
                    continue
                if "text" in part and part["text"]:
                    text_parts.append(part["text"])
                    content_parts.append({"type": "text", "text": part["text"]})
                elif "inline_data" in part:
                    inline_data = _to_dict(part["inline_data"])
                    if not isinstance(inline_data, dict):
                        continue
                    mime_type = inline_data.get("mime_type", "image/jpeg")
                    data = inline_data.get("data", b"")

                    if isinstance(data, bytes):
                        data_b64 = base64.b64encode(data).decode("utf-8")
                    else:
                        data_b64 = data

                    content_parts.append(
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:{mime_type};base64,{data_b64}",
                                "detail": "high",
                            },
                        }
                    )
                elif "functionResponse" in part:
                    function_response = part["functionResponse"]
                    content_parts.append(
                        {
                            "type": "function_response",
                            "function_response": {
                                "name": function_response.get("name"),
                                "response": function_response.get("response", {}),
                            },
                        }
                    )
                elif "function_call" in part or "functionCall" in part:
                    fc = _to_dict(part.get("function_call") or part.get("functionCall"))
                    if isinstance(fc, dict):
                        content_parts.append(
                            {
                                "type": "function_call",
                                "function_call": {
                                    "id": fc.get("id"),
                                    "name": fc.get("name"),
                                    "arguments": fc.get("args", {}),
                                },
                            }
                        )

            # If only text parts, use simple string format
            if content_parts and all(p.get("type") == "text" for p in content_parts):
                message_content: Union[str, list[dict[str, Any]]] = "\n".join(
                    text_parts
                )
            else:
                message_content = content_parts if content_parts else ""

            messages.append({"role": role, "content": message_content})

        return {
            "messages": messages,
            "model": inputs.get("model"),
            **({k: v for k, v in inputs.items() if k not in ("contents", "model")}),
        }

    # Fallback: return original inputs
    return inputs


def _infer_invocation_params(
    prepopulated_invocation_params: dict, kwargs: dict
) -> dict:
    """Extract invocation parameters for tracing."""
    stripped = _strip_none(kwargs)
    config = stripped.get("config", {})

    # Handle both dict config and GenerateContentConfig object
    if hasattr(config, "temperature"):
        temperature = config.temperature
        max_tokens = getattr(config, "max_output_tokens", None)
        stop = getattr(config, "stop_sequences", None)
    else:
        temperature = config.get("temperature")
        max_tokens = config.get("max_output_tokens")
        stop = config.get("stop_sequences")

    return {
        "ls_provider": "google",
        "ls_model_type": "chat",
        "ls_model_name": stripped.get("model"),
        "ls_temperature": temperature,
        "ls_max_tokens": max_tokens,
        "ls_stop": stop,
        "ls_invocation_params": prepopulated_invocation_params,
    }


def _create_usage_metadata(gemini_usage_metadata: dict) -> UsageMetadata:
    """Convert Gemini usage metadata to LangSmith format."""
    prompt_token_count = gemini_usage_metadata.get("prompt_token_count") or 0
    candidates_token_count = gemini_usage_metadata.get("candidates_token_count") or 0
    cached_content_token_count = (
        gemini_usage_metadata.get("cached_content_token_count") or 0
    )
    thoughts_token_count = gemini_usage_metadata.get("thoughts_token_count") or 0
    total_token_count = (
        gemini_usage_metadata.get("total_token_count")
        or prompt_token_count + candidates_token_count
    )

    input_token_details: dict = {}
    if cached_content_token_count:
        input_token_details["cache_read"] = min(cached_content_token_count, 200000)
        input_token_details["cache_read_over_200k"] = max(
            0, cached_content_token_count - 200000
        )
        input_token_details["over_200k"] = max(
            0, prompt_token_count - cached_content_token_count - 200000
        )

    output_token_details: dict = {}
    if thoughts_token_count:
        output_token_details["reasoning"] = thoughts_token_count

    if candidates_token_count:
        output_token_details["over_200k"] = max(0, candidates_token_count - 200000)

    return UsageMetadata(
        input_tokens=prompt_token_count,
        output_tokens=candidates_token_count,
        total_tokens=total_token_count,
        input_token_details=InputTokenDetails(
            **{k: v for k, v in input_token_details.items() if v is not None}
        ),
        output_token_details=OutputTokenDetails(
            **{k: v for k, v in output_token_details.items() if v is not None}
        ),
    )


def _process_generate_content_response(response: Any) -> dict:
    """Process Gemini response for tracing."""
    try:
        # Convert response to dictionary
        if hasattr(response, "to_dict"):
            rdict = response.to_dict()
        elif hasattr(response, "model_dump"):
            rdict = response.model_dump()
        else:
            rdict = {"text": getattr(response, "text", str(response))}

        # Extract content from candidates if available
        content_result = ""
        content_parts = []
        finish_reason: Optional[str] = None
        if "candidates" in rdict and rdict["candidates"]:
            candidate = rdict["candidates"][0]
            if "content" in candidate:
                content = candidate["content"]
                if "parts" in content and content["parts"]:
                    for part in content["parts"]:
                        # Handle text parts
                        if "text" in part and part["text"]:
                            content_result += part["text"]
                            content_parts.append({"type": "text", "text": part["text"]})
                        # Handle inline data (images) in response
                        elif "inline_data" in part and part["inline_data"] is not None:
                            inline_data = part["inline_data"]
                            mime_type = inline_data.get("mime_type", "image/jpeg")
                            data = inline_data.get("data", b"")

                            # Convert bytes to base64 string if needed
                            if isinstance(data, bytes):
                                data_b64 = base64.b64encode(data).decode("utf-8")
                            else:
                                data_b64 = data  # Already a string

                            content_parts.append(
                                {
                                    "type": "image_url",
                                    "image_url": {
                                        "url": f"data:{mime_type};base64,{data_b64}",
                                        "detail": "high",
                                    },
                                }
                            )
                        elif "function_call" in part or "functionCall" in part:
                            fc = _to_dict(
                                part.get("function_call") or part.get("functionCall")
                            )
                            if isinstance(fc, dict):
                                content_parts.append(
                                    {
                                        "type": "function_call",
                                        "function_call": {
                                            "id": fc.get("id"),
                                            "name": fc.get("name"),
                                            "arguments": fc.get("args", {}),
                                        },
                                    }
                                )
                if "finish_reason" in candidate and candidate["finish_reason"]:
                    finish_reason = candidate["finish_reason"]
        elif "text" in rdict:
            content_result = rdict["text"]
            content_parts.append({"type": "text", "text": content_result})

        # Build chat-like response format - use OpenAI-compatible format for tool calls
        tool_calls = [p for p in content_parts if p.get("type") == "function_call"]
        if tool_calls:
            # OpenAI-compatible format for LangSmith UI
            result = {
                "content": content_result or None,
                "role": "assistant",
                "finish_reason": finish_reason,
                "tool_calls": [
                    {
                        "id": tc["function_call"].get("id") or f"call_{i}",
                        "type": "function",
                        "index": i,
                        "function": {
                            "name": tc["function_call"]["name"],
                            "arguments": _dumps(
                                tc["function_call"]["arguments"]
                            ).decode(),
                        },
                    }
                    for i, tc in enumerate(tool_calls)
                ],
            }
        elif len(content_parts) > 1 or (
            content_parts and content_parts[0]["type"] != "text"
        ):
            # Use structured format for mixed non-tool content
            result = {
                "content": content_parts,
                "role": "assistant",
                "finish_reason": finish_reason,
            }
        else:
            # Use simple string format for text-only responses
            result = {
                "content": content_result,
                "role": "assistant",
                "finish_reason": finish_reason,
            }

        # Extract and convert usage metadata
        usage_metadata = rdict.get("usage_metadata")
        usage_dict: UsageMetadata = UsageMetadata(
            input_tokens=0, output_tokens=0, total_tokens=0
        )
        if usage_metadata:
            usage_dict = _create_usage_metadata(usage_metadata)

        # Return in a format that avoids stringification by LangSmith
        if result.get("tool_calls"):
            # For responses with tool calls, return structured format
            return {
                "content": result["content"],
                "role": "assistant",
                "finish_reason": finish_reason,
                "tool_calls": result["tool_calls"],
                "usage_metadata": usage_dict,
            }
        else:
            # For simple text responses, return minimal structure with usage metadata
            if isinstance(result["content"], str):
                return {
                    "content": result["content"],
                    "role": "assistant",
                    "finish_reason": finish_reason,
                    "usage_metadata": usage_dict,
                }
            else:
                # For multimodal content, return structured format with usage metadata
                return {
                    "content": result["content"],
                    "role": "assistant",
                    "finish_reason": finish_reason,
                    "usage_metadata": usage_dict,
                }
    except Exception as e:
        logger.debug(f"Error processing Gemini response: {e}")
        return {"output": response}


def _reduce_generate_content_chunks(all_chunks: list) -> dict:
    """Reduce streaming chunks into a single response."""
    if not all_chunks:
        return {
            "content": "",
            "usage_metadata": UsageMetadata(
                input_tokens=0, output_tokens=0, total_tokens=0
            ),
        }

    # Accumulate text from all chunks
    full_text = ""
    last_chunk = None

    for chunk in all_chunks:
        try:
            if hasattr(chunk, "text") and chunk.text:
                full_text += chunk.text
            last_chunk = chunk
        except Exception as e:
            logger.debug(f"Error processing chunk: {e}")

    # Extract usage metadata from the last chunk
    usage_metadata: UsageMetadata = UsageMetadata(
        input_tokens=0, output_tokens=0, total_tokens=0
    )
    if last_chunk:
        try:
            if hasattr(last_chunk, "usage_metadata") and last_chunk.usage_metadata:
                if hasattr(last_chunk.usage_metadata, "to_dict"):
                    usage_dict = last_chunk.usage_metadata.to_dict()
                elif hasattr(last_chunk.usage_metadata, "model_dump"):
                    usage_dict = last_chunk.usage_metadata.model_dump()
                else:
                    usage_dict = {
                        "prompt_token_count": getattr(
                            last_chunk.usage_metadata, "prompt_token_count", 0
                        ),
                        "candidates_token_count": getattr(
                            last_chunk.usage_metadata, "candidates_token_count", 0
                        ),
                        "cached_content_token_count": getattr(
                            last_chunk.usage_metadata, "cached_content_token_count", 0
                        ),
                        "thoughts_token_count": getattr(
                            last_chunk.usage_metadata, "thoughts_token_count", 0
                        ),
                        "total_token_count": getattr(
                            last_chunk.usage_metadata, "total_token_count", 0
                        ),
                    }
                # Add usage_metadata to both run.extra AND outputs
                usage_metadata = _create_usage_metadata(usage_dict)

        except Exception as e:
            logger.debug(f"Error extracting metadata from last chunk: {e}")

    # Return minimal structure with usage_metadata in outputs
    return {
        "content": full_text,
        "usage_metadata": usage_metadata,
    }


def _get_wrapper(
    original_generate: Callable,
    name: str,
    prepopulated_invocation_params: dict,
    tracing_extra: Optional[TracingExtra] = None,
    is_streaming: bool = False,
) -> Callable:
    """Create a wrapper for Gemini's `generate_content` methods."""
    textra = tracing_extra or {}

    @functools.wraps(original_generate)
    def generate(*args, **kwargs):
        # Handle config object before tracing setup
        _convert_config_for_tracing(kwargs)

        decorator = run_helpers.traceable(
            name=name,
            run_type="llm",
            reduce_fn=_reduce_generate_content_chunks if is_streaming else None,
            process_inputs=_process_gemini_inputs,
            process_outputs=(
                _process_generate_content_response if not is_streaming else None
            ),
            _invocation_params_fn=functools.partial(
                _infer_invocation_params, prepopulated_invocation_params
            ),
            **textra,
        )

        return decorator(original_generate)(*args, **kwargs)

    @functools.wraps(original_generate)
    async def agenerate(*args, **kwargs):
        # Handle config object before tracing setup
        _convert_config_for_tracing(kwargs)

        decorator = run_helpers.traceable(
            name=name,
            run_type="llm",
            reduce_fn=_reduce_generate_content_chunks if is_streaming else None,
            process_inputs=_process_gemini_inputs,
            process_outputs=(
                _process_generate_content_response if not is_streaming else None
            ),
            _invocation_params_fn=functools.partial(
                _infer_invocation_params, prepopulated_invocation_params
            ),
            **textra,
        )

        return await decorator(original_generate)(*args, **kwargs)

    return agenerate if run_helpers.is_async(original_generate) else generate


class TracingExtra(TypedDict, total=False):
    metadata: Optional[Mapping[str, Any]]
    tags: Optional[list[str]]
    client: Optional[ls_client.Client]


@warn_beta
def wrap_gemini(
    client: C,
    *,
    tracing_extra: Optional[TracingExtra] = None,
    chat_name: str = "ChatGoogleGenerativeAI",
) -> C:
    """Patch the Google Gen AI client to make it traceable.

    !!! warning

        **BETA**: This wrapper is in beta.

    Supports:
        - `generate_content` and `generate_content_stream` methods
        - Sync and async clients
        - Streaming and non-streaming responses
        - Tool/function calling with proper UI rendering
        - Multimodal inputs (text + images)
        - Image generation with `inline_data` support
        - Token usage tracking including reasoning tokens

    Args:
        client: The Google Gen AI client to patch.
        tracing_extra: Extra tracing information.
        chat_name: The run name for the chat endpoint.

    Returns:
        The patched client.

    Example:
        ```python
        from google import genai
        from google.genai import types
        from langsmith import wrappers

        # Use Google Gen AI client same as you normally would.
        client = wrappers.wrap_gemini(genai.Client(api_key="your-api-key"))

        # Basic text generation:
        response = client.models.generate_content(
            model="gemini-2.5-flash",
            contents="Why is the sky blue?",
        )
        print(response.text)

        # Streaming:
        for chunk in client.models.generate_content_stream(
            model="gemini-2.5-flash",
            contents="Tell me a story",
        ):
            print(chunk.text, end="")

        # Tool/Function calling:
        schedule_meeting_function = {
            "name": "schedule_meeting",
            "description": "Schedules a meeting with specified attendees.",
            "parameters": {
                "type": "object",
                "properties": {
                    "attendees": {"type": "array", "items": {"type": "string"}},
                    "date": {"type": "string"},
                    "time": {"type": "string"},
                    "topic": {"type": "string"},
                },
                "required": ["attendees", "date", "time", "topic"],
            },
        }

        tools = types.Tool(function_declarations=[schedule_meeting_function])
        config = types.GenerateContentConfig(tools=[tools])

        response = client.models.generate_content(
            model="gemini-2.5-flash",
            contents="Schedule a meeting with Bob and Alice tomorrow at 2 PM.",
            config=config,
        )

        # Image generation:
        response = client.models.generate_content(
            model="gemini-2.5-flash-image",
            contents=["Create a picture of a futuristic city"],
        )

        # Save generated image
        from io import BytesIO
        from PIL import Image

        for part in response.candidates[0].content.parts:
            if part.inline_data is not None:
                image = Image.open(BytesIO(part.inline_data.data))
                image.save("generated_image.png")
        ```

    !!! version-added "Added in `langsmith` 0.4.33"

        Initial beta release of Google Gemini wrapper.

    """
    tracing_extra = tracing_extra or {}

    # Extract ls_invocation_params from metadata
    metadata = dict(tracing_extra.get("metadata") or {})
    prepopulated_invocation_params = metadata.pop("ls_invocation_params", {})

    # Create new tracing_extra without ls_invocation_params in metadata
    tracing_extra_rest: TracingExtra = {  # type: ignore[assignment]
        k: v for k, v in tracing_extra.items() if k != "metadata"
    }
    if metadata:
        tracing_extra_rest["metadata"] = metadata  # type: ignore[typeddict-item]

    # Check if already wrapped to prevent double-wrapping
    if (
        hasattr(client, "models")
        and hasattr(client.models, "generate_content")
        and hasattr(client.models.generate_content, "__wrapped__")
    ):
        raise ValueError(
            "This Google Gen AI client has already been wrapped. "
            "Wrapping a client multiple times is not supported."
        )

    # Wrap synchronous methods
    if hasattr(client, "models") and hasattr(client.models, "generate_content"):
        client.models.generate_content = _get_wrapper(  # type: ignore[method-assign]
            client.models.generate_content,
            chat_name,
            prepopulated_invocation_params,
            tracing_extra=tracing_extra_rest,
            is_streaming=False,
        )

    if hasattr(client, "models") and hasattr(client.models, "generate_content_stream"):
        client.models.generate_content_stream = _get_wrapper(  # type: ignore[method-assign]
            client.models.generate_content_stream,
            chat_name,
            prepopulated_invocation_params,
            tracing_extra=tracing_extra_rest,
            is_streaming=True,
        )

    # Wrap async methods (aio namespace)
    if (
        hasattr(client, "aio")
        and hasattr(client.aio, "models")
        and hasattr(client.aio.models, "generate_content")
    ):
        client.aio.models.generate_content = _get_wrapper(  # type: ignore[method-assign]
            client.aio.models.generate_content,
            chat_name,
            prepopulated_invocation_params,
            tracing_extra=tracing_extra_rest,
            is_streaming=False,
        )

    if (
        hasattr(client, "aio")
        and hasattr(client.aio, "models")
        and hasattr(client.aio.models, "generate_content_stream")
    ):
        client.aio.models.generate_content_stream = _get_wrapper(  # type: ignore[method-assign]
            client.aio.models.generate_content_stream,
            chat_name,
            prepopulated_invocation_params,
            tracing_extra=tracing_extra_rest,
            is_streaming=True,
        )

    return client
