"""Metrics parsing and extraction."""

from __future__ import annotations

import re
from pathlib import Path

from abreka.errors import MetricParseError

# METRIC:train:loss=0.005  or  METRIC:loss=0.005 (defaults to val)
_METRIC_LINE_RE = re.compile(
    r"^METRIC:(?:([a-zA-Z_]+):)?([a-zA-Z_][a-zA-Z0-9_]*)=([0-9eE.+-]+)$"
)

# CLI spec: [split:]key=value
_METRIC_SPEC_RE = re.compile(
    r"^(?:([a-zA-Z_]+):)?([a-zA-Z_][a-zA-Z0-9_]*)=([0-9eE.+-]+)$"
)

# Reference: val:accuracy (no value)
_METRIC_REF_RE = re.compile(
    r"^(?:([a-zA-Z_]+):)?([a-zA-Z_][a-zA-Z0-9_]*)$"
)


def parse_metric_line(line: str) -> tuple[str, str, float] | None:
    """Parse a METRIC: line from stdout.

    Returns (split, key, value) or None if the line doesn't match.
    """
    line = line.strip()
    m = _METRIC_LINE_RE.match(line)
    if not m:
        return None
    split = m.group(1) or "val"
    key = m.group(2)
    value = float(m.group(3))
    return split, key, value


def parse_metric_spec(spec: str) -> tuple[str, str, float]:
    """Parse a CLI metric specification like 'val:accuracy=0.9912'.

    Returns (split, key, value). Split defaults to 'val'.
    """
    m = _METRIC_SPEC_RE.match(spec.strip())
    if not m:
        raise MetricParseError(f"Invalid metric spec: {spec!r}. Expected [split:]key=value")
    split = m.group(1) or "val"
    key = m.group(2)
    value = float(m.group(3))
    return split, key, value


def parse_metric_ref(ref: str) -> tuple[str, str]:
    """Parse a metric reference like 'val:accuracy' (no value).

    Returns (split, key). Split defaults to 'val'.
    """
    m = _METRIC_REF_RE.match(ref.strip())
    if not m:
        raise MetricParseError(f"Invalid metric reference: {ref!r}. Expected [split:]key")
    split = m.group(1) or "val"
    key = m.group(2)
    return split, key


def extract_metrics_from_output(output: str) -> dict[str, dict[str, float]]:
    """Scan multiline output for all METRIC: lines.

    Returns {split: {key: value}}. Later lines override earlier ones for the same split:key.
    """
    metrics: dict[str, dict[str, float]] = {}
    for line in output.splitlines():
        result = parse_metric_line(line)
        if result:
            split, key, value = result
            if split not in metrics:
                metrics[split] = {}
            metrics[split][key] = value
    return metrics


def extract_metrics_from_file(path: Path) -> dict[str, dict[str, float]]:
    """Scan a text file for METRIC: lines without loading the full file into memory."""
    metrics: dict[str, dict[str, float]] = {}
    with open(path, "r", errors="replace") as f:
        for line in f:
            result = parse_metric_line(line)
            if result:
                split, key, value = result
                if split not in metrics:
                    metrics[split] = {}
                metrics[split][key] = value
    return metrics
