"""
LLM utility functions for Noveum Trace SDK.

This module provides utility functions for working with LLM providers,
token counting, cost estimation, and metadata extraction.
"""

import re
from typing import Any, Optional, Union

from noveum_trace.utils import tokenizers
from noveum_trace.utils.model_registry import (
    MODEL_ALIASES,
    MODEL_REGISTRY,
    ModelInfo,
)

__all__ = [
    "ModelInfo",
    "MODEL_REGISTRY",
    "MODEL_ALIASES",
]


# MODEL_REGISTRY and MODEL_ALIASES are now in model_registry.py.
# They are imported above and re-exported so existing call sites are unaffected.


def detect_llm_provider(arguments: dict[str, Any]) -> Optional[str]:
    """
    Detect LLM provider from function arguments.

    Args:
        arguments: Function arguments

    Returns:
        Detected provider name or None
    """
    # Check for model parameter
    model = arguments.get("model", "")
    if isinstance(model, str) and model:
        model_info = get_model_info(model)
        if model_info:
            return model_info.provider

    # Check for provider-specific parameters
    if any(key in arguments for key in ["messages", "max_tokens", "temperature"]):
        # Common OpenAI parameters
        if "messages" in arguments and isinstance(arguments["messages"], list):
            return "openai"

    # Check for anthropic-specific patterns
    if any(key in arguments for key in ["max_tokens_to_sample", "stop_sequences"]):
        return "anthropic"

    # Check string representations for provider hints
    args_str = str(arguments).lower()
    if "anthropic" in args_str:
        return "anthropic"
    elif "openai" in args_str:
        return "openai"
    elif "google" in args_str or "gemini" in args_str:
        return "google"
    elif "claude" in args_str:
        return "anthropic"

    return None


def get_model_info(model_name: str) -> Optional[ModelInfo]:
    """
    Get model information by name.

    Args:
        model_name: Model name (with support for aliases)

    Returns:
        ModelInfo object or None if not found
    """
    if not isinstance(model_name, str):
        return None

    # Normalize model name
    normalized_name = normalize_model_name(model_name)

    # Check aliases first
    if normalized_name in MODEL_ALIASES:
        normalized_name = MODEL_ALIASES[normalized_name]

    # Return model info
    return MODEL_REGISTRY.get(normalized_name)


def normalize_model_name(model: str) -> str:
    """
    Normalize model name for consistent tracking.

    Args:
        model: Raw model name

    Returns:
        Normalized model name
    """
    if not isinstance(model, str):
        return str(model)

    # Remove version suffixes and normalize
    model = model.lower().strip()

    # Remove common prefixes/suffixes
    prefixes_to_remove = ["openai/", "anthropic/", "google/", "meta/", "microsoft/"]
    for prefix in prefixes_to_remove:
        if model.startswith(prefix):
            model = model[len(prefix) :]

    # Remove Bedrock-style version suffix (e.g. ":0", ":1") first, then the
    # "-vN" suffix, *before* the date regexes below. Bedrock IDs put the date
    # in the middle (e.g. "anthropic.claude-3-5-sonnet-20240620-v1:0"), so the
    # ":0" and "-v1" must be peeled off before the trailing-date regex can see
    # and strip the "-20240620" — otherwise the date is left in and the name
    # never matches the registry key.
    model = re.sub(r":\d+$", "", model)
    model = re.sub(r"-v\d+(\.\d+)*$", "", model)  # Remove version numbers

    # Handle date-based versions
    model = re.sub(r"-\d{8}$", "", model)  # Remove YYYYMMDD
    model = re.sub(r"-\d{4}-\d{2}-\d{2}$", "", model)  # Remove YYYY-MM-DD

    return model


