import time
from typing import Any, Dict
try:
    from llama_index.core.callbacks import CBEventType
    from llama_index.core.callbacks.base_handler import BaseCallbackHandler
except ImportError:
    raise ImportError("Install llama-index: pip install telona[llamaindex]")

class LlamaIndexHandler(BaseCallbackHandler):
    def __init__(self, client):
        super().__init__(event_starts_to_ignore=[], event_ends_to_ignore=[])
        self.client = client
        self._traces: Dict[str, Dict[str, Any]] = {}

    def on_event_start(self, event_type, payload=None, event_id="", parent_id="", **kwargs):
        if event_type == CBEventType.QUERY:
            self._traces[event_id] = {"query": payload.get("query_str", "") if payload else "", "start_time": time.time()}

    def on_event_end(self, event_type, payload=None, event_id="", **kwargs):
        if event_type == CBEventType.RETRIEVE and payload:
            for tid, trace in self._traces.items():
                if "retrieved_docs" not in trace:
                    trace["retrieved_docs"] = [{"content": n.node.text, "score": n.score} for n in payload.get("nodes", [])]
                    break
        if event_type == CBEventType.QUERY and event_id in self._traces:
            trace = self._traces.pop(event_id)
            trace["latency_ms"] = (time.time() - trace.pop("start_time", time.time())) * 1000
            if payload: trace["response"] = payload.get("response", str(payload))
            self.client.monitor(**trace)

    def start_trace(self, trace_id=None): pass
    def end_trace(self, trace_id=None, trace_map=None): pass
