"""
Google Gemini provider integration for Amplitude AI SDK.

Wraps the Google Gen AI SDK (``google-genai``) to automatically track LLM usage
in Amplitude.  Targets the **new** ``google-genai`` package — not the deprecated
``google-generativeai``.

Usage::

    from google import genai
    from amplitude_ai.providers.gemini import Gemini

    client = Gemini(amplitude=amp, api_key="...", model_name="gemini-2.5-flash")
    response = client.generate_content("Why is the sky blue?", amplitude_user_id="u1")
"""
from __future__ import annotations

import time
import uuid
from typing import Any, Dict, Iterator, List, Optional, Union

try:
    from google import genai
    from google.genai import types as genai_types

    GEMINI_AVAILABLE = True
except ImportError:
    GEMINI_AVAILABLE = False
    genai = None  # type: ignore[assignment]
    genai_types = None  # type: ignore[assignment]

from amplitude import Amplitude

from ..core.privacy import PrivacyConfig
from ..exceptions import ProviderError
from ..utils.costs import calculate_cost
from ..utils.logger import get_logger
from .base import BaseAIProvider, AmplitudeOrAI


class Gemini(BaseAIProvider):
    """
    Google Gemini client wrapper that automatically tracks usage in Amplitude.

    Uses the ``google-genai`` SDK (``from google import genai``).
    """

    def __init__(
        self,
        amplitude: AmplitudeOrAI,
        api_key: Optional[str] = None,
        model_name: str = "gemini-2.5-flash",
        privacy_mode: bool = False,
        content_mode: Optional[str] = None,
        **client_kwargs,
    ):
        if not GEMINI_AVAILABLE:
            raise ProviderError(
                "Google Gen AI SDK not installed. Install with: pip install google-genai"
            )

        privacy_config = PrivacyConfig(privacy_mode=privacy_mode, content_mode=content_mode)
        BaseAIProvider.__init__(self, amplitude=amplitude, privacy_config=privacy_config)

        self.model_name = model_name
        self.client = genai.Client(api_key=api_key, **client_kwargs)

    # ------------------------------------------------------------------
    # BaseAIProvider hooks
    # ------------------------------------------------------------------

    def _get_provider_name(self) -> str:
        return "gemini"

    def _extract_model_name(self, **kwargs) -> str:
        return kwargs.get("model", self.model_name)

    def _extract_response_model(self, response: Any) -> Optional[str]:
        return getattr(response, "model_version", None) or None

    def _extract_messages(self, **kwargs) -> List[Dict[str, Any]]:
        return kwargs.get("contents", [])

    def _extract_system_prompt(self, **kwargs) -> Optional[str]:
        config = kwargs.get("config")
        if config is None:
            return None
        si = getattr(config, "system_instruction", None) if not isinstance(config, dict) else config.get("system_instruction")
        if si is None:
            return None
        if isinstance(si, str):
            return si
        if hasattr(si, "parts"):
            parts = si.parts if hasattr(si.parts, "__iter__") else [si.parts]
            return "".join(getattr(p, "text", str(p)) for p in parts)
        return str(si)

    def _extract_model_params(self, **kwargs) -> Dict[str, Any]:
        result: Dict[str, Any] = {
            "temperature": None,
            "top_p": None,
            "max_output_tokens": None,
            "is_streaming": kwargs.get("stream"),
        }
        config = kwargs.get("config")
        if config is not None:
            if isinstance(config, dict):
                result["temperature"] = config.get("temperature")
                result["top_p"] = config.get("top_p")
                result["max_output_tokens"] = config.get("max_output_tokens")
            else:
                result["temperature"] = getattr(config, "temperature", None)
                result["top_p"] = getattr(config, "top_p", None)
                result["max_output_tokens"] = getattr(config, "max_output_tokens", None)
        return result

    def _format_messages_for_tracking(self, contents: Any) -> List[Dict[str, Any]]:
        from ..core.privacy import extract_text_from_gemini_contents, get_text_from_llm_message, sanitize_any_content

        messages: list[dict[str, Any]] = []

        if isinstance(contents, str):
            if contents:
                sanitized_result = sanitize_any_content(
                    contents,
                    privacy_mode=self.privacy_config.privacy_mode,
                    redact_pii=self.privacy_config.redact_pii,
                )
                message: dict[str, Any] = {"role": "user"}
                if self.privacy_config.privacy_mode:
                    message["content_hash"] = sanitized_result.get("content_hash")
                elif "$llm_message" in sanitized_result:
                    message["content"] = get_text_from_llm_message(sanitized_result["$llm_message"])
                messages.append(message)

        elif isinstance(contents, list):
            for content in contents:
                if isinstance(content, dict):
                    role = content.get("role", "user")
                    text_content = extract_text_from_gemini_contents([content])
                else:
                    role = "user"
                    text_content = str(content)

                message: dict[str, Any] = {"role": role}
                if text_content:
                    sanitized_result = sanitize_any_content(
                        text_content,
                        privacy_mode=self.privacy_config.privacy_mode,
                        redact_pii=self.privacy_config.redact_pii,
                    )
                    if self.privacy_config.privacy_mode:
                        message["content_hash"] = sanitized_result.get("content_hash")
                    elif "$llm_message" in sanitized_result:
                        message["content"] = get_text_from_llm_message(sanitized_result["$llm_message"])
                messages.append(message)

        return messages

    def _extract_response_content(self, response: Any) -> str:
        if hasattr(response, "text"):
            try:
                return response.text
            except (ValueError, AttributeError):
                pass
        if hasattr(response, "parts") and response.parts:
            return "".join(getattr(p, "text", "") for p in response.parts if hasattr(p, "text"))
        return ""

    def _extract_reasoning_content(self, response: Any) -> Optional[str]:
        if not hasattr(response, "candidates") or not response.candidates:
            return None
        candidate = response.candidates[0]
        content = getattr(candidate, "content", None)
        if content is None or not hasattr(content, "parts"):
            return None
        thinking_parts = [
            part.text for part in content.parts
            if getattr(part, "thought", False) and hasattr(part, "text")
        ]
        return "".join(thinking_parts) if thinking_parts else None

    def _extract_finish_reason(self, response: Any) -> Optional[str]:
        if not hasattr(response, "candidates") or not response.candidates:
            return None
        reason = str(getattr(response.candidates[0], "finish_reason", ""))
        if "STOP" in reason:
            return "stop"
        if "MAX_TOKENS" in reason:
            return "length"
        if "SAFETY" in reason:
            return "content_filter"
        return reason.lower() if reason else None

    def _extract_tool_calls(self, response: Any) -> Optional[List[Dict[str, Any]]]:
        func_calls = getattr(response, "function_calls", None)
        if func_calls:
            return [
                {
                    "type": "function",
                    "id": str(uuid.uuid4()),
                    "function": {
                        "name": getattr(fc, "name", ""),
                        "arguments": str(getattr(fc, "args", {})),
                    },
                }
                for fc in func_calls
            ]
        if not hasattr(response, "candidates") or not response.candidates:
            return None
        candidate = response.candidates[0]
        content = getattr(candidate, "content", None)
        if content is None or not hasattr(content, "parts"):
            return None
        tool_calls = []
        for part in content.parts:
            fc = getattr(part, "function_call", None)
            if fc is not None:
                tool_calls.append({
                    "type": "function",
                    "id": str(uuid.uuid4()),
                    "function": {
                        "name": getattr(fc, "name", ""),
                        "arguments": str(getattr(fc, "args", {})),
                    },
                })
        return tool_calls if tool_calls else None

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def generate_content(
        self,
        contents: Union[str, List[Any]],
        *,
        model: Optional[str] = None,
        config: Any = None,
        amplitude_user_id: Optional[str] = None,
        amplitude_conversation_id: Optional[str] = None,
        amplitude_trace_id: Optional[str] = None,
        amplitude_turn_id: Optional[int] = None,
        amplitude_event_properties: Optional[Dict[str, Any]] = None,
        amplitude_user_properties: Optional[Dict[str, Any]] = None,
        amplitude_groups: Optional[Dict[str, Any]] = None,
        stream: bool = False,
        **kwargs,
    ):
        """
        Generate content with automatic Amplitude tracking.

        Args:
            contents: The input prompt or conversation.
            model: Model name override (defaults to ``self.model_name``).
            config: ``types.GenerateContentConfig`` or dict.
            amplitude_user_id: User ID for tracking.
            amplitude_conversation_id: Conversation ID for linking events.
            amplitude_trace_id: Trace ID for linking related events.
            amplitude_turn_id: Turn number in conversation.
            amplitude_event_properties: Additional event properties.
            amplitude_user_properties: User properties to set.
            amplitude_groups: Group analytics properties.
            stream: Whether to stream the response.
            **kwargs: Additional arguments forwarded to the SDK.
        """
        # Only auto-generate IDs when no active session context can supply them
        if not amplitude_trace_id:
            amplitude_trace_id = str(uuid.uuid4())
        if not amplitude_conversation_id:
            from ..context import get_active_context
            ctx = get_active_context()
            if ctx is None or not ctx.session_id:
                amplitude_conversation_id = str(uuid.uuid4())

        effective_model = model or self.model_name

        if stream:
            return self._generate_streaming(
                contents=contents,
                model=effective_model,
                config=config,
                amplitude_user_id=amplitude_user_id,
                amplitude_conversation_id=amplitude_conversation_id,
                amplitude_trace_id=amplitude_trace_id,
                amplitude_turn_id=amplitude_turn_id,
                amplitude_event_properties=amplitude_event_properties,
                amplitude_user_properties=amplitude_user_properties,
                amplitude_groups=amplitude_groups,
                **kwargs,
            )
        return self._generate_non_streaming(
            contents=contents,
            model=effective_model,
            config=config,
            amplitude_user_id=amplitude_user_id,
            amplitude_conversation_id=amplitude_conversation_id,
            amplitude_trace_id=amplitude_trace_id,
            amplitude_turn_id=amplitude_turn_id,
            amplitude_event_properties=amplitude_event_properties,
            amplitude_user_properties=amplitude_user_properties,
            amplitude_groups=amplitude_groups,
            **kwargs,
        )

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    def _generate_non_streaming(
        self,
        contents,
        model: str,
        config: Any = None,
        amplitude_user_id=None,
        amplitude_conversation_id=None,
        amplitude_trace_id=None,
        amplitude_turn_id=None,
        amplitude_event_properties=None,
        amplitude_user_properties=None,
        amplitude_groups=None,
        **kwargs,
    ):
        start_time = time.time()
        try:
            response = self.client.models.generate_content(
                model=model, contents=contents, config=config, **kwargs,
            )
            latency_ms = (time.time() - start_time) * 1000

            if amplitude_user_id:
                token_usage = self._extract_token_usage(response)
                total_cost_usd = None
                if token_usage.get("input_tokens") is not None and token_usage.get("output_tokens") is not None:
                    total_cost_usd = calculate_cost(
                        model_name=model,
                        input_tokens=token_usage.get("input_tokens") or 0,
                        output_tokens=token_usage.get("output_tokens") or 0,
                        reasoning_tokens=token_usage.get("reasoning_tokens") or 0,
                        cache_read_input_tokens=token_usage.get("cache_read_input_tokens") or 0,
                        cache_creation_input_tokens=token_usage.get("cache_creation_input_tokens") or 0,
                        default_provider='google',
                    )
                self.track_completion(
                    user_id=amplitude_user_id,
                    response=response,
                    latency_ms=latency_ms,
                    conversation_id=amplitude_conversation_id,
                    trace_id=amplitude_trace_id,
                    turn_id=amplitude_turn_id,
                    event_properties=amplitude_event_properties,
                    user_properties=amplitude_user_properties,
                    groups=amplitude_groups,
                    total_cost_usd=total_cost_usd,
                    contents=contents,
                    config=config,
                    **kwargs,
                )
            return response

        except Exception as e:
            latency_ms = (time.time() - start_time) * 1000
            if amplitude_user_id:
                self.track_completion(
                    user_id=amplitude_user_id,
                    response=None,
                    latency_ms=latency_ms,
                    conversation_id=amplitude_conversation_id,
                    trace_id=amplitude_trace_id,
                    turn_id=amplitude_turn_id,
                    is_error=True,
                    error_message=str(e),
                    event_properties=amplitude_event_properties,
                    user_properties=amplitude_user_properties,
                    groups=amplitude_groups,
                    contents=contents,
                    config=config,
                    **kwargs,
                )
            raise

    def _generate_streaming(
        self,
        contents,
        model: str,
        config: Any = None,
        amplitude_user_id=None,
        amplitude_conversation_id=None,
        amplitude_trace_id=None,
        amplitude_turn_id=None,
        amplitude_event_properties=None,
        amplitude_user_properties=None,
        amplitude_groups=None,
        **kwargs,
    ) -> Iterator:
        if not amplitude_user_id:
            return self.client.models.generate_content_stream(
                model=model, contents=contents, config=config, **kwargs,
            )

        streaming_tracker = self.create_streaming_tracker(
            user_id=amplitude_user_id,
            conversation_id=amplitude_conversation_id,
            trace_id=amplitude_trace_id,
            turn_id=amplitude_turn_id,
            event_properties=amplitude_event_properties,
            user_properties=amplitude_user_properties,
            groups=amplitude_groups,
            contents=contents,
            config=config,
            **kwargs,
        )

        try:
            stream = self.client.models.generate_content_stream(
                model=model, contents=contents, config=config, **kwargs,
            )
        except Exception as e:
            streaming_tracker.set_error(str(e))
            streaming_tracker.finalize(model)
            raise

        try:
            for chunk in stream:
                try:
                    if not getattr(streaming_tracker, "_response_model", None):
                        _mv = getattr(chunk, "model_version", None)
                        if isinstance(_mv, str) and _mv:
                            streaming_tracker._response_model = _mv

                    if hasattr(chunk, "candidates") and chunk.candidates:
                        candidate_content = getattr(chunk.candidates[0], "content", None)
                        for part in getattr(candidate_content, "parts", []) if candidate_content else []:
                            if getattr(part, "thought", False) and hasattr(part, "text"):
                                streaming_tracker.add_reasoning_chunk(part.text)
                            elif hasattr(part, "text") and part.text:
                                streaming_tracker.add_content_chunk(part.text)
                    elif hasattr(chunk, "text") and chunk.text:
                        streaming_tracker.add_content_chunk(chunk.text)

                    usage = getattr(chunk, "usage_metadata", None)
                    if usage:
                        token_update = {}
                        if hasattr(usage, "prompt_token_count"):
                            token_update["input_tokens"] = usage.prompt_token_count
                        if hasattr(usage, "candidates_token_count"):
                            token_update["output_tokens"] = usage.candidates_token_count
                        if hasattr(usage, "total_token_count"):
                            token_update["total_tokens"] = usage.total_token_count
                        streaming_tracker.update_token_usage(token_update)

                    if hasattr(chunk, "candidates") and chunk.candidates:
                        finish_reason = str(getattr(chunk.candidates[0], "finish_reason", ""))
                        if "STOP" in finish_reason:
                            streaming_tracker.set_finish_reason("stop")
                        elif "MAX_TOKENS" in finish_reason:
                            streaming_tracker.set_finish_reason("length")
                        elif finish_reason:
                            streaming_tracker.set_finish_reason(finish_reason.lower())

                except Exception as tracking_error:
                    get_logger(self.amplitude).warning(f"Tracking failed for chunk: {tracking_error}")

                yield chunk

            if streaming_tracker.token_usage:
                try:
                    total_cost_usd = calculate_cost(
                        model_name=model,
                        input_tokens=streaming_tracker.token_usage.get("input_tokens", 0),
                        output_tokens=streaming_tracker.token_usage.get("output_tokens", 0),
                        reasoning_tokens=streaming_tracker.token_usage.get("reasoning_tokens", 0),
                        cache_read_input_tokens=streaming_tracker.token_usage.get("cache_read_input_tokens", 0),
                        cache_creation_input_tokens=streaming_tracker.token_usage.get("cache_creation_input_tokens", 0),
                        default_provider='google',
                    )
                    streaming_tracker.token_usage["total_cost_usd"] = total_cost_usd
                except Exception as cost_error:
                    get_logger(self.amplitude).warning(f"Cost calculation failed: {cost_error}")

            try:
                streaming_tracker.finalize(model)
            except Exception as finalize_error:
                get_logger(self.amplitude).warning(f"Streaming tracker finalization failed: {finalize_error}")

        except Exception as e:
            try:
                streaming_tracker.set_error(str(e))
                streaming_tracker.finalize(model)
            except Exception as tracking_error:
                get_logger(self.amplitude).warning(f"Error tracking failed: {tracking_error}")
            raise
        finally:
            try:
                streaming_tracker.finalize_on_abandon(model)
            except Exception:
                pass

    def __getattr__(self, name):
        """Forward other attributes to the underlying client."""
        return getattr(self.client, name)


# TODO: Add async Gemini wrapper with async streaming support.
# When implementing, use the same finalize_on_abandon guard pattern
# from the sync Gemini above (and OpenAI's WrappedAsyncCompletions).