def estimate_token_count(
    text: Union[str, list[dict[str, Any]], Any],
    *,
    model: Optional[str] = None,
    provider: Optional[str] = None,
) -> int:
    """Estimate token count for text or messages.

    Args:
        text: Text string, list of messages, or other content
        model: Optional model name to leverage provider-specific tokenizers
        provider: Optional provider override (e.g., ``"openai"``)

    Returns:
        Estimated token count
    """

    if isinstance(text, list):
        total_tokens = 0
        for message in text:
            if isinstance(message, dict):
                if "content" in message:
                    total_tokens += estimate_token_count(
                        message["content"], model=model, provider=provider
                    )
                total_tokens += 10  # Overhead for message structure
            else:
                total_tokens += estimate_token_count(
                    message, model=model, provider=provider
                )
        return max(1, total_tokens)

    if isinstance(text, dict):
        total_tokens = 0
        for key, value in text.items():
            total_tokens += estimate_token_count(key, model=model, provider=provider)
            total_tokens += estimate_token_count(value, model=model, provider=provider)
        return max(1, total_tokens)

    if not isinstance(text, str):
        text = str(text)

    try:
        provider_tokens = tokenizers.count_tokens(text, model=model, provider=provider)
    except Exception:
        provider_tokens = None

    if provider_tokens is not None:
        return provider_tokens

    if not text:
        return 1

    words = len(text.split())
    chars = len(text)

    if text.strip().startswith("{") and text.strip().endswith("}"):
        return max(1, int(chars / 3))
    if text.count("\n") > 5:
        return max(1, int(chars / 3.5))
    if len([w for w in text.split() if len(w) > 10]) > words * 0.3:
        return max(1, int(chars / 3))
    return max(1, int(chars / 4))


def estimate_cost(
    model: str,
    input_tokens: int,
    output_tokens: int,
) -> dict[str, Any]:
    """
    Estimate API cost based on model and token usage.

    Args:
        model: Model name
        input_tokens: Number of input tokens
        output_tokens: Number of output tokens

    Returns:
        Dictionary with cost information
    """
    model_info = get_model_info(model)

    if not model_info:
        # Fallback pricing for unknown models
        input_cost_per_1m = 1.0
        output_cost_per_1m = 2.0
        provider = "unknown"
    else:
        input_cost_per_1m = model_info.input_cost_per_1m
        output_cost_per_1m = model_info.output_cost_per_1m
        provider = model_info.provider

    # Calculate costs (prices are per 1M tokens)
    input_cost = (input_tokens / 1_000_000) * input_cost_per_1m
    output_cost = (output_tokens / 1_000_000) * output_cost_per_1m
    total_cost = input_cost + output_cost

    return {
        "input_cost": round(input_cost, 8),
        "output_cost": round(output_cost, 8),
        "total_cost": round(total_cost, 8),
        "currency": "USD",
        "model": model,
        "provider": provider,
        "input_cost_per_1m": input_cost_per_1m,
        "output_cost_per_1m": output_cost_per_1m,
    }


def _get_value(obj: Any, name: str) -> Any:
    if obj is None:
        return None
    if isinstance(obj, dict):
        return obj.get(name)
    return getattr(obj, name, None)


def _to_dict(obj: Any) -> dict[str, Any]:
    if obj is None:
        return {}
    if isinstance(obj, dict):
        return dict(obj)
    if hasattr(obj, "model_dump") and callable(obj.model_dump):
        return obj.model_dump()
    if hasattr(obj, "dict") and callable(obj.dict):
        return obj.dict()
    if hasattr(obj, "__dict__"):
        return dict(obj.__dict__)
    return {}


def _coerce_int(value: Any) -> Optional[int]:
    if value is None:
        return None
    if isinstance(value, bool):
        return int(value)
    if isinstance(value, (int, float)):
        return int(value)
    try:
        return int(str(value))
    except (TypeError, ValueError):
        return None


