"""
Google Generative AI integration — wraps GenerativeModel to auto-trace.

Supports ``generate_content()``, ``generate_content_async()``, and
streaming via ``stream=True``.

Usage:
    import google.generativeai as genai
    from syntropy import Syntropy
    from syntropy.integrations.google import wrap_google_ai

    syn = Syntropy(api_key="syn_...", agent_id="my-agent")
    model = genai.GenerativeModel("gemini-2.0-flash")
    model = wrap_google_ai(model, syntropy=syn)

    # Non-streaming (auto-traced):
    response = model.generate_content("Explain quantum computing")

    # Streaming (auto-traced on completion):
    for chunk in model.generate_content("Tell me a story", stream=True):
        print(chunk.text, end="")
"""

from __future__ import annotations

import time
from typing import Any, Optional

from syntropy.client import Syntropy
from syntropy.trace_context import is_inside_trace


# ----------------------------------------------------------------
# Helpers
# ----------------------------------------------------------------

def _model_name_of(model: Any) -> str:
    return (
        getattr(model, "model_name", None)
        or getattr(model, "_model_name", None)
        or "gemini"
    )


def _extract_text(response: Any) -> str:
    try:
        return response.text
    except (AttributeError, ValueError):
        if hasattr(response, "candidates") and response.candidates:
            parts = response.candidates[0].content.parts
            return " ".join(p.text for p in parts if hasattr(p, "text"))
    return ""


def _extract_usage(response: Any) -> tuple[int, int]:
    usage = getattr(response, "usage_metadata", None)
    tokens_in = getattr(usage, "prompt_token_count", 0) if usage else 0
    tokens_out = getattr(usage, "candidates_token_count", 0) if usage else 0
    return tokens_in, tokens_out


def _wrap_stream(
    stream: Any,
    prompt: str,
    start: float,
    model_name: str,
    syntropy: Syntropy,
    agent_id: str,
) -> Any:
    """Yield chunks from a Gemini streaming response and trace on end."""
    text_parts: list[str] = []
    tokens_in = 0
    tokens_out = 0
    stream_error: BaseException | None = None

    try:
        for chunk in stream:
            try:
                t = chunk.text
                if t:
                    text_parts.append(t)
            except (AttributeError, ValueError):
                pass
            # Last chunk carries usage
            usage = getattr(chunk, "usage_metadata", None)
            if usage:
                tokens_in = getattr(usage, "prompt_token_count", tokens_in)
                tokens_out = getattr(usage, "candidates_token_count", tokens_out)
            yield chunk
    except BaseException as exc:
        stream_error = exc
        raise
    finally:
        elapsed_ms = int((time.perf_counter() - start) * 1000)
        syntropy.trace(
            prompt=prompt,
            response="".join(text_parts),
            model=model_name,
            agent_id=agent_id,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            latency_ms=elapsed_ms,
            status="error" if stream_error else "success",
            metadata={
                "provider": "google",
                "stream": True,
                **({
                    "error": str(stream_error),
                    "partialResponse": True,
                } if stream_error else {}),
            },
        )


# ----------------------------------------------------------------
# Traced class
# ----------------------------------------------------------------


class TracedGenerativeModel:
    """Proxy around google.generativeai.GenerativeModel."""

    def __init__(self, model: Any, syntropy: Syntropy, agent_id: str):
        self._model = model
        self._syntropy = syntropy
        self._agent_id = agent_id

    def generate_content(self, contents: Any, **kwargs: Any) -> Any:
        # Skip if already inside an orchestration-level trace (e.g. LangChain)
        if is_inside_trace():
            return self._model.generate_content(contents, **kwargs)

        start = time.perf_counter()
        prompt = contents if isinstance(contents, str) else str(contents)
        model_name = _model_name_of(self._model)
        is_stream = kwargs.get("stream", False)

        response = self._model.generate_content(contents, **kwargs)

        # ---- Streaming path ----
        if is_stream:
            return _wrap_stream(
                response, prompt, start, model_name,
                self._syntropy, self._agent_id,
            )

        # ---- Non-streaming path ----
        elapsed_ms = int((time.perf_counter() - start) * 1000)

        output = _extract_text(response)
        tokens_in, tokens_out = _extract_usage(response)

        self._syntropy.trace(
            prompt=prompt,
            response=output,
            model=model_name,
            agent_id=self._agent_id,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            latency_ms=elapsed_ms,
            metadata={
                "provider": "google",
                "stream": False,
                "safety_ratings": str(getattr(response, "prompt_feedback", None)),
            },
        )

        return response

    async def generate_content_async(self, contents: Any, **kwargs: Any) -> Any:
        start = time.perf_counter()
        response = await self._model.generate_content_async(contents, **kwargs)
        elapsed_ms = int((time.perf_counter() - start) * 1000)

        prompt = contents if isinstance(contents, str) else str(contents)
        output = _extract_text(response)
        tokens_in, tokens_out = _extract_usage(response)
        model_name = _model_name_of(self._model)

        self._syntropy.trace(
            prompt=prompt,
            response=output,
            model=model_name,
            agent_id=self._agent_id,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            latency_ms=elapsed_ms,
            metadata={"provider": "google", "stream": False},
        )

        return response

    def __getattr__(self, name: str) -> Any:
        return getattr(self._model, name)


def wrap_google_ai(
    model: Any,
    *,
    syntropy: Syntropy,
    agent_id: str = "",
) -> TracedGenerativeModel:
    """
    Wrap a Google GenerativeModel to auto-trace generate_content() calls.
    Supports sync, async, and streaming.

    Args:
        model: A google.generativeai.GenerativeModel instance.
        syntropy: A Syntropy client instance.
        agent_id: Override agent_id.

    Returns:
        A drop-in replacement with tracing enabled.
    """
    return TracedGenerativeModel(
        model, syntropy, agent_id or syntropy._config.agent_id
    )
