import time
from typing import Any, Dict
try:
    from langchain_core.callbacks import BaseCallbackHandler
except ImportError:
    raise ImportError("Install langchain: pip install telona[langchain]")

class LangChainCallback(BaseCallbackHandler):
    def __init__(self, client):
        self.client = client
        self._runs: Dict[str, Dict[str, Any]] = {}

    def on_chain_start(self, serialized, inputs, *, run_id, **kwargs):
        self._runs[str(run_id)] = {"query": inputs.get("query", inputs.get("question", str(inputs))), "start_time": time.time()}

    def on_retriever_end(self, documents, *, run_id, parent_run_id, **kwargs):
        pid = str(parent_run_id) if parent_run_id else None
        if pid and pid in self._runs:
            self._runs[pid]["retrieved_docs"] = [{"content": doc.page_content, "metadata": doc.metadata} for doc in documents]

    def on_llm_end(self, response, *, run_id, parent_run_id, **kwargs):
        pid = str(parent_run_id) if parent_run_id else None
        if pid and pid in self._runs:
            gen = response.generations[0][0] if response.generations else None
            if gen: self._runs[pid]["response"] = gen.text
            if hasattr(response, "llm_output") and response.llm_output:
                usage = response.llm_output.get("token_usage", {})
                self._runs[pid].update({"prompt_tokens": usage.get("prompt_tokens"), "completion_tokens": usage.get("completion_tokens"), "tokens_used": usage.get("total_tokens"), "model": response.llm_output.get("model_name")})

    def on_chain_end(self, outputs, *, run_id, **kwargs):
        rid = str(run_id)
        if rid in self._runs:
            run_data = self._runs.pop(rid)
            run_data["latency_ms"] = (time.time() - run_data.pop("start_time", time.time())) * 1000
            if "response" not in run_data: run_data["response"] = outputs.get("result", outputs.get("answer", str(outputs)))
            self.client.monitor(**run_data)
