"""@private"""

import logging
import typing
from datetime import datetime, timezone

try:
    import pydantic.v1 as pydantic  # type: ignore
except ImportError:
    import pydantic  # type: ignore

from langfuse.model import ModelUsage, PromptClient

log = logging.getLogger("langfuse")


def _get_timestamp():
    return datetime.now(timezone.utc)


def _create_prompt_context(
    prompt: typing.Optional[PromptClient] = None,
):
    if prompt is not None and not prompt.is_fallback:
        return {"prompt_version": prompt.version, "prompt_name": prompt.name}

    return {"prompt_version": None, "prompt_name": None}


T = typing.TypeVar("T")


def extract_by_priority(
    usage: dict, keys: typing.List[str], target_type: typing.Type[T]
) -> typing.Optional[T]:
    """Extracts the first key that exists in usage and converts its value to target_type"""
    for key in keys:
        if key in usage:
            value = usage[key]
            try:
                if value is None:
                    return None
                return target_type(value)
            except Exception:
                continue
    return None


def _convert_usage_input(usage: typing.Union[pydantic.BaseModel, ModelUsage]):
    """Converts any usage input to a usage object"""
    if isinstance(usage, pydantic.BaseModel):
        usage = usage.dict()

    # sometimes we do not match the pydantic usage object
    # in these cases, we convert to dict manually
    if hasattr(usage, "__dict__"):
        usage = usage.__dict__

    # validate that usage object has input, output, total, usage
    is_langfuse_usage = any(k in usage for k in ("input", "output", "total", "unit"))

    if is_langfuse_usage:
        return usage

    is_openai_usage = any(
        k in usage
        for k in (
            "promptTokens",
            "prompt_tokens",
            "completionTokens",
            "completion_tokens",
            "totalTokens",
            "total_tokens",
            "inputCost",
            "input_cost",
            "outputCost",
            "output_cost",
            "totalCost",
            "total_cost",
        )
    )

    if is_openai_usage:
        # convert to langfuse usage
        usage = {
            "input": extract_by_priority(usage, ["promptTokens", "prompt_tokens"], int),
            "output": extract_by_priority(
                usage,
                ["completionTokens", "completion_tokens"],
                int,
            ),
            "total": extract_by_priority(usage, ["totalTokens", "total_tokens"], int),
            "unit": "TOKENS",
            "inputCost": extract_by_priority(usage, ["inputCost", "input_cost"], float),
            "outputCost": extract_by_priority(
                usage, ["outputCost", "output_cost"], float
            ),
            "totalCost": extract_by_priority(usage, ["totalCost", "total_cost"], float),
        }
        return usage

    if not is_langfuse_usage and not is_openai_usage:
        raise ValueError(
            "Usage object must have either {input, output, total, unit} or {promptTokens, completionTokens, totalTokens}"
        )