def _apply_usage_key_conversions(usage_model: dict[str, Any]) -> dict[str, Any]:
    conversion_list = [
        ("input_tokens", "input_tokens"),  # Anthropic
        ("output_tokens", "output_tokens"),
        ("total_tokens", "total_tokens"),
        ("prompt_tokens", "input_tokens"),  # OpenAI/Bedrock
        ("completion_tokens", "output_tokens"),
        ("prompt_token_count", "input_tokens"),  # Vertex AI
        ("candidates_token_count", "output_tokens"),
        ("total_token_count", "total_tokens"),
        ("inputTokenCount", "input_tokens"),  # Bedrock
        ("outputTokenCount", "output_tokens"),
        ("totalTokenCount", "total_tokens"),
        ("inputTokens", "input_tokens"),  # Bedrock Converse
        ("outputTokens", "output_tokens"),
        ("totalTokens", "total_tokens"),
        ("input_token_count", "input_tokens"),  # Watsonx
        ("generated_token_count", "output_tokens"),
    ]
    normalized = dict(usage_model)
    for source_key, target_key in conversion_list:
        if source_key in normalized and target_key not in normalized:
            normalized[target_key] = normalized.get(source_key)
    return normalized


def _set_usage_attr(
    usage_attrs: dict[str, Any], key: str, value: Any
) -> dict[str, Any]:
    if key in usage_attrs:
        return usage_attrs
    coerced = _coerce_int(value)
    if coerced is not None:
        usage_attrs[key] = coerced
    return usage_attrs


def _flatten_openai_token_details(usage_model: dict[str, Any]) -> dict[str, Any]:
    details: dict[str, Any] = {}
    prompt_details = usage_model.get("prompt_tokens_details") or {}
    if isinstance(prompt_details, dict):
        for key, value in prompt_details.items():
            details_key = f"llm.input_{key}"
            details[details_key] = _coerce_int(value)
    completion_details = usage_model.get("completion_tokens_details") or {}
    if isinstance(completion_details, dict):
        for key, value in completion_details.items():
            details_key = f"llm.output_{key}"
            details[details_key] = _coerce_int(value)
    return {k: v for k, v in details.items() if v is not None}


def _flatten_vertex_token_details(usage_model: dict[str, Any]) -> dict[str, Any]:
    details: dict[str, Any] = {}
    prompt_details = usage_model.get("prompt_tokens_details") or []
    for item in prompt_details if isinstance(prompt_details, list) else []:
        if not isinstance(item, dict):
            continue
        modality = item.get("modality")
        token_count = _coerce_int(item.get("token_count"))
        if modality and token_count is not None:
            details[f"llm.input_modality_{modality}"] = token_count
    candidates_details = usage_model.get("candidates_tokens_details") or []
    for item in candidates_details if isinstance(candidates_details, list) else []:
        if not isinstance(item, dict):
            continue
        modality = item.get("modality")
        token_count = _coerce_int(item.get("token_count"))
        if modality and token_count is not None:
            details[f"llm.output_modality_{modality}"] = token_count
    cache_details = usage_model.get("cache_tokens_details") or []
    for item in cache_details if isinstance(cache_details, list) else []:
        if not isinstance(item, dict):
            continue
        modality = item.get("modality")
        token_count = _coerce_int(item.get("token_count"))
        if modality and token_count is not None:
            details[f"llm.cache_modality_{modality}"] = token_count
    return details


def _extract_generation_usage_metadata(response: Any) -> list[dict[str, Any]]:
    sources: list[dict[str, Any]] = []
    generations = _get_value(response, "generations") or []
    for generation_list in generations:
        for gen in generation_list:
            generation_info = _get_value(gen, "generation_info")
            if not generation_info:
                generation_info = {}
            usage_metadata = _get_value(generation_info, "usage_metadata")
            if usage_metadata:
                sources.append(_to_dict(usage_metadata))
            message = _get_value(gen, "message")
            # Canonical langchain_core location: AIMessage.usage_metadata. Read it
            # before response_metadata so the normalized field wins (first source
            # to set a token key wins via _set_usage_attr).
            message_usage = _get_value(message, "usage_metadata")
            if message_usage:
                sources.append(_to_dict(message_usage))
            response_metadata = _get_value(message, "response_metadata")
            if response_metadata:
                metadata_dict = _to_dict(response_metadata)
                usage_from_message = metadata_dict.get("usage_metadata")
                if usage_from_message:
                    sources.append(_to_dict(usage_from_message))
    return sources


