"""
Prometheus metrics endpoint for Kubiya Control Plane.

This router exposes the /metrics endpoint for Prometheus scraping.
All gauge metrics are calculated on-the-fly from the database when
the endpoint is called, ensuring fresh data on every scrape.

Metrics collected:
- Active tasks by type and status
- Task failures by type
- Execution duration by type
- Worker queue depth
- LLM requests, latency, and token usage
- Tool execution metrics
- Streaming connections
- Webhook requests
- Organization-level metrics (agents, executions)
- Scheduled jobs status
"""

from datetime import datetime, timedelta
from fastapi import APIRouter, Depends
from fastapi.responses import Response
from sqlalchemy.orm import Session
from sqlalchemy import select, func, and_, case, extract
import structlog

from control_plane_api.app.database import get_db
from control_plane_api.app.observability.metrics import (
    get_metrics_response,
    # Original
    update_active_tasks,
    update_task_failures,
    # HIGH priority
    update_execution_duration,
    update_worker_queue_depth,
    update_llm_requests,
    update_llm_latency,
    update_llm_tokens,
    # MEDIUM priority
    update_streaming_connections,
    update_tool_execution_duration,
    update_tool_executions,
    update_execution_wait_time,
    update_webhook_requests,
    # LOWER priority
    update_executions_by_org,
    update_agents_active,
    update_scheduled_jobs,
    update_conversation_turns,
)

logger = structlog.get_logger(__name__)

router = APIRouter(tags=["monitoring"])

# Constants
ACTIVE_STATUSES = ['pending', 'queued', 'running', 'waiting_for_input', 'paused']
EXECUTION_TYPES = ['agent', 'team', 'workflow']
TERMINAL_STATUSES = ['completed', 'failed', 'cancelled']


# ==================== ORIGINAL METRICS ====================

def _collect_active_tasks(db: Session) -> None:
    """Query database for active task counts."""
    try:
        from control_plane_api.app.models.execution import Execution
        
        query = (
            select(
                Execution.execution_type,
                Execution.status,
                func.count(Execution.id).label('count')
            )
            .where(Execution.status.in_(ACTIVE_STATUSES))
            .group_by(Execution.execution_type, Execution.status)
        )
        
        result = db.execute(query)
        seen = set()
        
        for row in result.fetchall():
            execution_type = row.execution_type or 'unknown'
            status = row.status or 'unknown'
            update_active_tasks(execution_type, status, int(row.count))
            seen.add((execution_type, status))
        
        # Reset missing combinations to zero
        for et in EXECUTION_TYPES:
            for st in ACTIVE_STATUSES:
                if (et, st) not in seen:
                    update_active_tasks(et, st, 0)
                    
    except Exception as e:
        logger.error("active_tasks_collection_failed", error=str(e))


def _collect_task_failures(db: Session) -> None:
    """Query database for failed execution counts."""
    try:
        from control_plane_api.app.models.execution import Execution
        
        query = (
            select(
                Execution.execution_type,
                func.count(Execution.id).label('count')
            )
            .where(Execution.status == 'failed')
            .group_by(Execution.execution_type)
        )
        
        result = db.execute(query)
        seen = set()
        
        for row in result.fetchall():
            execution_type = row.execution_type or 'unknown'
            update_task_failures(execution_type, int(row.count))
            seen.add(execution_type)
        
        for et in EXECUTION_TYPES:
            if et not in seen:
                update_task_failures(et, 0)
                
    except Exception as e:
        logger.error("task_failures_collection_failed", error=str(e))


# ==================== HIGH PRIORITY METRICS ====================

def _collect_execution_duration(db: Session) -> None:
    """Calculate average execution duration for completed executions."""
    try:
        from control_plane_api.app.models.execution import Execution
        
        # Average duration for terminal state executions (last 24h)
        cutoff = datetime.utcnow() - timedelta(hours=24)
        
        query = (
            select(
                Execution.execution_type,
                Execution.status,
                func.avg(
                    extract('epoch', Execution.completed_at) - 
                    extract('epoch', Execution.started_at)
                ).label('avg_duration')
            )
            .where(
                and_(
                    Execution.status.in_(TERMINAL_STATUSES),
                    Execution.started_at.isnot(None),
                    Execution.completed_at.isnot(None),
                    Execution.completed_at >= cutoff,
                )
            )
            .group_by(Execution.execution_type, Execution.status)
        )
        
        result = db.execute(query)
        
        for row in result.fetchall():
            if row.avg_duration is not None:
                update_execution_duration(
                    row.execution_type or 'unknown',
                    row.status or 'unknown',
                    float(row.avg_duration)
                )
                
    except Exception as e:
        logger.error("execution_duration_collection_failed", error=str(e))


