from __future__ import annotations

from typing import Any, Optional, Dict, List
from uuid import UUID


def ZappyBeeCallbackHandler() -> Any:
    """Create a LangChain callback handler that traces all chain/LLM/tool runs.

    Usage:
        from zappybee.wrappers.langchain import ZappyBeeCallbackHandler
        from langchain.chains import LLMChain

        handler = ZappyBeeCallbackHandler()
        chain.invoke({"input": "hello"}, config={"callbacks": [handler]})
    """
    try:
        from langchain_core.callbacks import AsyncCallbackHandler
    except ImportError:
        print("[ZappyBee] langchain-core not installed. Install with: pip install langchain-core")
        return None

    class _Handler(AsyncCallbackHandler):
        def __init__(self) -> None:
            self._traces: Dict[UUID, Any] = {}
            self._steps: Dict[UUID, Any] = {}
            # Keep the model used for LLM steps so we can compute cost on step.end().
            self._llm_models: Dict[UUID, str] = {}

        @property
        def always_verbose(self) -> bool:
            return True

        async def on_chain_start(
            self,
            serialized: Dict[str, Any],
            inputs: Dict[str, Any],
            *,
            run_id: UUID,
            parent_run_id: Optional[UUID] = None,
            **kwargs: Any,
        ) -> None:
            from ..client import TokenCat

            name = serialized.get("name", serialized.get("id", ["unknown"])[-1])

            if parent_run_id is None:
                trace = await TokenCat.start_trace(
                    f"langchain/{name}",
                    input=_safe_serialize(inputs),
                )
                self._traces[run_id] = trace
            else:
                parent_trace = self._find_trace(parent_run_id)
                if parent_trace:
                    step = await parent_trace.start_step(name, "tool_call", input=_safe_serialize(inputs))
                    self._steps[run_id] = step

        async def on_chain_end(
            self,
            outputs: Dict[str, Any],
            *,
            run_id: UUID,
            **kwargs: Any,
        ) -> None:
            if run_id in self._steps:
                await self._steps.pop(run_id).end(output=_safe_serialize(outputs), status="success")
            elif run_id in self._traces:
                await self._traces.pop(run_id).end(output=_safe_serialize(outputs), status="success")

        async def on_chain_error(
            self,
            error: BaseException,
            *,
            run_id: UUID,
            **kwargs: Any,
        ) -> None:
            if run_id in self._steps:
                await self._steps.pop(run_id).end(status="error", error_message=str(error))
            elif run_id in self._traces:
                await self._traces.pop(run_id).end(status="error", error_message=str(error))

        async def on_llm_start(
            self,
            serialized: Dict[str, Any],
            prompts: List[str],
            *,
            run_id: UUID,
            parent_run_id: Optional[UUID] = None,
            **kwargs: Any,
        ) -> None:
            model = kwargs.get("invocation_params", {}).get("model_name", "unknown")
            parent_trace = self._find_trace(parent_run_id)
            if parent_trace:
                step = await parent_trace.start_step(
                    "llm_call",
                    "llm_call",
                    input={"prompts": prompts[:2]},
                    model=model,
                )
                self._steps[run_id] = step
                self._llm_models[run_id] = model

        async def on_llm_end(
            self,
            response: Any,
            *,
            run_id: UUID,
            **kwargs: Any,
        ) -> None:
            if run_id not in self._steps:
                return
            step = self._steps.pop(run_id)
            model = self._llm_models.pop(run_id, "unknown")

            output = None
            prompt_tokens = None
            completion_tokens = None

            if hasattr(response, "llm_output") and response.llm_output:
                usage = response.llm_output.get("token_usage", {})
                prompt_tokens = usage.get("prompt_tokens")
                completion_tokens = usage.get("completion_tokens")

            if hasattr(response, "generations") and response.generations:
                first = response.generations[0]
                if first:
                    output = getattr(first[0], "text", None)

            await step.end(
                output=output,
                status="success",
                model=model,
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
            )

        async def on_llm_error(
            self,
            error: BaseException,
            *,
            run_id: UUID,
            **kwargs: Any,
        ) -> None:
            if run_id in self._steps:
                await self._steps.pop(run_id).end(status="error", error_message=str(error))
            self._llm_models.pop(run_id, None)

        async def on_tool_start(
            self,
            serialized: Dict[str, Any],
            input_str: str,
            *,
            run_id: UUID,
            parent_run_id: Optional[UUID] = None,
            **kwargs: Any,
        ) -> None:
            name = serialized.get("name", "tool")
            parent_trace = self._find_trace(parent_run_id)
            if parent_trace:
                step = await parent_trace.start_step(name, "tool_call", input=input_str)
                self._steps[run_id] = step

        async def on_tool_end(
            self,
            output: str,
            *,
            run_id: UUID,
            **kwargs: Any,
        ) -> None:
            if run_id in self._steps:
                await self._steps.pop(run_id).end(output=output, status="success")

        async def on_tool_error(
            self,
            error: BaseException,
            *,
            run_id: UUID,
            **kwargs: Any,
        ) -> None:
            if run_id in self._steps:
                await self._steps.pop(run_id).end(status="error", error_message=str(error))

        def _find_trace(self, run_id: Optional[UUID]) -> Any:
            if run_id is None:
                return None
            if run_id in self._traces:
                return self._traces[run_id]
            # Walk up parent chain — for deeply nested runs, use the first available trace
            for trace in self._traces.values():
                return trace
            return None

    return _Handler()


# Backwards-compatible alias
TokenCatCallbackHandler = ZappyBeeCallbackHandler


def _safe_serialize(obj: Any) -> Any:
    """Convert to JSON-safe format, truncating large values."""
    if isinstance(obj, dict):
        return {k: _safe_serialize(v) for k, v in list(obj.items())[:20]}
    if isinstance(obj, (list, tuple)):
        return [_safe_serialize(v) for v in obj[:10]]
    if isinstance(obj, str) and len(obj) > 2000:
        return obj[:2000] + "...[truncated]"
    if isinstance(obj, (int, float, bool, type(None))):
        return obj
    return str(obj)[:500]
