"""
Prometheus metrics for Kubiya Control Plane.

This module provides comprehensive metrics for monitoring the AI orchestration layer:
- HTTP request metrics (latency, error rates, request counts)
- Execution metrics (active count, failures, duration, wait time)
- Worker and queue metrics
- LLM/AI metrics (requests, latency, tokens)
- Tool execution metrics
- Business metrics (agents, jobs, organizations)

Metrics naming convention follows Prometheus best practices:
- Prefix: kubiya_control_plane_ for HTTP metrics
- Prefix: kubiya_ for business metrics
- Suffix: _total for counters, _seconds for durations, _count for gauges
"""

import re
import os
import structlog
from typing import Tuple, Optional
from functools import lru_cache

from prometheus_client import (
    Counter,
    Histogram,
    Gauge,
    REGISTRY,
    generate_latest,
    CONTENT_TYPE_LATEST,
    CollectorRegistry,
)

logger = structlog.get_logger(__name__)

PROMETHEUS_MULTIPROC_DIR = os.environ.get('PROMETHEUS_MULTIPROC_DIR')

HTTP_REQUESTS_TOTAL = Counter(
    'kubiya_control_plane_http_requests_total',
    'Total number of HTTP requests received by the Control Plane API',
    ['method', 'endpoint', 'status_code'],
)

HTTP_REQUEST_DURATION_SECONDS = Histogram(
    'kubiya_control_plane_http_request_duration_seconds',
    'HTTP request duration in seconds (latency)',
    ['method', 'endpoint'],
    buckets=(0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0, 120.0),
)


ACTIVE_TASKS_COUNT = Gauge(
    'kubiya_active_tasks_count',
    'Number of currently active tasks/executions by type and status',
    ['execution_type', 'status'],
)

TASK_FAILURE_TOTAL = Gauge(
    'kubiya_task_failure_total',
    'Total number of failed task/executions by type',
    ['execution_type'],
)


# Execution duration - how long executions take
EXECUTION_DURATION_SECONDS = Gauge(
    'kubiya_execution_duration_seconds',
    'Average execution duration in seconds by type and status (from completed executions)',
    ['execution_type', 'status'],
)

# Worker queue depth - pending executions per queue
WORKER_QUEUE_DEPTH = Gauge(
    'kubiya_worker_queue_depth',
    'Number of pending executions waiting in queue',
    ['queue_id'],
)

# LLM requests total
LLM_REQUESTS_TOTAL = Gauge(
    'kubiya_llm_requests_total',
    'Total number of LLM API requests by model and status',
    ['model', 'status'],
)

# LLM latency
LLM_LATENCY_SECONDS = Gauge(
    'kubiya_llm_latency_seconds',
    'Average LLM request latency in seconds by model',
    ['model'],
)

# LLM tokens total
LLM_TOKENS_TOTAL = Gauge(
    'kubiya_llm_tokens_total',
    'Total tokens used in LLM requests by model and type',
    ['model', 'token_type'],
)


# Active streaming connections
STREAMING_CONNECTIONS_ACTIVE = Gauge(
    'kubiya_streaming_connections_active',
    'Number of active SSE streaming connections',
)

# Tool execution duration
TOOL_EXECUTION_DURATION_SECONDS = Gauge(
    'kubiya_tool_execution_duration_seconds',
    'Average tool execution duration in seconds by tool name',
    ['tool_name'],
)

# Tool executions total
TOOL_EXECUTIONS_TOTAL = Gauge(
    'kubiya_tool_executions_total',
    'Total number of tool executions by tool name and status',
    ['tool_name', 'status'],
)

# Execution wait time (time from created to running)
EXECUTION_WAIT_TIME_SECONDS = Gauge(
    'kubiya_execution_wait_time_seconds',
    'Average wait time from creation to running in seconds by type',
    ['execution_type'],
)

# Webhook requests total
WEBHOOK_REQUESTS_TOTAL = Gauge(
    'kubiya_webhook_requests_total',
    'Total number of webhook trigger requests by status',
    ['status'],
)


# Executions by organization
EXECUTIONS_BY_ORG_TOTAL = Gauge(
    'kubiya_executions_by_org_total',
    'Total number of executions by organization',
    ['organization_id'],
)

# Active agents
AGENTS_ACTIVE = Gauge(
    'kubiya_agents_active',
    'Number of active agents by organization',
    ['organization_id'],
)

# Scheduled jobs total
SCHEDULED_JOBS_TOTAL = Gauge(
    'kubiya_scheduled_jobs_total',
    'Total number of scheduled jobs by status',
    ['status'],
)

# Conversation turns total
CONVERSATION_TURNS_TOTAL = Gauge(
    'kubiya_conversation_turns_total',
    'Total conversation turns by execution type',
    ['execution_type'],
)