def _collect_worker_queue_depth(db: Session) -> None:
    """Count pending executions per worker queue."""
    try:
        from control_plane_api.app.models.execution import Execution
        
        query = (
            select(
                Execution.worker_queue_id,
                func.count(Execution.id).label('count')
            )
            .where(
                and_(
                    Execution.status.in_(['pending', 'queued']),
                    Execution.worker_queue_id.isnot(None),
                )
            )
            .group_by(Execution.worker_queue_id)
        )
        
        result = db.execute(query)
        
        for row in result.fetchall():
            queue_id = str(row.worker_queue_id) if row.worker_queue_id else 'default'
            update_worker_queue_depth(queue_id, int(row.count))
            
    except Exception as e:
        logger.error("worker_queue_depth_collection_failed", error=str(e))


def _collect_llm_metrics(db: Session) -> None:
    """Collect LLM request, latency, and token metrics."""
    try:
        from control_plane_api.app.models.analytics import ExecutionTurn
        
        # Last 24h for performance
        cutoff = datetime.utcnow() - timedelta(hours=24)
        
        # LLM requests by model and status (success vs error based on error_message)
        query = (
            select(
                ExecutionTurn.model,
                case(
                    (ExecutionTurn.error_message.isnot(None), 'error'),
                    else_='success'
                ).label('status'),
                func.count(ExecutionTurn.id).label('count'),
                func.avg(ExecutionTurn.duration_ms).label('avg_latency_ms'),
                func.sum(ExecutionTurn.input_tokens).label('input_tokens'),
                func.sum(ExecutionTurn.output_tokens).label('output_tokens'),
            )
            .where(ExecutionTurn.created_at >= cutoff)
            .group_by(ExecutionTurn.model, 'status')
        )
        
        result = db.execute(query)
        
        for row in result.fetchall():
            model = row.model or 'unknown'
            status = row.status
            
            update_llm_requests(model, status, int(row.count))
            
            if row.avg_latency_ms:
                update_llm_latency(model, float(row.avg_latency_ms) / 1000.0)
            
            if row.input_tokens:
                update_llm_tokens(model, 'input', int(row.input_tokens))
            if row.output_tokens:
                update_llm_tokens(model, 'output', int(row.output_tokens))
                
    except Exception as e:
        logger.error("llm_metrics_collection_failed", error=str(e))


# ==================== MEDIUM PRIORITY METRICS ====================

def _collect_tool_metrics(db: Session) -> None:
    """Collect tool execution metrics."""
    try:
        from control_plane_api.app.models.analytics import ExecutionToolCall
        
        cutoff = datetime.utcnow() - timedelta(hours=24)
        
        query = (
            select(
                ExecutionToolCall.tool_name,
                ExecutionToolCall.success,
                func.count(ExecutionToolCall.id).label('count'),
                func.avg(ExecutionToolCall.duration_ms).label('avg_duration_ms'),
            )
            .where(ExecutionToolCall.created_at >= cutoff)
            .group_by(ExecutionToolCall.tool_name, ExecutionToolCall.success)
        )
        
        result = db.execute(query)
        
        for row in result.fetchall():
            tool_name = row.tool_name or 'unknown'
            status = 'success' if row.success else 'failed'
            
            update_tool_executions(tool_name, status, int(row.count))
            
            if row.avg_duration_ms:
                update_tool_execution_duration(tool_name, float(row.avg_duration_ms) / 1000.0)
                
    except Exception as e:
        logger.error("tool_metrics_collection_failed", error=str(e))


def _collect_execution_wait_time(db: Session) -> None:
    """Calculate average wait time (created_at to started_at)."""
    try:
        from control_plane_api.app.models.execution import Execution
        
        cutoff = datetime.utcnow() - timedelta(hours=24)
        
        query = (
            select(
                Execution.execution_type,
                func.avg(
                    extract('epoch', Execution.started_at) - 
                    extract('epoch', Execution.created_at)
                ).label('avg_wait')
            )
            .where(
                and_(
                    Execution.started_at.isnot(None),
                    Execution.created_at.isnot(None),
                    Execution.started_at >= cutoff,
                )
            )
            .group_by(Execution.execution_type)
        )
        
        result = db.execute(query)
        
        for row in result.fetchall():
            if row.avg_wait is not None:
                avg_wait = float(row.avg_wait)
                if avg_wait >= 0:
                    update_execution_wait_time(
                        row.execution_type or 'unknown',
                        avg_wait
                    )
                
    except Exception as e:
        logger.error("execution_wait_time_collection_failed", error=str(e))


def _collect_webhook_metrics(db: Session) -> None:
    """Count webhook-triggered executions."""
    try:
        from control_plane_api.app.models.execution import Execution
        
        query = (
            select(
                Execution.status,
                func.count(Execution.id).label('count')
            )
            .where(Execution.trigger_source == 'job_webhook')
            .group_by(Execution.status)
        )
        
        result = db.execute(query)
        
        for row in result.fetchall():
            update_webhook_requests(row.status or 'unknown', int(row.count))
            
    except Exception as e:
        logger.error("webhook_metrics_collection_failed", error=str(e))


