"""Event construction from OpenAI request/response pairs."""

from __future__ import annotations

import logging
import uuid
from datetime import datetime, timezone
from typing import Any

logger = logging.getLogger("llm_observatory")


def build_event(
    *,
    request_kwargs: dict[str, Any],
    response: Any,
    start_time: float,
    end_time: float,
    time_to_first_token_ms: int | None = None,
    tokens_per_second: float | None = None,
    error: dict[str, str] | None = None,
    metadata: dict[str, Any] | None = None,
    redact_messages: bool = False,
) -> dict[str, Any]:
    latency_ms = int((end_time - start_time) * 1000)

    # Build request section
    request = {
        "model": request_kwargs.get("model", ""),
        "messages": _maybe_redact(request_kwargs.get("messages"), redact_messages),
        "stream": request_kwargs.get("stream", False),
    }
    if "temperature" in request_kwargs:
        request["temperature"] = request_kwargs["temperature"]
    if "max_tokens" in request_kwargs:
        request["max_tokens"] = request_kwargs["max_tokens"]
    if "tools" in request_kwargs:
        request["tools"] = request_kwargs["tools"]

    # Build response section
    resp_section: dict[str, Any] = {}
    usage_section: dict[str, int] = {
        "prompt_tokens": 0,
        "completion_tokens": 0,
        "total_tokens": 0,
    }

    if error is not None:
        pass  # response may be None
    elif response is not None:
        resp_section = {
            "id": getattr(response, "id", ""),
            "model": getattr(response, "model", ""),
            "finish_reason": _extract_finish_reason(response),
        }
        choices_raw = getattr(response, "choices", None)
        if choices_raw is not None:
            resp_section["choices"] = [
                _choice_to_dict(c, redact_messages) for c in choices_raw
            ]

        usage = getattr(response, "usage", None)
        if usage is not None:
            usage_section = {
                "prompt_tokens": getattr(usage, "prompt_tokens", 0) or 0,
                "completion_tokens": getattr(usage, "completion_tokens", 0) or 0,
                "total_tokens": getattr(usage, "total_tokens", 0) or 0,
            }
            # Granular token breakdowns (OpenAI reasoning + cached tokens)
            prompt_details = getattr(usage, "prompt_tokens_details", None)
            if prompt_details is not None:
                cached = getattr(prompt_details, "cached_tokens", None)
                if cached:
                    usage_section["cached_tokens"] = cached
            completion_details = getattr(usage, "completion_tokens_details", None)
            if completion_details is not None:
                reasoning = getattr(completion_details, "reasoning_tokens", None)
                if reasoning:
                    usage_section["reasoning_tokens"] = reasoning
            # Some providers (e.g. Keywords AI) send reasoning/cached tokens
            # at the top level of usage instead of nested under *_details.
            if "reasoning_tokens" not in usage_section:
                top_reasoning = getattr(usage, "reasoning_tokens", None)
                if top_reasoning:
                    usage_section["reasoning_tokens"] = top_reasoning
            if "cached_tokens" not in usage_section:
                top_cached = getattr(usage, "cached_tokens", None)
                if top_cached:
                    usage_section["cached_tokens"] = top_cached
            # system_fingerprint for model determinism tracking
            sf = getattr(response, "system_fingerprint", None)
            if sf:
                usage_section["system_fingerprint"] = sf

    # Build metrics section
    metrics: dict[str, Any] = {"latency_ms": latency_ms}
    if time_to_first_token_ms is not None:
        metrics["time_to_first_token_ms"] = time_to_first_token_ms
    if tokens_per_second is not None:
        metrics["tokens_per_second"] = round(tokens_per_second, 1)

    # Build metadata section
    meta = metadata or {}
    metadata_section: dict[str, Any] = {}
    if meta.get("user_id"):
        metadata_section["user_id"] = meta["user_id"]
    if meta.get("session_id"):
        metadata_section["session_id"] = meta["session_id"]
    if meta.get("tags"):
        metadata_section["tags"] = meta["tags"]
    if meta.get("custom"):
        metadata_section["custom"] = {k: str(v) for k, v in meta["custom"].items()}

    event: dict[str, Any] = {
        "event_id": str(uuid.uuid4()),
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "request": request,
        "response": resp_section,
        "usage": usage_section,
        "metrics": metrics,
        "metadata": metadata_section,
        "span_type": "generation",
    }
    if error is not None:
        event["error"] = error

    # Pick up trace context from active span
    try:
        from llm_observatory.tracing import get_current_trace_id, get_current_span_id

        trace_id = get_current_trace_id()
        parent_id = get_current_span_id()
        if trace_id:
            event["trace_id"] = trace_id
        if parent_id:
            event["parent_id"] = parent_id
    except Exception as exc:
        logger.debug("llm_observatory: failed to pick up trace context: %s", exc)

    return event


