#!/usr/bin/env python3
"""
Latency exporter - mesure complete latency Aurelie input→output
Exporte vers Prometheus sur http://localhost:9090/metrics

Usage dans le workflow:
  from exporters.latency_exporter import measure_endpoint_latency
  result = measure_endpoint_latency('ollama_chat', ollama.generate, prompt)
  
Ou integration directe avec EdgeTTS:
  latency_ms = measure_endpoint_latency('tts', edgeTTS.synthesize, tokens)
"""

import time
import json
import logging
import os
from prometheus_client import start_http_server, Histogram, Summary, Counter, Gauge

logger = logging.getLogger('aurelie_latency_exporter')
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')

# Prometheus metrics pour la latence
RESPONSE_LATENCY_MS = Histogram(
    'aurelie_response_latency_ms',
    'Latence complete (input->TTS fini) en ms',
    ['endpoint'],  # ollama, tts, pgvector
    namespace='aurelie'
)

TTS_QUALITY_DB = Gauge(
    'aurelie_tts_quality_dbfs',
    'Peak dBFS de ma voix EdgeTTS',
    ['status'],  # generated, degraded, error
    namespace='aurelie'
)

LARGEUR_SCORE_TOTAL = Counter(
    'aurelie_latency_total_count',
    'Nombre total de réponses générées (latency tracking)',
    namespace='aurelie'
)


def measure_endpoint_latency(endpoint: str, callback_fn, *args, **kwargs):
    """Wrapper pour mesurer latence d'un endpoint Aurelie."""
    t0 = time.monotonic()
    try:
        result = callback_fn(*args, **kwargs)
        latency_ms = (time.monotonic() - t0) * 1000
        
        RESPONSE_LATENCY_MS.labels(endpoint=endpoint).observe(latency_ms)
        LARGEUR_SCORE_TOTAL.inc()
        
        logger.info(f"Endpoint: {endpoint}, Latence: {latency_ms:.2f}ms")
        return result
        
    except Exception as e:
        logger.error(f"Erreur endpoint {endpoint}: {e}")
        raise


def measure_tts_quality(audio_path: str) -> dict:
    """Calcule RMS, peak dBFS, SNR du signal audio généré par EdgeTTS."""
    try:
        import numpy as np
        rate, data = None, None
        
        # Support WAV (scipy) ou raw array numpy directement
        if hasattr(data, 'ndim'):  # c'est déjà un numpy array
            data_float = data.astype(float)
        else:
            from scipy.io import wavfile
            rate, data = wavfile.read(audio_path)
            data_float = data.astype(float)
        
        rms_db = 20 * np.log10(np.sqrt(np.mean(data_float**2)) / 32768 + 1e-9)
        peak_dbfs = 20 * np.log10(max(np.abs(data_float)) / 32768 + 1e-9)
        
        TTS_QUALITY_DB.labels(status='generated').set(peak_dbfs)
        
        return {
            'rms_db': rms_db,
            'peak_dbfs': peak_dbfs,
            'duration_samples': len(data),
            'sample_rate': rate
        }
    except Exception as e:
        logger.error(f"Erreur mesure qualité TTS: {e}")
        return {}


def compute_health_score():
    """Calcule le health score composite de la chaine Aurelie (0-100)."""
    # Récupérer les scores des composants (exemple)
    ollama_status = get_component_score('ollama')     # HTTP /api/tags
    pgvector_status = get_component_score('pgvector')  # SELECT 1
    tts_status = get_component_score('edgetts')       # HTTP health check
    
    gpu_utilization = get_gpu_utilization()           # nvidia-smi GPU utilization
    cpu_load = get_cpu_load_avg()                     # os.getloadavg() load 1 min
    
    # Score compose (poids definis plus tot par Christian)
    weights = {
        'ollama':     25,  # % weight for LLM API status
        'pgvector':   20,  # % weight for PostgreSQL/pqvector
        'edgetts':    15,  # % weight for Text-to-Speech health
        'gpu_health': 20,  # % weight for GPU availability + temp
        'cpu_load':   20,  # % weight for CPU load (lower is better)
    }
    
    score = sum(
        component_score * weights[metric] / 100.0
        for metric, component_score in [
            ('ollama', ollama_status),
            ('pgvector', pgvector_status),
            ('edgetts', tts_status),
            ('gpu_health', gpu_health),
            ('cpu_load', cpu_load),
        ]
    )
    
    # Log the score to Prometheus via Gauge
    HEALTH_SCORE.set(score)
    if score < 85:
        logger.warning(f"Health Score LOW: {score:.1f}/100 — attention nécessaire")
    else:
        logger.info(f"Health Score OK: {score:.1f}/100")
    
    return score


# Prometheus metric pour le composite health score
HEALTH_SCORE = Gauge(
    'aurelie_health_score',
    'Sante composite de la chaine Aurelie (0-100)',
    namespace='aurelie'
)


if __name__ == '__main__':
    start_http_server(8100)  # Prometheus scrape cet exporter sur port 8100/metrics
    
    logger.info('Latency-exporter started on :8100')
    
    while True:
        time.sleep(60)  # Chaque minute, log un point
        
        try:
            cpu_load = os.getloadavg()[0]  # load 1 min
            health_score = compute_health_score()
        except Exception as e:
            logger.error(f"Erreur mesure CPU: {e}")