def _collect_streaming_connections(db: Session) -> None:
    """Count currently streaming executions (proxy for active SSE connections)."""
    try:
        from control_plane_api.app.models.execution import Execution
        
        # Running executions are likely streaming
        query = (
            select(func.count(Execution.id))
            .where(Execution.status == 'running')
        )
        
        result = db.execute(query)
        count = result.scalar() or 0
        update_streaming_connections(count)
        
    except Exception as e:
        logger.error("streaming_connections_collection_failed", error=str(e))


# ==================== LOWER PRIORITY METRICS ====================

def _collect_executions_by_org(db: Session) -> None:
    """Count total executions by organization (top 50)."""
    try:
        from control_plane_api.app.models.execution import Execution
        
        query = (
            select(
                Execution.organization_id,
                func.count(Execution.id).label('count')
            )
            .group_by(Execution.organization_id)
            .order_by(func.count(Execution.id).desc())
            .limit(50)  # Limit to prevent high cardinality
        )
        
        result = db.execute(query)
        
        for row in result.fetchall():
            org_id = row.organization_id or 'unknown'
            # Truncate org_id to reduce cardinality
            org_id_short = org_id[:36] if len(org_id) > 36 else org_id
            update_executions_by_org(org_id_short, int(row.count))
            
    except Exception as e:
        logger.error("executions_by_org_collection_failed", error=str(e))


def _collect_agents_active(db: Session) -> None:
    """Count active agents by organization (top 50)."""
    try:
        from control_plane_api.app.models.agent import Agent
        
        query = (
            select(
                Agent.organization_id,
                func.count(Agent.id).label('count')
            )
            .where(Agent.status == 'active')
            .group_by(Agent.organization_id)
            .order_by(func.count(Agent.id).desc())
            .limit(50)
        )
        
        result = db.execute(query)
        
        for row in result.fetchall():
            org_id = row.organization_id or 'unknown'
            org_id_short = org_id[:36] if len(org_id) > 36 else org_id
            update_agents_active(org_id_short, int(row.count))
            
    except Exception as e:
        logger.error("agents_active_collection_failed", error=str(e))


def _collect_scheduled_jobs(db: Session) -> None:
    """Count scheduled jobs by status."""
    try:
        from control_plane_api.app.models.job import Job
        
        query = (
            select(
                Job.status,
                func.count(Job.id).label('count')
            )
            .group_by(Job.status)
        )
        
        result = db.execute(query)
        
        for row in result.fetchall():
            update_scheduled_jobs(row.status or 'unknown', int(row.count))
            
    except Exception as e:
        logger.error("scheduled_jobs_collection_failed", error=str(e))


def _collect_conversation_turns(db: Session) -> None:
    """Count total conversation turns by execution type."""
    try:
        from control_plane_api.app.models.analytics import ExecutionTurn
        from control_plane_api.app.models.execution import Execution
        
        query = (
            select(
                Execution.execution_type,
                func.count(ExecutionTurn.id).label('count')
            )
            .join(Execution, ExecutionTurn.execution_id == Execution.id)
            .group_by(Execution.execution_type)
        )
        
        result = db.execute(query)
        
        for row in result.fetchall():
            update_conversation_turns(row.execution_type or 'unknown', int(row.count))
            
    except Exception as e:
        logger.error("conversation_turns_collection_failed", error=str(e))


# ==================== ENDPOINT ====================

@router.get("/metrics", include_in_schema=False)
async def prometheus_metrics(db: Session = Depends(get_db)) -> Response:
    """
    Prometheus metrics endpoint.
    
    Returns all application metrics in Prometheus text format.
    All gauge metrics are calculated on-the-fly from the database.
    """
    try:
        # Original metrics
        _collect_active_tasks(db)
        _collect_task_failures(db)
        
        # HIGH priority metrics
        _collect_execution_duration(db)
        _collect_worker_queue_depth(db)
        _collect_llm_metrics(db)
        
        # MEDIUM priority metrics
        _collect_tool_metrics(db)
        _collect_execution_wait_time(db)
        _collect_webhook_metrics(db)
        _collect_streaming_connections(db)
        
        # LOWER priority metrics
        _collect_executions_by_org(db)
        _collect_agents_active(db)
        _collect_scheduled_jobs(db)
        _collect_conversation_turns(db)
        
        # Generate response
        output, content_type = get_metrics_response()
        
        return Response(
            content=output,
            media_type=content_type,
            headers={
                "Cache-Control": "no-cache, no-store, must-revalidate",
                "Pragma": "no-cache",
                "Expires": "0",
            }
        )
        
    except Exception as e:
        logger.error("metrics_endpoint_failed", error=str(e))
        return Response(
            content=f"# Error: {str(e)}\n".encode('utf-8'),
            media_type="text/plain; charset=utf-8",
            status_code=500,
        )
