from __future__ import annotations

import asyncio
import time
from typing import Any, Dict, List, Optional, Union

from ..client import VantageClient

try:
    from langchain_core.callbacks.base import BaseCallbackHandler
    from langchain_core.outputs import LLMResult
    _LANGCHAIN_AVAILABLE = True
except ImportError:
    _LANGCHAIN_AVAILABLE = False
    BaseCallbackHandler = object  # type: ignore[assignment,misc]
    LLMResult = Any  # type: ignore[assignment,misc]


class VantageCallbackHandler(BaseCallbackHandler):
    """LangChain callback handler — logs LLM calls to Vantage automatically."""

    def __init__(self, client: VantageClient) -> None:
        if _LANGCHAIN_AVAILABLE:
            super().__init__()
        self._client = client
        self._llm_start_times: Dict[str, float] = {}
        self._current_workflow_id: Optional[str] = None

    def _fire(self, coro) -> None:
        try:
            loop = asyncio.get_event_loop()
            if loop.is_running():
                loop.create_task(coro)
            else:
                loop.run_until_complete(coro)
        except RuntimeError:
            asyncio.run(coro)

    def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: List[str],
        *,
        run_id: Any = None,
        **kwargs: Any,
    ) -> None:
        run_id_str = str(run_id) if run_id is not None else ""
        self._llm_start_times[run_id_str] = time.monotonic()
        if not self._current_workflow_id:
            self._current_workflow_id = run_id_str

    def on_llm_end(
        self,
        response: Any,
        *,
        run_id: Any = None,
        **kwargs: Any,
    ) -> None:
        run_id_str = str(run_id) if run_id is not None else ""
        start = self._llm_start_times.pop(run_id_str, None)
        duration_ms = int((time.monotonic() - start) * 1000) if start is not None else None

        generations = []
        if hasattr(response, "generations"):
            for gen_list in response.generations:
                generations.extend(gen_list)
        output_text = "".join(
            g.text if hasattr(g, "text") else str(g) for g in generations
        )
        llm_output = getattr(response, "llm_output", {}) or {}

        self._fire(
            self._client.log_event(
                type="llm_call",
                input={"runId": run_id_str},
                output={"text": output_text, "llmOutput": llm_output},
                metadata={"durationMs": duration_ms},
            )
        )

    def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        *,
        run_id: Any = None,
        **kwargs: Any,
    ) -> None:
        run_id_str = str(run_id) if run_id is not None else ""
        tool_name = (serialized.get("id") or ["unknown"])[-1]
        self._fire(
            self._client.log_event(
                type="tool_use",
                input={"tool": tool_name, "input": input_str},
                metadata={"runId": run_id_str},
            )
        )

    def on_tool_end(
        self,
        output: str,
        *,
        run_id: Any = None,
        **kwargs: Any,
    ) -> None:
        run_id_str = str(run_id) if run_id is not None else ""
        self._fire(
            self._client.log_event(
                type="tool_use",
                output={"result": output},
                metadata={"runId": run_id_str},
            )
        )

    def on_agent_finish(
        self,
        finish: Any,
        *,
        run_id: Any = None,
        **kwargs: Any,
    ) -> None:
        run_id_str = str(run_id) if run_id is not None else ""
        return_values = getattr(finish, "return_values", {}) or {}
        log = getattr(finish, "log", "")
        self._fire(
            self._client.log_event(
                type="chain_end",
                output=return_values,
                metadata={"runId": run_id_str, "log": log},
            )
        )
        self._current_workflow_id = None