UUID_PATTERN = re.compile(
    r'[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}',
    re.IGNORECASE
)
NUMERIC_ID_PATTERN = re.compile(r'/\d+(?=/|$)')


@lru_cache(maxsize=1000)
def normalize_endpoint(path: str) -> str:
    """Normalize endpoint paths to reduce cardinality in metrics."""
    normalized = UUID_PATTERN.sub('{id}', path)
    normalized = NUMERIC_ID_PATTERN.sub('/{id}', normalized)
    return normalized


def get_metrics_response() -> Tuple[bytes, str]:
    """Generate Prometheus metrics response."""
    try:
        if PROMETHEUS_MULTIPROC_DIR:
            from prometheus_client.multiprocess import MultiProcessCollector
            registry = CollectorRegistry()
            MultiProcessCollector(registry)
            output = generate_latest(registry)
        else:
            output = generate_latest(REGISTRY)
        return output, CONTENT_TYPE_LATEST
    except Exception as e:
        logger.error("metrics_generation_failed", error=str(e))
        return b"# Error generating metrics\n", CONTENT_TYPE_LATEST


def record_http_request(
    method: str,
    endpoint: str,
    status_code: int,
    duration_seconds: float,
    trace_id: Optional[str] = None,
):
    """Record an HTTP request in Prometheus metrics."""
    HTTP_REQUESTS_TOTAL.labels(
        method=method,
        endpoint=endpoint,
        status_code=str(status_code),
    ).inc()
    
    HTTP_REQUEST_DURATION_SECONDS.labels(
        method=method,
        endpoint=endpoint,
    ).observe(duration_seconds)
    
    if status_code >= 500:
        logger.error(
            "prometheus_5xx_recorded",
            method=method,
            endpoint=endpoint,
            status_code=status_code,
            duration_seconds=round(duration_seconds, 4),
            trace_id=trace_id or "unknown",
        )


def update_active_tasks(execution_type: str, status: str, count: int):
    """Update the active tasks gauge."""
    ACTIVE_TASKS_COUNT.labels(execution_type=execution_type, status=status).set(count)


def update_task_failures(execution_type: str, count: int):
    """Update the task failures gauge."""
    TASK_FAILURE_TOTAL.labels(execution_type=execution_type).set(count)


def update_execution_duration(execution_type: str, status: str, avg_seconds: float):
    """Update average execution duration."""
    EXECUTION_DURATION_SECONDS.labels(execution_type=execution_type, status=status).set(avg_seconds)


def update_worker_queue_depth(queue_id: str, count: int):
    """Update worker queue depth."""
    WORKER_QUEUE_DEPTH.labels(queue_id=queue_id).set(count)


def update_llm_requests(model: str, status: str, count: int):
    """Update LLM requests total."""
    LLM_REQUESTS_TOTAL.labels(model=model, status=status).set(count)


def update_llm_latency(model: str, avg_seconds: float):
    """Update average LLM latency."""
    LLM_LATENCY_SECONDS.labels(model=model).set(avg_seconds)


def update_llm_tokens(model: str, token_type: str, count: int):
    """Update LLM tokens total."""
    LLM_TOKENS_TOTAL.labels(model=model, token_type=token_type).set(count)


def update_streaming_connections(count: int):
    """Update active streaming connections."""
    STREAMING_CONNECTIONS_ACTIVE.set(count)


def update_tool_execution_duration(tool_name: str, avg_seconds: float):
    """Update average tool execution duration."""
    TOOL_EXECUTION_DURATION_SECONDS.labels(tool_name=tool_name).set(avg_seconds)


def update_tool_executions(tool_name: str, status: str, count: int):
    """Update tool executions total."""
    TOOL_EXECUTIONS_TOTAL.labels(tool_name=tool_name, status=status).set(count)


def update_execution_wait_time(execution_type: str, avg_seconds: float):
    """Update average execution wait time."""
    EXECUTION_WAIT_TIME_SECONDS.labels(execution_type=execution_type).set(avg_seconds)


def update_webhook_requests(status: str, count: int):
    """Update webhook requests total."""
    WEBHOOK_REQUESTS_TOTAL.labels(status=status).set(count)


def update_executions_by_org(organization_id: str, count: int):
    """Update executions by organization."""
    EXECUTIONS_BY_ORG_TOTAL.labels(organization_id=organization_id).set(count)


def update_agents_active(organization_id: str, count: int):
    """Update active agents count."""
    AGENTS_ACTIVE.labels(organization_id=organization_id).set(count)


def update_scheduled_jobs(status: str, count: int):
    """Update scheduled jobs total."""
    SCHEDULED_JOBS_TOTAL.labels(status=status).set(count)


def update_conversation_turns(execution_type: str, count: int):
    """Update conversation turns total."""
    CONVERSATION_TURNS_TOTAL.labels(execution_type=execution_type).set(count)
