import logging
import threading
from typing import Dict, Optional, Any

from mca_sdk.core.client import MCAClient, _sanitize_error_message

logger = logging.getLogger(__name__)

# Constants for cost estimation (per 1k tokens)
# These should ideally be fetched from a config or registry,
# but hardcoded here for the prototype phase as requested.
OPENAI_PRICING = {
    "gpt-3.5-turbo": {"prompt": 0.0015, "completion": 0.002},
    "gpt-3.5-turbo-16k": {"prompt": 0.003, "completion": 0.004},
    "gpt-4": {"prompt": 0.03, "completion": 0.06},
    "gpt-4-32k": {"prompt": 0.06, "completion": 0.12},
    "text-embedding-ada-002": {"prompt": 0.0001, "completion": 0.0},
}


def create_genai_client(
    service_name: str,
    model_id: str,
    team_name: Optional[str] = None,
    capture_phi: bool = True,
    **kwargs,
) -> MCAClient:
    """Create a pre-configured MCAClient specifically for GenAI workloads.

    This helper function sets up an MCAClient with defaults optimized for
    Generative AI, such as forcing the taxonomy category to "generative".

    Args:
        service_name: Name of the consuming application/service
        model_id: Unique identifier for the foundation model
        team_name: Optional team identifier for billing/tracking
        capture_phi: Whether to capture and log prompt/completion text
            If True, payloads are logged. If False, payloads are dropped.
            NOTE: Dropping PHI requires downstream components to respect the flag.
        **kwargs: Additional arguments to pass to MCAClient constructor

    Returns:
        Configured MCAClient instance
    """
    # Force taxonomy for GenAI to ensure proper routing
    kwargs["model_category"] = kwargs.get("model_category", "vendor")
    kwargs["model_type"] = "generative"

    client = MCAClient(service_name=service_name, model_id=model_id, team_name=team_name, **kwargs)

    # Attach phi flag to client config dynamically for use in decorators
    client._capture_phi = capture_phi

    return client


def estimate_openai_cost(model: str, prompt_tokens: int, completion_tokens: int) -> float:
    """Estimate cost in USD based on OpenAI pricing tiers.

    Args:
        model: Model name (e.g. 'gpt-4')
        prompt_tokens: Number of tokens sent in prompt
        completion_tokens: Number of tokens received in response

    Returns:
        Estimated cost in USD
    """
    # Default to gpt-3.5 pricing if model unknown
    pricing = OPENAI_PRICING.get(model, OPENAI_PRICING["gpt-3.5-turbo"])

    prompt_cost = (prompt_tokens / 1000) * pricing["prompt"]
    completion_cost = (completion_tokens / 1000) * pricing["completion"]

    return prompt_cost + completion_cost


# Global cache for pre-created OpenTelemetry instruments.
# Uses client ID as key to support multiple clients safely.
_instruments_cache: Dict[int, Dict[str, Any]] = {}
_instruments_lock = threading.Lock()


def _get_instruments(client: MCAClient) -> Dict[str, Any]:
    """Retrieve or initialize pre-created OpenTelemetry instruments for high performance.

    This fixes the major performance bug where `.record_metric()` was dynamically creating
    instruments on every single LLM call.
    """
    client_id = id(client)

    # We must explicitly clean up the cache if a client is destroyed to prevent memory leaks.
    # However, Python's GC is non-deterministic. For now, we rely on the fact that
    # long-lived applications use a singleton client.

    with _instruments_lock:
        if client_id not in _instruments_cache:
            prefix = getattr(client, "_metric_prefix", "genai")
            if not isinstance(prefix, str):
                prefix = "genai"

            meter = client.meter

            # Create all metric instruments once with descriptions and units matching LitellmCallback
            instruments = {
                "requests_total": meter.create_counter(
                    f"{prefix}.requests_total",
                    description="Total number of GenAI requests",
                ),
                "requests_failed_total": meter.create_counter(
                    f"{prefix}.requests_failed_total",
                    description="Total number of failed GenAI requests",
                ),
                "tokens_total": meter.create_counter(
                    f"{prefix}.tokens_total",
                    description="Total tokens consumed. Use token_type attribute to differentiate prompt/completion.",
                ),
                "latency_seconds": meter.create_histogram(
                    f"{prefix}.latency_seconds",
                    description="GenAI request latency",
                    unit="s",
                ),
                "cost_dollars": meter.create_histogram(
                    f"{prefix}.cost_dollars",
                    description="Estimated cost of GenAI requests",
                    unit="usd",
                ),
                # Embedding metrics
                "embedding_requests_total": meter.create_counter(
                    f"{prefix}.embedding_requests_total",
                    description="Total number of embedding requests",
                ),
                "embedding_requests_failed_total": meter.create_counter(
                    f"{prefix}.embedding_requests_failed_total",
                    description="Total number of failed embedding requests",
                ),
                "embedding_tokens_total": meter.create_counter(
                    f"{prefix}.embedding_tokens_total",
                    description="Total tokens consumed for embeddings",
                ),
                "embedding_latency_seconds": meter.create_histogram(
                    f"{prefix}.embedding_latency_seconds",
                    description="Embedding request latency",
                    unit="s",
                ),
                "embedding_dimensions_count": meter.create_histogram(
                    f"{prefix}.embedding_dimensions_count",
                    description="Dimensions returned by embedding model",
                ),
            }
            _instruments_cache[client_id] = instruments

        return _instruments_cache[client_id]