def _extract_response_usage_sources(response: Any) -> list[dict[str, Any]]:
    sources: list[dict[str, Any]] = []
    llm_output = _get_value(response, "llm_output")
    if llm_output:
        token_usage = _get_value(llm_output, "token_usage")
        if token_usage:
            sources.append(_to_dict(token_usage))
        output_usage = _get_value(llm_output, "usage")
        if output_usage:
            sources.append(_to_dict(output_usage))
        output_usage_metadata = _get_value(llm_output, "usage_metadata")
        if output_usage_metadata:
            sources.append(_to_dict(output_usage_metadata))
    sources.extend(_extract_generation_usage_metadata(response))
    usage = _get_value(response, "usage")
    if usage:
        sources.append(_to_dict(usage))
    usage_metadata = _get_value(response, "usage_metadata")
    if usage_metadata:
        sources.append(_to_dict(usage_metadata))
    meta = _get_value(response, "meta")
    if meta:
        tokens_info = _get_value(meta, "tokens")
        if tokens_info:
            sources.append(_to_dict(tokens_info))
    return sources


def _extract_token_details(usage_model: dict[str, Any]) -> dict[str, Any]:
    details: dict[str, Any] = {}
    if (
        "prompt_tokens_details" in usage_model
        or "completion_tokens_details" in usage_model
    ):
        details.update(_flatten_openai_token_details(usage_model))
    if (
        "prompt_tokens_details" in usage_model
        or "candidates_tokens_details" in usage_model
        or "cache_tokens_details" in usage_model
    ):
        details.update(_flatten_vertex_token_details(usage_model))
    return details


def parse_usage_from_response(
    response: Any, *, provider: Optional[str] = None, model: Optional[str] = None
) -> dict[str, Any]:
    """
    Parse token usage from a response object, normalized across providers.

    Returns attributes like:
    - llm.input_tokens
    - llm.output_tokens
    - llm.total_tokens
    - llm.input_cached_tokens
    - llm.output_reasoning_tokens
    - llm.input_modality_TEXT
    """
    _ = provider, model
    usage_attrs: dict[str, Any] = {}
    sources = _extract_response_usage_sources(response)
    for usage_model in sources:
        if not isinstance(usage_model, dict) or not usage_model:
            continue
        normalized = _apply_usage_key_conversions(usage_model)
        usage_attrs = _set_usage_attr(
            usage_attrs, "llm.input_tokens", normalized.get("input_tokens")
        )
        usage_attrs = _set_usage_attr(
            usage_attrs, "llm.output_tokens", normalized.get("output_tokens")
        )
        usage_attrs = _set_usage_attr(
            usage_attrs, "llm.total_tokens", normalized.get("total_tokens")
        )
        details = _extract_token_details(usage_model)
        for key, value in details.items():
            if key not in usage_attrs and value is not None:
                usage_attrs[key] = value
    return usage_attrs


