"""
DSPy module wrapper for Raindrop observability.

Wraps a DSPy module (Predict, ChainOfThought, etc.) to capture
input, output, model name, and token usage.

Usage:
    from raindrop_dspy import create_raindrop_dspy
    import dspy

    raindrop = create_raindrop_dspy(api_key="rk_...")
    lm = dspy.LM("openai/gpt-4o-mini")
    dspy.configure(lm=lm)

    predict = dspy.Predict("question -> answer")
    wrapped = raindrop["wrap"](predict)

    result = wrapped(question="What is the capital of France?")
    raindrop["flush"]()
"""

from __future__ import annotations

import json
import uuid
from typing import Any, Dict, Optional

import raindrop.analytics as raindrop


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

    def wrap(module: Any) -> Any:
        return wrap_module(module, user_id=user_id, convo_id=convo_id)

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


def wrap_module(
    module: Any,
    user_id: Optional[str] = None,
    convo_id: Optional[str] = None,
) -> Any:
    """Wrap a DSPy module to capture forward/call invocations."""
    try:
        original_forward = getattr(module, "forward", None)
        if original_forward is None or not callable(original_forward):
            return module

        def wrapped_forward(*args: Any, **kwargs: Any) -> Any:
            return _instrument_forward(original_forward, args, kwargs, user_id, convo_id)

        module.forward = wrapped_forward
    except Exception:
        pass

    return module


def _instrument_forward(
    original_fn: Any,
    args: tuple,
    kwargs: Dict[str, Any],
    user_id: Optional[str],
    convo_id: Optional[str],
) -> Any:
    try:
        input_text = _extract_input(kwargs)
        model_name = _get_model_name()
        event_id = str(uuid.uuid4())
    except Exception:
        return original_fn(*args, **kwargs)

    try:
        result = original_fn(*args, **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_name,
                convo_id=convo_id,
            )
        except Exception:
            pass
        raise

    try:
        output_text = _extract_output(result)
        properties = _extract_usage(result)

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

    return result


def _extract_input(kwargs: Dict[str, Any]) -> Optional[str]:
    """Extract a readable input from DSPy module kwargs."""
    if not kwargs:
        return None
    try:
        parts = []
        for key, value in kwargs.items():
            if isinstance(value, str):
                parts.append(value)
            else:
                try:
                    parts.append(json.dumps(value, default=str))
                except Exception:
                    parts.append(str(value))
        return "\n".join(parts) if parts else None
    except Exception:
        return None


def _extract_output(result: Any) -> Optional[str]:
    """Extract output text from a DSPy Prediction result."""
    if result is None:
        return None
    try:
        # DSPy Prediction objects have completions or direct attribute access
        if hasattr(result, "toDict"):
            d = result.toDict()
            values = [str(v) for v in d.values() if v is not None]
            return "\n".join(values) if values else None

        # Fallback: try common output attribute names
        for attr in ("answer", "output", "response", "text", "completion"):
            val = getattr(result, attr, None)
            if val is not None:
                return str(val)

        return str(result)
    except Exception:
        return None


def _get_model_name() -> Optional[str]:
    """Try to get the configured DSPy LM model name."""
    try:
        import dspy
        settings = getattr(dspy, "settings", None)
        if settings:
            lm = getattr(settings, "lm", None)
            if lm:
                model = getattr(lm, "model", None)
                if model:
                    return str(model)
    except Exception:
        pass
    return None


def _extract_usage(result: Any) -> Optional[Dict[str, Any]]:
    """Extract token usage from DSPy Prediction result if available."""
    try:
        # DSPy Prediction uses get_lm_usage() which returns {lm_name: {prompt_tokens, completion_tokens}}
        get_usage = getattr(result, "get_lm_usage", None)
        if callable(get_usage):
            usage_dict = get_usage()
            if isinstance(usage_dict, dict):
                props: Dict[str, Any] = {}
                total_prompt = 0
                total_completion = 0
                for lm_usage in usage_dict.values():
                    if isinstance(lm_usage, dict):
                        total_prompt += lm_usage.get("prompt_tokens", 0) or 0
                        total_completion += lm_usage.get("completion_tokens", 0) or 0
                if total_prompt > 0:
                    props["ai.usage.prompt_tokens"] = total_prompt
                if total_completion > 0:
                    props["ai.usage.completion_tokens"] = total_completion
                return props if props else None

        # Fallback: check for direct usage attribute
        usage = getattr(result, "usage", None)
        if usage is None:
            return None
        props = {}
        pt = getattr(usage, "prompt_tokens", None)
        ct = getattr(usage, "completion_tokens", None)
        if isinstance(usage, dict):
            pt = pt if pt is not None else usage.get("prompt_tokens")
            ct = ct if ct is not None else usage.get("completion_tokens")
        if pt is not None:
            props["ai.usage.prompt_tokens"] = pt
        if ct is not None:
            props["ai.usage.completion_tokens"] = ct
        return props if props else None
    except Exception:
        return None
