import time
from openai import OpenAI as OpenAIClient
from .client import LlumoClient
from .llumoSessionContext import LlumoSessionContext
from .llumoLogger import LlumoLogger


# evaluation function that uses LlumoClient
def performEvaluation(data, api_key=None, evals=["Response Correctness"], **kwargs):
    try:
        client = LlumoClient(api_key=api_key)
        results = client.evaluateMultiple(
            data,
            evals=evals,
            prompt_template="Give answer to the query: {{query}}, using context: {{context}}",
            getDataFrame=False,
        )
        return results
    except Exception as e:
        print(f"Error in perform_evaluation: {e}")
        raise


# Wrapper around ChatCompletion to allow custom fields like `.evaluation`
class ChatCompletionWithEval:
    def __init__(
        self,
        response,
        #  , evaluation
    ):
        self._response = response
        # self.evaluation = evaluation

    def __getattr__(self, name):
        return getattr(self._response, name)

    def __getitem__(self, key):
        return self._response[key]

    def __repr__(self):
        return repr(self._response)


class OpenAI(OpenAIClient):
    def __init__(self, api_key: str, session):
        super().__init__(api_key=api_key)
        self.session = session
        self.llumo_key = session.apiKey

        original_create = self.chat.completions.create

        def create_wrapper(*args, **kwargs):
            context = kwargs.pop("context", None)
            model = kwargs["model"]
            create_experiment = kwargs.pop("createExperiment", False)

            messages = kwargs.get("messages", [])
            user_message = next(
                (
                    m.get("content")
                    for m in reversed(messages)
                    if m.get("role") == "user"
                ),
                "",
            )

            if not context or context.strip() == "":
                context = user_message

            # Get IDs from the session logger
            playground_id = self.session.logger.getPlaygroundID()
            workspace_id = self.session.logger.getWorkspaceID()

            # Input Bias Evaluation
            # eval_input_bias = [
            #     {
            #         "query": user_message,
            #         "context": context,
            #         "output": "",  # No output yet
            #     }
            # ]
            # try:
            #     start_time = time.time()
            #     bias_evaluation_result = performEvaluation(
            #         eval_input_bias,
            #         api_key=self.llumo_key,
            #         evals=["Input Bias"],
            #         playgroundID=playground_id,
            #         workspaceID=workspace_id,
            #         createExperiment=create_experiment,
            #     )
            #     latency = int((time.time() - start_time) * 1000)
            #     # Access the first result object
            #     bias_evaluation = bias_evaluation_result[0]
            #     message = "-".join(
            #         getattr(bias_evaluation, "edgeCases", {}).get("value", [])
            #     )
            #     self.session.logEvalStep(
            #         stepName=f"EVAL-Input Bias",
            #         output="",
            #         context=context,
            #         query=user_message,
            #         messageHistory="",
            #         tools="",
            #         intermediateSteps="",
            #         groundTruth="",
            #         analyticsScore=getattr(bias_evaluation, "analyticsScore", {}),
            #         reasoning=getattr(bias_evaluation, "reasoning", {}),
            #         classification=getattr(bias_evaluation, "classification", {}),
            #         evalLabel=getattr(bias_evaluation, "evalLabel", {}),
            #         latencyMs=latency,
            #         status="SUCCESS",
            #         message=message,
            #     )
            # except Exception as e:
            #     print(f"Input Bias evaluation failed: {e}")
            #     self.session.logEvalStep(
            #         stepName=f"EVAL-FAILURE",
            #         output="",
            #         context=context,
            #         query=user_message,
            #         messageHistory="",
            #         tools="",
            #         intermediateSteps="",
            #         groundTruth="",
            #         analyticsScore={},
            #         reasoning={},
            #         classification={},
            #         evalLabel={},
            #         latencyMs=0,
            #         status="FAILURE",
            #         message="EVAL_ERROR",
            #     )

            start_time = time.time()
            response = original_create(*args, **kwargs)
            latency = int((time.time() - start_time) * 1000)
            output_text = response.choices[0].message.content
            self.session.logQueryStep(
                stepName="Query Invocation",
                model=model,
                provider="openai",
                inputTokens=response.usage.prompt_tokens,
                query=user_message,
                status = "SUCCESS")

            self.session.logLlmStep(
                stepName=f"LLM-{user_message[:30]}",
                model=model,
                provider="openai",
                inputTokens=response.usage.prompt_tokens,
                outputTokens=response.usage.completion_tokens,
                # temperature=kwargs.get("temperature", 0.0),
                # promptTruncated=False,
                latencyMs=latency,
                prompt=user_message,
                output=output_text,
                status="SUCCESS",
                # message="",
            )

            # Response Correctness Evaluation
            # eval_input_correctness = [
            #     {
            #         "query": user_message,
            #         "context": context,
            #         "output": output_text,
            #     }
            # ]
            # try:
            #     start_time = time.time()
            #     correctness_evaluation_result = performEvaluation(
            #         eval_input_correctness,
            #         api_key=self.llumo_key,
            #         evals=["Response Correctness"],
            #         playgroundID=playground_id,
            #         workspaceID=workspace_id,
            #         createExperiment=create_experiment,
            #     )
            #     latency = int((time.time() - start_time) * 1000)
            #     # Access the first result object
            #     correctness_evaluation = correctness_evaluation_result[0]
            #     message = "-".join(
            #         getattr(correctness_evaluation, "edgeCases", {}).get("value", [])
            #     )
            #     self.session.logEvalStep(
            #         stepName=f"EVAL-Response Correctness",
            #         output=output_text,
            #         context=context,
            #         query=user_message,
            #         messageHistory="",
            #         tools="",
            #         intermediateSteps="",
            #         groundTruth="",
            #         analyticsScore=getattr(
            #             correctness_evaluation, "analyticsScore", {}
            #         ),
            #         reasoning=getattr(correctness_evaluation, "reasoning", {}),
            #         classification=getattr(
            #             correctness_evaluation, "classification", {}
            #         ),
            #         evalLabel=getattr(correctness_evaluation, "evalLabel", {}),
            #         latencyMs=latency,
            #         status="SUCCESS",
            #         message=message,
            #     )
            # except Exception as e:
            #     print(f"Response Correctness evaluation failed: {e}")
            #     correctness_evaluation = None
            #     self.session.logEvalStep(
            #         stepName=f"EVAL-FAILURE",
            #         output=output_text,
            #         context=context,
            #         query=user_message,
            #         messageHistory="",
            #         tools="",
            #         intermediateSteps="",
            #         groundTruth="",
            #         analyticsScore={},
            #         reasoning={},
            #         classification={},
            #         evalLabel={},
            #         latencyMs=0,
            #         status="FAILURE",
            #         message="EVAL_ERROR",
            #     )

            # if correctness_evaluation is None:
            #     return response

            return ChatCompletionWithEval(
                response
                # , correctness_evaluation
            )

        self.chat.completions.create = create_wrapper
