"""Utility functions for improvement mechanisms."""
from __future__ import annotations

from datetime import UTC, datetime
from fnmatch import fnmatch
from math import exp
from typing import Any

from logos_adk.cost import DEFAULT_RATES


def utcnow() -> str:
    """Return current UTC timestamp in ISO format."""
    return datetime.now(UTC).isoformat()


def _mean(values: list[float], default: float = 0.0) -> float:
    """Calculate mean of values."""
    if not values:
        return default
    return sum(values) / len(values)


def _estimate_cost_usd(model: str | None, usage: dict[str, Any] | None) -> float:
    """Estimate LLM cost in USD."""
    if not model or not usage:
        return 0.0
    input_tokens = int(usage.get("input_tokens", 0) or 0)
    output_tokens = int(usage.get("output_tokens", 0) or 0)
    rates = None
    if model in DEFAULT_RATES:
        rates = DEFAULT_RATES[model]
    else:
        for pattern, candidate in DEFAULT_RATES.items():
            if "*" in pattern and fnmatch(model, pattern):
                rates = candidate
                break
    if rates is None:
        return 0.0
    return round((input_tokens * rates[0] + output_tokens * rates[1]) / 1_000_000, 8)


def _human_feedback_score(feedback: list[Any]) -> float:
    """Calculate human feedback score from feedback records."""
    values: list[float] = []
    for item in feedback:
        if item.kind == "thumbs_up":
            values.append(1.0)
        elif item.kind == "thumbs_down":
            values.append(0.0)
        elif item.kind == "numeric_rating" and item.rating is not None:
            values.append(max(0.0, min(1.0, (float(item.rating) - 1.0) / 4.0)))
        elif item.kind == "edited_answer":
            values.append(0.15)
        elif item.kind == "suggestion":
            values.append(0.8 if item.suggestion_outcome == "accepted" else 0.25)
        if item.failure_types:
            values.append(max(0.0, 1.0 - (0.2 * len(item.failure_types))))
    return round(_mean(values, default=0.5), 4)


def _business_outcome_score(outcomes: list[Any]) -> float:
    """Calculate business outcome score from outcome records."""
    values: list[float] = []
    for item in outcomes:
        if item.kind in {"conversion", "resolution"}:
            values.append(1.0 if bool(item.value) else 0.0)
        elif item.kind in {"ctr", "reply_rate"} and isinstance(item.value, (int, float)):
            values.append(max(0.0, min(1.0, float(item.value))))
        elif item.kind in {"escalation", "human_takeover"}:
            values.append(0.0 if bool(item.value) else 1.0)
    return round(_mean(values, default=0.5), 4)


def _clip_unit(value: float) -> float:
    """Clip value to [0.0, 1.0] range."""
    return max(0.0, min(1.0, float(value)))


def _reward_feature_values(example: Any) -> dict[str, float]:
    """Extract reward feature values from learning example."""
    metadata = example.metadata if isinstance(example.metadata, dict) else {}
    usage = metadata.get("usage") if isinstance(metadata, dict) else None
    cost_penalty = min(1.0, _estimate_cost_usd(example.model, usage) * 100)
    if cost_penalty == 0.0:
        cost_penalty = min(1.0, float(metadata.get("estimated_cost_usd", 0.0) or 0.0) * 100)
    latency_ms = float(metadata.get("latency_ms", 0.0) or 0.0)
    latency_penalty = min(1.0, latency_ms / 10_000.0)
    safety_penalty = 1.0 if "policy_violation" in example.failure_types else 0.0
    correction_penalty = 0.0
    if example.corrected_output or any(item.expected_output for item in example.feedback):
        correction_penalty = 1.0
    elif any(item.kind == "edited_answer" for item in example.feedback):
        correction_penalty = 0.7
    return {
        "human_feedback": _human_feedback_score(example.feedback),
        "business_outcome": _business_outcome_score(example.outcomes),
        "cost_penalty": cost_penalty,
        "latency_penalty": latency_penalty,
        "safety_penalty": safety_penalty,
        "correction_penalty": correction_penalty,
    }


def _preference_target(example: Any) -> float:
    """Calculate preference target from learning example."""
    explicit_values: list[float] = []
    for item in example.feedback:
        if item.kind == "thumbs_up":
            explicit_values.append(1.0)
        elif item.kind == "thumbs_down":
            explicit_values.append(0.0)
        elif item.kind == "numeric_rating" and item.rating is not None:
            explicit_values.append(_clip_unit((float(item.rating) - 1.0) / 4.0))
        elif item.kind == "edited_answer":
            explicit_values.append(0.1)
        elif item.kind == "suggestion":
            explicit_values.append(0.85 if item.suggestion_outcome == "accepted" else 0.15)
        if item.corrected_output or item.expected_output:
            explicit_values.append(0.1)
    for item in example.outcomes:
        if item.kind in {"conversion", "resolution"}:
            explicit_values.append(1.0 if bool(item.value) else 0.0)
        elif item.kind in {"ctr", "reply_rate"} and isinstance(item.value, (int, float)):
            explicit_values.append(_clip_unit(float(item.value)))
        elif item.kind in {"escalation", "human_takeover"}:
            explicit_values.append(0.0 if bool(item.value) else 1.0)
    if explicit_values:
        return round(_mean(explicit_values, default=0.5), 4)
    features = _reward_feature_values(example)
    fallback = (
        0.45 * features["human_feedback"]
        + 0.35 * features["business_outcome"]
        + 0.08 * (1.0 - features["cost_penalty"])
        + 0.06 * (1.0 - features["latency_penalty"])
        + 0.06 * (1.0 - max(features["safety_penalty"], features["correction_penalty"]))
    )
    return round(_clip_unit(fallback), 4)


def _sigmoid(value: float) -> float:
    """Sigmoid activation function."""
    if value >= 0:
        exp_value = exp(-value)
        return 1.0 / (1.0 + exp_value)
    exp_value = exp(value)
    return exp_value / (1.0 + exp_value)


__all__ = [
    "utcnow",
    "_mean",
    "_estimate_cost_usd",
    "_human_feedback_score",
    "_business_outcome_score",
    "_clip_unit",
    "_reward_feature_values",
    "_preference_target",
    "_sigmoid",
]
