from __future__ import annotations

import json
import math
from numbers import Integral, Real
import re
import threading
import uuid
from collections.abc import Mapping
from typing import Any

from models import MetricRecord, RunContext
from storage import LocalSessionStore, utc_now

_METRIC_SEGMENT_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]*$")


class MetricsLogger:
    """Buffered scalar metric writer and reader for a single run."""

    def __init__(
        self,
        session: Any,
        *,
        batch_size: int = 100,
    ) -> None:
        if batch_size <= 0:
            raise ValueError("batch_size must be a positive integer.")

        self._context, self._store = self._resolve_session_dependencies(session)
        self._session = session if callable(getattr(session, "_register_flushable", None)) else None
        self._batch_size = batch_size
        self._buffer: list[MetricRecord] = []
        self._lock = threading.RLock()
        self._closed = False
        if self._session is not None:
            self._session._register_flushable(self)

    @property
    def run_id(self) -> str:
        return self._context.run_id

    @property
    def buffered_count(self) -> int:
        with self._lock:
            return len(self._buffer)

    def log(
        self,
        metric: str,
        value: float,
        *,
        step: int | None = None,
        epoch: int | None = None,
        namespace: str | None = None,
        description: str | None = None,
        metadata: Mapping[str, Any] | None = None,
        timestamp: str | None = None,
    ) -> MetricRecord:
        """Queue a scalar metric for batched persistence."""

        self._ensure_open()
        normalized_step = self._normalize_index("step", step)
        normalized_epoch = self._normalize_index("epoch", epoch)
        if normalized_step is None and normalized_epoch is None:
            raise ValueError("Metrics must include at least one of step or epoch.")

        normalized_namespace, metric_name, metric_key = self._normalize_metric_identity(
            metric,
            namespace=namespace,
        )
        normalized_value = self._normalize_value(value)
        normalized_description = self._normalize_description(description)
        normalized_metadata = self._normalize_metadata(metadata)
        record = MetricRecord(
            id=uuid.uuid4().hex,
            run_id=self.run_id,
            sequence_id=self._allocate_sequence_id(),
            namespace=normalized_namespace,
            metric_name=metric_name,
            metric_key=metric_key,
            value=normalized_value,
            step=normalized_step,
            epoch=normalized_epoch,
            timestamp=timestamp or utc_now(),
            description=normalized_description,
            metadata=normalized_metadata,
        )

        should_flush = False
        with self._lock:
            self._buffer.append(record)
            should_flush = len(self._buffer) >= self._batch_size

        if should_flush:
            self.flush()

        return record

    def flush(self) -> tuple[MetricRecord, ...]:
        """Persist any buffered metrics as a single storage transaction."""

        self._ensure_open()
        with self._lock:
            if not self._buffer:
                return ()
            batch = tuple(self._buffer)
            self._buffer.clear()

        try:
            return self._store.write_scalar_metrics(self.run_id, batch)
        except Exception:
            with self._lock:
                self._buffer = list(batch) + self._buffer
            raise

    def read(
        self,
        *,
        metric_key: str | None = None,
        namespace: str | None = None,
        metric_name: str | None = None,
        step: int | None = None,
        step_min: int | None = None,
        step_max: int | None = None,
        epoch: int | None = None,
        epoch_min: int | None = None,
        epoch_max: int | None = None,
        limit: int | None = None,
    ) -> tuple[MetricRecord, ...]:
        """Flush pending metrics and return ordered persisted records for the run."""

        if not self._closed:
            self.flush()
        return self._store.read_scalar_metrics(
            self.run_id,
            metric_key=metric_key,
            namespace=namespace,
            metric_name=metric_name,
            step=step,
            step_min=step_min,
            step_max=step_max,
            epoch=epoch,
            epoch_min=epoch_min,
            epoch_max=epoch_max,
            limit=limit,
        )

    def close(self) -> None:
        if self._closed:
            return
        try:
            self.flush()
        finally:
            self._closed = True
            if self._session is not None:
                self._session._unregister_flushable(self)

    def __enter__(self) -> MetricsLogger:
        return self

    def __exit__(self, exc_type, exc, exc_tb) -> None:
        self.close()

    def _allocate_sequence_id(self) -> int:
        return self._store.allocate_metric_sequence(self.run_id)

    def _ensure_open(self) -> None:
        if self._closed:
            raise RuntimeError("MetricsLogger is closed.")

    @staticmethod
    def _resolve_session_dependencies(session: Any) -> tuple[RunContext, LocalSessionStore]:
        if isinstance(session, RunContext):
            return session, LocalSessionStore(session.storage_dir)

        context = getattr(session, "context", None)
        store = getattr(session, "_store", None)
        if isinstance(context, RunContext) and isinstance(store, LocalSessionStore):
            return context, store

        raise TypeError(
            "MetricsLogger expects a TrainingSession or RunContext created by the Quark SDK."
        )

    @staticmethod
    def _normalize_metric_identity(metric: str, *, namespace: str | None) -> tuple[str, str, str]:
        candidate = metric.strip()
        if not candidate:
            raise ValueError("Metric names must not be empty.")

        if namespace is not None:
            normalized_namespace = MetricsLogger._normalize_namespace(namespace)
            if "/" in candidate:
                raise ValueError(
                    "Metric names must not include '/' when an explicit namespace is provided."
                )
            normalized_name = MetricsLogger._normalize_segment(candidate, "Metric name")
            return (
                normalized_namespace,
                normalized_name,
                f"{normalized_namespace}/{normalized_name}",
            )

        parts = candidate.split("/")
        if len(parts) == 1:
            normalized_name = MetricsLogger._normalize_segment(parts[0], "Metric name")
            return "", normalized_name, normalized_name

        normalized_namespace = "/".join(
            MetricsLogger._normalize_segment(part, "Namespace segment") for part in parts[:-1]
        )
        normalized_name = MetricsLogger._normalize_segment(parts[-1], "Metric name")
        return (
            normalized_namespace,
            normalized_name,
            f"{normalized_namespace}/{normalized_name}",
        )

    @staticmethod
    def _normalize_namespace(namespace: str) -> str:
        parts = [segment.strip() for segment in namespace.split("/")]
        if not parts or any(not part for part in parts):
            raise ValueError("Metric namespaces must not be empty.")
        return "/".join(
            MetricsLogger._normalize_segment(part, "Namespace segment") for part in parts
        )

    @staticmethod
    def _normalize_segment(value: str, label: str) -> str:
        candidate = value.strip()
        if not candidate:
            raise ValueError(f"{label} must not be empty.")
        if _METRIC_SEGMENT_PATTERN.fullmatch(candidate) is None:
            raise ValueError(
                f"{label} '{candidate}' may only contain letters, digits, '.', '_' or '-'."
            )
        return candidate

    @staticmethod
    def _normalize_index(label: str, value: int | None) -> int | None:
        if value is None:
            return None
        if isinstance(value, bool) or not isinstance(value, Integral):
            raise TypeError(f"{label} must be an integer when provided.")
        normalized_value = int(value)
        if normalized_value < 0:
            raise ValueError(f"{label} must not be negative.")
        return normalized_value

    @staticmethod
    def _normalize_value(value: float) -> float:
        if isinstance(value, bool) or not isinstance(value, Real):
            raise TypeError("Metric values must be numeric scalars.")
        normalized_value = float(value)
        if not math.isfinite(normalized_value):
            raise ValueError("Metric values must be finite real numbers.")
        return normalized_value

    @staticmethod
    def _normalize_description(description: str | None) -> str | None:
        if description is None:
            return None
        normalized = description.strip()
        return normalized or None

    @staticmethod
    def _normalize_metadata(metadata: Mapping[str, Any] | None) -> dict[str, Any] | None:
        if metadata is None:
            return None
        payload = dict(metadata)
        json.dumps(payload, sort_keys=True)
        return payload
