from langchain_openai import OpenAIEmbeddings
from typing import List, Any
import time
from common.utils.business_event_context import current_event


class TrackedOpenAIEmbeddings(OpenAIEmbeddings):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        start_time = time.time()
        result = super().embed_documents(texts)
        end_time = time.time()

        # Estimate token usage (OpenAI uses tiktoken for this)
        import tiktoken
        enc = tiktoken.encoding_for_model(self.model)
        total_tokens = sum(len(enc.encode(text)) for text in texts)

        metrics = {
            'total_tokens': total_tokens,
            'prompt_tokens': total_tokens,  # For embeddings, all tokens are prompt tokens
            'completion_tokens': 0,
            'time_elapsed': end_time - start_time,
            'interaction_type': 'Embedding',
            }
        current_event.log_llm_metrics(metrics)

        return result

    def embed_query(self, text: str) -> List[float]:
        start_time = time.time()
        result = super().embed_query(text)
        end_time = time.time()

        # Estimate token usage
        import tiktoken
        enc = tiktoken.encoding_for_model(self.model)
        total_tokens = len(enc.encode(text))

        metrics = {
            'total_tokens': total_tokens,
            'prompt_tokens': total_tokens,
            'completion_tokens': 0,
            'time_elapsed': end_time - start_time,
            'interaction_type': 'Embedding',
        }
        current_event.log_llm_metrics(metrics)

        return result