from __future__ import annotations

# standard
from enum import Enum
from dataclasses import dataclass, field
from typing import Any, Callable

# third party
# custom


class StopReason(Enum):
    END_TURN = "end_turn"
    TOOL_USE = "tool_use"
    MAX_TOKENS = "max_tokens"


class Role(Enum):
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"
    TOOL = "tool"


@dataclass
class ToolCall:
    id: str
    name: str
    arguments: dict[str, Any]
    error: str | None = field(default=None)
    duration: float = field(default=0.0)
    results: list[dict] = field(default_factory=list)


@dataclass
class Message:
    role: Role
    content: str | None = None
    reasoning_content: str | None = None
    # NOTE anthropic: signature // google: thoughtSignature
    reasoning_signature: str | None = None
    tool_call_id: str | None = None
    tool_calls: list[ToolCall] | None = None


@dataclass
class Tool:
    name: str
    description: str
    parameters: dict[str, Any]
    # NOTE callable function associated with the tool
    fn: Callable[..., Any] | None = field(default=None, repr=False)


def resolve_tokens(
    input_tokens: int, output_tokens: int, total_tokens: int
) -> tuple[int, int, int]:
    """Normalize token counts across providers.

    Some providers (xAI, Google with thinking) exclude reasoning tokens from
    output_tokens but include them in total_tokens. When total is provided and
    doesn't match input+output we trust total and back-compute output_tokens.
    """

    if total_tokens == 0:
        total_tokens = input_tokens + output_tokens
    # NOTE sometimes input + output != total tokens, so recompute output tokens
    elif input_tokens + output_tokens != total_tokens:
        output_tokens = total_tokens - input_tokens
    return input_tokens, output_tokens, total_tokens


@dataclass
class Error:
    message: str = ""
    status_code: int = 0


@dataclass
class Performance:
    latency: float = 0.0
    reasoning_duration: float = 0.0
    content_duration: float = 0.0
    total_duration: float = 0.0
    throughput: int = 0


@dataclass
class Cost:
    input: float = 0.0
    output: float = 0.0
    cache_read: float = 0.0
    cache_write: float = 0.0
    total: float = 0.0


@dataclass
class Usage:
    input_tokens: int = 0
    output_tokens: int = 0
    total_tokens: int = 0
    cache_read_tokens: int = 0
    cache_write_tokens: int = 0


def make_non_stream_performance(
    response: Response, total_duration: float
) -> Performance:
    """Estimate reasoning/content split from character counts (non-streaming only)."""

    r_len = len(response.reasoning_content or "")
    c_len = len(response.content or "")
    total_chars = r_len + c_len
    if total_chars > 0:
        reasoning_duration = r_len / total_chars * total_duration
        content_duration = c_len / total_chars * total_duration
    else:  # pragma: no cover
        reasoning_duration = 0.0
        content_duration = 0.0
    throughput = (
        int(response.usage.output_tokens / total_duration)
        if response.usage and total_duration > 0
        else 0
    )
    return Performance(
        latency=round(total_duration, 2),
        reasoning_duration=round(reasoning_duration, 2),
        content_duration=round(content_duration, 2),
        total_duration=round(total_duration, 2),
        throughput=throughput,
    )


@dataclass
class Response:
    provider: str = ""
    model: str = ""
    streaming: bool = False
    synthetic: bool = False
    content: str | None = None
    reasoning_content: str | None = None
    reasoning_signature: str | None = None
    tool_calls: list[ToolCall] | None = None
    stop_reason: StopReason | None = None
    error: Error | None = None
    cost: Cost | None = None
    performance: Performance | None = None
    usage: Usage | None = None