def extract_llm_metadata(response: Any) -> dict[str, Any]:
    """
    Extract metadata from LLM response objects.

    Args:
        response: LLM response object

    Returns:
        Dictionary of extracted metadata
    """
    metadata: dict[str, Any] = {}

    def _get_value(obj: Any, name: str) -> Any:
        if obj is None:
            return None
        if isinstance(obj, dict):
            return obj.get(name)
        return getattr(obj, name, None)

    def _update_usage(usage_obj: Any) -> None:
        prompt = _get_value(usage_obj, "prompt_tokens")
        if prompt is not None:
            metadata.setdefault("llm.usage.prompt_tokens", prompt)
            metadata.setdefault("llm.usage.input_tokens", prompt)

        completion = _get_value(usage_obj, "completion_tokens")
        if completion is not None:
            metadata.setdefault("llm.usage.completion_tokens", completion)
            metadata.setdefault("llm.usage.output_tokens", completion)

        total = _get_value(usage_obj, "total_tokens")
        if total is not None:
            metadata.setdefault("llm.usage.total_tokens", total)

        input_tokens = _get_value(usage_obj, "input_tokens")
        if input_tokens is not None:
            metadata.setdefault("llm.usage.input_tokens", input_tokens)

        output_tokens = _get_value(usage_obj, "output_tokens")
        if output_tokens is not None:
            metadata.setdefault("llm.usage.output_tokens", output_tokens)

    def _update_usage_metadata(usage_obj: Any) -> None:
        prompt = _get_value(usage_obj, "prompt_token_count")
        if prompt is not None:
            metadata.setdefault("llm.usage.prompt_tokens", prompt)
            metadata.setdefault("llm.usage.input_tokens", prompt)

        completion = _get_value(usage_obj, "candidates_token_count")
        if completion is not None:
            metadata.setdefault("llm.usage.completion_tokens", completion)
            metadata.setdefault("llm.usage.output_tokens", completion)

        total = _get_value(usage_obj, "total_token_count")
        if total is not None:
            metadata.setdefault("llm.usage.total_tokens", total)

    usage = _get_value(response, "usage")
    if usage is None and isinstance(response, dict):
        usage = response.get("usage")
    if usage is not None:
        _update_usage(usage)

    usage_metadata = _get_value(response, "usage_metadata")
    if usage_metadata is None and isinstance(response, dict):
        usage_metadata = response.get("usage_metadata")
    if usage_metadata is not None:
        _update_usage_metadata(usage_metadata)

    meta = _get_value(response, "meta")
    if meta is not None:
        tokens_info = _get_value(meta, "tokens")
        if tokens_info is not None:
            input_tokens = _get_value(tokens_info, "input_tokens")
            if input_tokens is not None:
                metadata.setdefault("llm.usage.input_tokens", input_tokens)
            output_tokens = _get_value(tokens_info, "output_tokens")
            if output_tokens is not None:
                metadata.setdefault("llm.usage.output_tokens", output_tokens)
            prompt_tokens = _get_value(tokens_info, "prompt_tokens")
            if prompt_tokens is not None:
                metadata.setdefault("llm.usage.prompt_tokens", prompt_tokens)

    # Extract model information
    model_name = _get_value(response, "model")
    if model_name is not None:
        metadata["llm.model"] = model_name
        model_info = get_model_info(model_name)
        if model_info:
            metadata.setdefault("llm.provider", model_info.provider)
            metadata.setdefault("llm.context_window", model_info.context_window)
            metadata.setdefault("llm.max_output_tokens", model_info.max_output_tokens)

    # Extract system fingerprint for OpenAI
    system_fingerprint = _get_value(response, "system_fingerprint")
    if system_fingerprint is not None:
        metadata["llm.system_fingerprint"] = system_fingerprint

    # Extract finish reason
    choices = _get_value(response, "choices")
    if choices:
        choice = choices[0]
        finish_reason = _get_value(choice, "finish_reason")
        if finish_reason is not None:
            metadata["llm.finish_reason"] = finish_reason

    created = _get_value(response, "created")
    if created is not None:
        metadata["llm.created"] = created

    if "llm.usage.total_tokens" not in metadata:
        prompt_tokens = metadata.get("llm.usage.input_tokens") or metadata.get(
            "llm.usage.prompt_tokens"
        )
        completion_tokens = metadata.get("llm.usage.output_tokens") or metadata.get(
            "llm.usage.completion_tokens"
        )
        if prompt_tokens is not None and completion_tokens is not None:
            metadata["llm.usage.total_tokens"] = prompt_tokens + completion_tokens

    return metadata


