"""
Vertex AI / Google Gen AI client wrapper for Raindrop observability.

Wraps a Google Gen AI client to capture models.generate_content() calls.

Usage:
    from raindrop_vertex_ai import create_raindrop_vertex_ai
    from google import genai

    raindrop = create_raindrop_vertex_ai(api_key="rk_...")
    client = genai.Client(api_key="...")
    wrapped = raindrop["wrap"](client)

    response = wrapped.models.generate_content(model="gemini-2.0-flash", contents="Hello!")
    raindrop["flush"]()
"""

from __future__ import annotations

import uuid
from typing import Any, Dict, Optional

import raindrop.analytics as raindrop


def create_raindrop_vertex_ai(
    api_key: str,
    user_id: Optional[str] = None,
    convo_id: Optional[str] = None,
) -> dict:
    """Create a Raindrop-instrumented Vertex AI wrapper."""
    raindrop.init(api_key=api_key)

    def wrap(client: Any) -> Any:
        return wrap_genai_client(client, user_id=user_id, convo_id=convo_id)

    return {
        "wrap": wrap,
        "flush": raindrop.flush,
        "shutdown": raindrop.shutdown,
    }


def wrap_genai_client(
    client: Any,
    user_id: Optional[str] = None,
    convo_id: Optional[str] = None,
) -> Any:
    """Wrap a Google Gen AI client to capture generate_content calls."""
    try:
        models = client.models
        original_generate = models.generate_content

        def wrapped_generate_content(**kwargs: Any) -> Any:
            return _instrument_generate(original_generate, kwargs, user_id, convo_id)

        models.generate_content = wrapped_generate_content
    except Exception:
        pass

    return client


def _instrument_generate(
    original_fn: Any,
    kwargs: Dict[str, Any],
    user_id: Optional[str],
    convo_id: Optional[str],
) -> Any:
    try:
        model_id = kwargs.get("model")
        input_text = _extract_input(kwargs.get("contents"))
        event_id = str(uuid.uuid4())
    except Exception:
        return original_fn(**kwargs)

    try:
        response = original_fn(**kwargs)
    except Exception:
        try:
            raindrop.track_ai(
                user_id=user_id or "unknown",
                event="ai_generation",
                event_id=event_id,
                input=input_text,
                model=model_id,
                convo_id=convo_id,
            )
        except Exception:
            pass
        raise

    try:
        output_text = _extract_output(response)
        usage = getattr(response, "usage_metadata", None)
        prompt_tokens = getattr(usage, "prompt_token_count", None) if usage else None
        completion_tokens = getattr(usage, "candidates_token_count", None) if usage else None
        response_model = getattr(response, "model_version", model_id)

        properties: Dict[str, Any] = {}
        if prompt_tokens is not None:
            properties["ai.usage.prompt_tokens"] = prompt_tokens
        if completion_tokens is not None:
            properties["ai.usage.completion_tokens"] = completion_tokens

        raindrop.track_ai(
            user_id=user_id or "unknown",
            event="ai_generation",
            event_id=event_id,
            input=input_text,
            output=output_text,
            model=response_model,
            properties=properties or None,
            convo_id=convo_id,
        )
    except Exception:
        pass

    return response


def _extract_input(contents: Any) -> Optional[str]:
    if isinstance(contents, str):
        return contents
    if not isinstance(contents, list):
        return None
    try:
        user_contents = []
        all_contents = []
        for c in contents:
            role = c.get("role") if isinstance(c, dict) else getattr(c, "role", None)
            raw_parts = c.get("parts", []) if isinstance(c, dict) else getattr(c, "parts", [])
            texts = []
            for p in (raw_parts if isinstance(raw_parts, list) else []):
                t = p.get("text") if isinstance(p, dict) else getattr(p, "text", None)
                if t:
                    texts.append(t)
            if texts:
                all_contents.append(" ".join(texts))
                if role == "user":
                    user_contents.append(" ".join(texts))
        last = user_contents[-1] if user_contents else (all_contents[-1] if all_contents else None)
        return last
    except Exception:
        return None


def _extract_output(response: Any) -> Optional[str]:
    try:
        text = getattr(response, "text", None)
        if isinstance(text, str):
            return text
        candidates = getattr(response, "candidates", None)
        if candidates and len(candidates) > 0:
            content = getattr(candidates[0], "content", None)
            if content:
                parts = getattr(content, "parts", [])
                texts = [getattr(p, "text", "") for p in parts if getattr(p, "text", None)]
                if texts:
                    return "\n".join(texts)
    except Exception:
        pass
    return None