def _extract_finish_reason(response: Any) -> str:
    choices = getattr(response, "choices", None)
    if choices and len(choices) > 0:
        return getattr(choices[0], "finish_reason", "") or ""
    return ""


def _choice_to_dict(choice: Any, redact: bool = False) -> dict[str, Any]:
    d: dict[str, Any] = {}
    if hasattr(choice, "index"):
        d["index"] = choice.index
    if hasattr(choice, "message"):
        msg = choice.message
        content = getattr(msg, "content", "")
        if redact and content:
            content = "[REDACTED]"
        d["message"] = {
            "role": getattr(msg, "role", ""),
            "content": content,
        }
        reasoning = getattr(msg, "reasoning_content", None)
        if reasoning:
            if redact:
                reasoning = "[REDACTED]"
            d["message"]["reasoning_content"] = reasoning
        if getattr(msg, "tool_calls", None):
            tool_calls_list = []
            for tc in msg.tool_calls:
                tc_dict: dict[str, Any] = {
                    "id": getattr(tc, "id", ""),
                    "type": getattr(tc, "type", "function"),
                    "function": {
                        "name": getattr(getattr(tc, "function", None), "name", ""),
                        "arguments": (
                            "[REDACTED]"
                            if redact
                            else getattr(
                                getattr(tc, "function", None), "arguments", ""
                            )
                        ),
                    },
                }
                tool_calls_list.append(tc_dict)
            d["message"]["tool_calls"] = tool_calls_list
    if hasattr(choice, "finish_reason"):
        d["finish_reason"] = choice.finish_reason
    return d


def _maybe_redact(messages: Any, redact: bool) -> Any:
    """Deep redaction of message content.

    Handles:
    - String content → "[REDACTED]"
    - List content (multi-part messages) → redacts each block's text/image_url
    - tool_calls[].function.arguments → "[REDACTED]" (preserves name)
    - function_call.arguments → "[REDACTED]"
    """
    if not redact or messages is None:
        return messages
    redacted = []
    for msg in messages:
        if isinstance(msg, dict):
            redacted.append(_redact_message(msg))
        else:
            redacted.append(msg)
    return redacted


def _redact_message(msg: dict[str, Any]) -> dict[str, Any]:
    """Redact a single message dict deeply."""
    out = dict(msg)

    # Redact content
    content = out.get("content")
    if isinstance(content, str):
        out["content"] = "[REDACTED]"
    elif isinstance(content, list):
        out["content"] = [_redact_content_block(block) for block in content]

    # Redact tool_calls
    if "tool_calls" in out and isinstance(out["tool_calls"], list):
        redacted_tcs = []
        for tc in out["tool_calls"]:
            if isinstance(tc, dict):
                tc_copy = dict(tc)
                fn = tc_copy.get("function")
                if isinstance(fn, dict):
                    tc_copy["function"] = {**fn, "arguments": "[REDACTED]"}
                redacted_tcs.append(tc_copy)
            else:
                redacted_tcs.append(tc)
        out["tool_calls"] = redacted_tcs

    # Redact function_call (legacy OpenAI format)
    if "function_call" in out and isinstance(out["function_call"], dict):
        out["function_call"] = {**out["function_call"], "arguments": "[REDACTED]"}

    return out


def _redact_content_block(block: Any) -> Any:
    """Redact a single content block in a multi-part message."""
    if not isinstance(block, dict):
        return block
    out = dict(block)
    block_type = out.get("type", "")
    if block_type == "text":
        out["text"] = "[REDACTED]"
    elif block_type == "image_url":
        image_url = out.get("image_url")
        if isinstance(image_url, dict):
            out["image_url"] = {**image_url, "url": "[REDACTED]"}
    return out