def track_llm_completion(
    client: MCAClient,
    model: str,
    prompt_tokens: int,
    completion_tokens: int,
    latency: float,
    status: str = "success",
    error: Optional[str] = None,
    provider: str = "openai",
):
    """Track LLM completion metrics.

    Records standard GenAI metrics for an LLM completion request.

    Args:
        client: MCAClient instance
        model: LLM model name
        prompt_tokens: Number of prompt tokens
        completion_tokens: Number of completion tokens
        latency: Request latency in seconds
        status: Request status ("success" or "error")
        error: Error message if status is "error"
        provider: The provider platform (e.g. "openai", "google"). Defaults to "openai" for backward compatibility.
    """
    # Get metric prefix from client
    prefix = getattr(client, "_metric_prefix", "genai")

    # Common attributes for all metrics
    attrs = {"model": model, "provider": provider}

    if status == "error" or error is not None:
        # Failure path: only record the failure to avoid skewing stats and charging for errors
        fail_attrs = {**attrs, "status": "error"}

        # Record error metric using client.record_metric for consistency
        client.record_metric(f"{prefix}.errors_total", 1, **fail_attrs)

        # Log the error safely
        error_msg = _sanitize_error_message(str(error)) if error else "Unknown error"
        client._sdk_logger.error(f"GenAI request failed. Model: {model}. Error: {error_msg}")
        return

    # Success path: Record all positive metrics
    success_attrs = {**attrs, "status": "success"}

    client.record_metric(f"{prefix}.requests_total", 1, **success_attrs)

    # Token Counts - minimal attributes {model, token_type} to distinguish from litellm_callback path
    total_tokens = prompt_tokens + completion_tokens
    client.record_metric(f"{prefix}.tokens_total", total_tokens, token_type="total", model=model)
    client.record_metric(f"{prefix}.tokens_total", prompt_tokens, token_type="prompt", model=model)
    client.record_metric(
        f"{prefix}.tokens_total", completion_tokens, token_type="completion", model=model
    )

    # Latency - include provider for filtering
    client.record_metric(f"{prefix}.latency_seconds", latency, model=model, provider=provider)

    # Cost Estimation - include provider for filtering
    if model in OPENAI_PRICING:
        cost = estimate_openai_cost(model, prompt_tokens, completion_tokens)
        client.record_metric(f"{prefix}.cost_dollars", cost, model=model, provider=provider)


def calculate_token_rate(
    prompt_tokens: int, completion_tokens: int, latency: float
) -> Dict[str, float]:
    """Calculate token processing rates (tokens per second).

    Args:
        prompt_tokens: Number of prompt tokens
        completion_tokens: Number of completion tokens
        latency: Request latency in seconds

    Returns:
        Dict containing total_rate, prompt_rate, and completion_rate
    """
    if latency <= 0:
        return {"total_rate": 0.0, "prompt_rate": 0.0, "completion_rate": 0.0}

    total_tokens = prompt_tokens + completion_tokens

    return {
        "total_rate": total_tokens / latency,
        "prompt_rate": prompt_tokens / latency,
        "completion_rate": completion_tokens / latency,
    }


def track_llm_embedding(
    client: MCAClient,
    model: str,
    tokens: int,
    latency: float,
    dimensions: Optional[int] = None,
    status: str = "success",
    error: Optional[str] = None,
    provider: str = "openai",
):
    """Track metrics specifically for text embedding models.

    Args:
        client: MCAClient instance
        model: Embedding model name (e.g. 'text-embedding-ada-002')
        tokens: Number of tokens processed
        latency: Request latency in seconds
        dimensions: Number of dimensions in resulting embedding vector
        status: Request status ("success" or "error")
        error: Error message if status is "error"
        provider: The provider platform. Defaults to "openai".
    """
    # Get metric prefix from client
    prefix = getattr(client, "_metric_prefix", "genai")

    attrs = {"model": model, "provider": provider}

    if status == "error" or error is not None:
        fail_attrs = {**attrs, "status": "error"}
        client.record_metric(f"{prefix}.embedding_errors_total", 1, **fail_attrs)
        error_msg = _sanitize_error_message(str(error)) if error else "Unknown error"
        client._sdk_logger.error(f"Embedding request failed. Model: {model}. Error: {error_msg}")
        return

    success_attrs = {**attrs, "status": "success"}
    client.record_metric(f"{prefix}.embedding_requests_total", 1, **success_attrs)
    client.record_metric(f"{prefix}.embedding_tokens_total", tokens, **attrs)
    client.record_metric(f"{prefix}.embedding_latency_seconds", latency, **attrs)

    if dimensions:
        client.record_metric(f"{prefix}.embedding_dimensions_count", dimensions, **attrs)
