"""
AgentBase Auto-Instrumentation
Monkey-patches popular LLM libraries to automatically trace all calls.
"""

import functools
import logging
import json

from agentbase_sdk.tracer import Span

logger = logging.getLogger("agentbase")


def _safe_str(obj, max_len=5000):
    try:
        if isinstance(obj, str):
            return obj[:max_len]
        if isinstance(obj, (dict, list)):
            return json.dumps(obj, default=str)[:max_len]
        return str(obj)[:max_len]
    except Exception:
        return "<unserializable>"


def instrument_openai():
    """Auto-instrument OpenAI SDK."""
    try:
        import openai
    except ImportError:
        logger.warning("AgentBase: openai not installed, skipping")
        return

    original_create = openai.resources.chat.completions.Completions.create

    @functools.wraps(original_create)
    def wrapped(self, *args, **kwargs):
        model = kwargs.get("model", "unknown")
        messages = kwargs.get("messages", [])
        with Span("openai.chat.completions", span_type="llm") as s:
            s.set_model(model)
            s.set_input(_safe_str(messages))
            try:
                response = original_create(self, *args, **kwargs)
                if hasattr(response, "usage") and response.usage:
                    s.set_tokens(
                        prompt_tokens=response.usage.prompt_tokens or 0,
                        completion_tokens=response.usage.completion_tokens or 0,
                    )
                if hasattr(response, "choices") and response.choices:
                    s.set_output(response.choices[0].message.content or "")
                return response
            except Exception as e:
                s.set_error(type(e).__name__, str(e))
                raise

    openai.resources.chat.completions.Completions.create = wrapped
    logger.debug("AgentBase: OpenAI instrumented")


def instrument_anthropic():
    """Auto-instrument Anthropic SDK."""
    try:
        import anthropic
    except ImportError:
        logger.warning("AgentBase: anthropic not installed, skipping")
        return

    original_create = anthropic.resources.messages.Messages.create

    @functools.wraps(original_create)
    def wrapped(self, *args, **kwargs):
        model = kwargs.get("model", "unknown")
        messages = kwargs.get("messages", [])
        with Span("anthropic.messages.create", span_type="llm") as s:
            s.set_model(model)
            s.set_input(_safe_str(messages))
            try:
                response = original_create(self, *args, **kwargs)
                if hasattr(response, "usage"):
                    s.set_tokens(
                        prompt_tokens=getattr(response.usage, "input_tokens", 0),
                        completion_tokens=getattr(response.usage, "output_tokens", 0),
                    )
                if hasattr(response, "content") and response.content:
                    text = "".join(b.text for b in response.content if hasattr(b, "text"))
                    s.set_output(text)
                return response
            except Exception as e:
                s.set_error(type(e).__name__, str(e))
                raise

    anthropic.resources.messages.Messages.create = wrapped
    logger.debug("AgentBase: Anthropic instrumented")


def instrument_langchain():
    """Auto-instrument LangChain using callbacks."""
    try:
        from langchain_core.callbacks import BaseCallbackHandler
        from langchain_core.callbacks.manager import CallbackManager
    except ImportError:
        logger.warning("AgentBase: langchain not installed, skipping")
        return

    class ABHandler(BaseCallbackHandler):
        def __init__(self):
            self._spans = {}

        def on_llm_start(self, serialized, prompts, *, run_id, **kwargs):
            s = Span("langchain.llm", span_type="llm")
            s.__enter__()
            s.set_input(_safe_str(prompts))
            self._spans[str(run_id)] = s

        def on_llm_end(self, response, *, run_id, **kwargs):
            s = self._spans.pop(str(run_id), None)
            if s:
                s.set_output(_safe_str(response.generations) if response.generations else "")
                s.__exit__(None, None, None)

        def on_llm_error(self, error, *, run_id, **kwargs):
            s = self._spans.pop(str(run_id), None)
            if s:
                s.set_error(type(error).__name__, str(error))
                s.__exit__(type(error), error, None)

        def on_tool_start(self, serialized, input_str, *, run_id, **kwargs):
            name = serialized.get("name", "unknown_tool")
            s = Span(f"langchain.tool.{name}", span_type="tool")
            s.__enter__()
            s.set_tool(name)
            s.set_input(_safe_str(input_str))
            self._spans[str(run_id)] = s

        def on_tool_end(self, output, *, run_id, **kwargs):
            s = self._spans.pop(str(run_id), None)
            if s:
                s.set_output(_safe_str(output))
                s.__exit__(None, None, None)

        def on_tool_error(self, error, *, run_id, **kwargs):
            s = self._spans.pop(str(run_id), None)
            if s:
                s.set_error(type(error).__name__, str(error))
                s.__exit__(type(error), error, None)

    handler = ABHandler()
    _orig = CallbackManager.configure

    @classmethod
    def _patched(cls, *args, **kwargs):
        mgr = _orig.__func__(cls, *args, **kwargs)
        if not any(isinstance(h, ABHandler) for h in mgr.handlers):
            mgr.add_handler(handler)
        return mgr

    CallbackManager.configure = _patched
    logger.debug("AgentBase: LangChain instrumented")


def instrument_crewai():
    """Auto-instrument CrewAI."""
    try:
        from crewai import Crew
    except ImportError:
        logger.warning("AgentBase: crewai not installed, skipping")
        return

    original_kickoff = Crew.kickoff

    @functools.wraps(original_kickoff)
    def wrapped(self, *args, **kwargs):
        from agentbase_sdk.tracer import Tracer
        crew_name = getattr(self, "name", "unnamed-crew")
        with Tracer(f"crewai.{crew_name}") as t:
            t.set_framework("crewai")
            try:
                result = original_kickoff(self, *args, **kwargs)
                t.set_output(_safe_str(result))
                return result
            except Exception as e:
                t.set_error(type(e).__name__, str(e))
                raise

    Crew.kickoff = wrapped
    logger.debug("AgentBase: CrewAI instrumented")