def validate_model_compatibility(model: str, messages: list[Any]) -> dict[str, Any]:
    """
    Validate model compatibility and provide suggestions.

    Args:
        model: Model name
        messages: List of messages

    Returns:
        Dictionary with validation results and suggestions
    """
    model_info = get_model_info(model)

    result: dict[str, Any] = {
        "valid": True,
        "warnings": [],
        "suggestions": [],
        "model_info": model_info,
    }

    if not model_info:
        result["valid"] = False
        result["warnings"].append(f"Model '{model}' not found in registry")

        # Suggest similar models
        normalized = normalize_model_name(model)
        suggestions: list[str] = []
        for known_model in MODEL_REGISTRY.keys():
            if normalized in known_model or known_model in normalized:
                suggestions.append(known_model)

        if suggestions:
            result["suggestions"] = suggestions[:3]  # Top 3 suggestions
        else:
            # Fallback suggestions based on common patterns
            if "gpt" in normalized:
                result["suggestions"] = ["gpt-4o", "gpt-4", "gpt-3.5-turbo"]
            elif "claude" in normalized:
                result["suggestions"] = ["claude-3.5-sonnet", "claude-3.5-haiku"]
            elif "gemini" in normalized:
                result["suggestions"] = ["gemini-1.5-pro", "gemini-1.5-flash"]

    if model_info:
        # Estimate token usage
        estimated_tokens = estimate_token_count(
            messages, model=model, provider=model_info.provider
        )

        if estimated_tokens > model_info.context_window * 0.9:
            result["warnings"].append(
                f"Input may exceed context window ({estimated_tokens} tokens estimated, "
                f"limit is {model_info.context_window})"
            )

        # Check for vision content without vision support
        if not model_info.supports_vision:
            for message in messages:
                if isinstance(message, dict) and "content" in message:
                    content = message["content"]
                    if isinstance(content, list):
                        for item in content:
                            if (
                                isinstance(item, dict)
                                and item.get("type") == "image_url"
                            ):
                                result["warnings"].append(
                                    f"Model '{model}' does not support vision/image inputs"
                                )
                                break

    return result


def get_supported_models(provider: Optional[str] = None) -> list[str]:
    """
    Get list of supported models, optionally filtered by provider.

    Args:
        provider: Optional provider filter

    Returns:
        List of model names
    """
    if provider:
        return [
            name for name, info in MODEL_REGISTRY.items() if info.provider == provider
        ]
    return list(MODEL_REGISTRY.keys())


def extract_prompt_template_variables(prompt: str) -> list[str]:
    """
    Extract template variables from a prompt string.

    Args:
        prompt: Prompt string with template variables

    Returns:
        List of variable names found in the prompt
    """
    variables = set()

    # Find variables in {variable} format
    variables.update(re.findall(r"\{([^}]+)\}", prompt))

    # Find variables in {{variable}} format (Jinja2 style)
    variables.update(re.findall(r"\{\{([^}]+)\}\}", prompt))

    # Find variables in ${variable} format
    variables.update(re.findall(r"\$\{([^}]+)\}", prompt))

    return [var.strip() for var in variables]


def sanitize_llm_content(content: str) -> str:
    """
    Sanitize LLM content for safe logging.

    Args:
        content: Content to sanitize

    Returns:
        Sanitized content
    """
    if not isinstance(content, str):
        content = str(content)

    # Remove or mask sensitive patterns
    patterns = [
        (r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b", "[EMAIL]"),
        (r"\b\d{3}[-.]?\d{3}[-.]?\d{4}\b", "[PHONE]"),
        (r"\b\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b", "[CARD]"),
        (r"\b[A-Z0-9]{20,}\b", "[TOKEN]"),  # API keys, tokens
        (r"sk-[a-zA-Z0-9]{32,}", "[API_KEY]"),  # OpenAI API keys
        (r"xoxb-[0-9]{12}-[0-9]{12}-[a-zA-Z0-9]{24}", "[SLACK_TOKEN]"),
    ]

    for pattern, replacement in patterns:
        content = re.sub(pattern, replacement, content)

    return content