from __future__ import annotations

import json
import re

from hytop.gpu.metrics import JSON_KEY_BY_METRIC
from hytop.gpu.models import Sample

ANSI_RE = re.compile(r"\x1B\[[0-?]*[ -/]*[@-~]")
CARD_KEY_RE = re.compile(r"^card(\d+)$")


def strip_ansi(text: str) -> str:
    """Strip ANSI escape sequences from text.

    Args:
        text: Input text possibly containing ANSI escapes.

    Returns:
        Text with ANSI escape sequences removed.
    """

    return ANSI_RE.sub("", text)


def parse_number(text: str) -> float:
    """Extract the first numeric token from text.

    Args:
        text: Input text that contains at least one number.

    Returns:
        Parsed float value.

    Raises:
        ValueError: If no number can be found.
    """

    match = re.search(r"-?\d+(?:\.\d+)?", text)
    if not match:
        raise ValueError(f"cannot parse number from {text!r}")
    return float(match.group(0))


def parse_hy_smi_output(raw: str, sample_ts: float) -> dict[int, Sample]:
    """Parse hy-smi JSON output into GPU keyed samples.

    Args:
        raw: Raw hy-smi stdout text.
        sample_ts: Monotonic timestamp assigned to parsed rows.

    Returns:
        Mapping from GPU id to parsed sample.
    """

    cleaned = strip_ansi(raw).strip()
    if not cleaned:
        return {}
    try:
        payload = json.loads(cleaned)
    except json.JSONDecodeError:
        return {}
    if not isinstance(payload, dict):
        return {}

    result: dict[int, Sample] = {}
    for card_key, card_data in payload.items():
        if not isinstance(card_key, str):
            continue
        card_match = CARD_KEY_RE.match(card_key)
        if card_match is None or not isinstance(card_data, dict):
            continue
        gpu_id = int(card_match.group(1))
        sample = Sample(ts=sample_ts)
        for metric_name, json_key in JSON_KEY_BY_METRIC.items():
            raw_value = card_data.get(json_key)
            if raw_value is None:
                continue
            try:
                parsed_value = parse_number(str(raw_value))
            except ValueError:
                continue
            setattr(sample, metric_name, parsed_value)
        result[gpu_id] = sample
    return result
