from __future__ import annotations

from typing import Any, Dict, Iterable, Optional, Tuple

PROCESS_RSS_BYTES = "hl_process_rss_bytes"
PROCESS_CPU_UTILIZATION = "hl_process_cpu_utilization"
PROCESS_CPU_SECONDS_TOTAL = "hl_process_cpu_seconds_total"
PROCESS_THREAD_COUNT = "hl_process_thread_count"
PROCESS_RUNTIME_GC_COLLECTIONS = "hl_process_runtime_gc_collections"
PROCESS_RUNTIME_GC_COLLECTED = "hl_process_runtime_gc_collected"
PROCESS_RUNTIME_GC_UNCOLLECTABLE = "hl_process_runtime_gc_uncollectable"

PIPELINE_MEMORY_DELTA_BYTES = "hl_pipeline_memory_delta_bytes"
PIPELINE_TIME_US = "hl_pipeline_time_us"
INFERENCE_TIME_US = "hl_inference_time_us"
ELEMENT_TIME_US = "hl_element_time_us"

FILES_QUEUED = "hl_files_queued"
FILES_IN_FLIGHT = "hl_files_in_flight"
DATAFILE_SAVE_LOCAL_US = "hl_datafile_save_local_us"
DATAFILE_BYTES = "hl_datafile_bytes"
DATAFILE_SAMPLES = "hl_datafile_samples"
DATAFILE_UPLOAD_US = "hl_datafile_upload_us"
DATAFILE_UPLOAD_BYTES = "hl_datafile_upload_bytes"
DATAFILE_UPLOAD_MB_PER_S = "hl_datafile_upload_mb_per_s"

RUNTIME_SHUTDOWN_US = "hl_runtime_shutdown_us"

_FRAME_METRIC_ALIASES: Dict[str, Tuple[str, ...]] = {
    "pipeline_memory_delta_bytes": (
        "pipeline_memory_delta_bytes",
        "pipeline_memory_delta",
        "pipeline_memory",
    ),
    "pipeline_time_us": ("pipeline_time_us", "pipeline_time_ms", "pipeline_time_ns", "pipeline_time"),
    "inference_time_us": ("inference_time_us", "inference_time_ms", "inference_time_ns", "inference_time"),
    "process_rss_bytes": ("process_rss_bytes", "process_rss"),
}


def _lookup_metric(metrics: Dict[str, Any], keys: Iterable[str]) -> Tuple[Optional[Any], Optional[str]]:
    for key in keys:
        if key in metrics:
            return metrics[key], key
    return None, None


def duration_us_from_value(value: Any, key: str) -> Optional[float]:
    try:
        value_float = float(value)
    except (TypeError, ValueError):
        return None
    if key.endswith("_ms"):
        return value_float * 1_000
    if key.endswith("_us"):
        return value_float
    if key.endswith("_ns"):
        return value_float / 1_000
    return value_float * 1_000_000


def normalize_frame_metrics(metrics: Optional[Dict[str, Any]]) -> Dict[str, Any]:
    if not metrics:
        return {}
    normalized: Dict[str, Any] = {}
    for canonical, keys in _FRAME_METRIC_ALIASES.items():
        value, key = _lookup_metric(metrics, keys)
        if value is None or key is None:
            continue
        if canonical.endswith("_time_us"):
            duration_us = duration_us_from_value(value, key)
            if duration_us is not None:
                normalized[canonical] = duration_us
        else:
            normalized[canonical] = value
    return normalized
