from __future__ import annotations

import threading
import uuid
from collections.abc import Mapping, Sequence
from typing import Any

from log_utils import (
    LOG_EVENT_TYPE,
    build_log_payload,
    normalize_event_type,
    normalize_index,
    normalize_json_payload,
    normalize_status,
    severity_for_log_identifier,
    utc_now,
)
from models import EventRecord, RunContext
from storage import LocalSessionStore

_EVENT_SEVERITIES = {"debug", "info", "warning", "error", "critical"}
_DEFAULT_SEVERITY_BY_TYPE = {
    "run_started": "info",
    "run_resumed": "info",
    "run_ended": "info",
    "checkpoint_saved": "info",
    "scheduler_step": "info",
    "learning_rate_changed": "info",
    "gradient_clipping_applied": "warning",
    "overflow_skipped_step": "warning",
    "anomaly_detected": "warning",
    "probe_started": "info",
    "probe_completed": "info",
    "user_log": "info",
    "runtime_error": "error",
    "runtime_interruption": "warning",
}


class EventLogger:
    """Buffered structured event 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[EventRecord] = []
        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,
        event_type: str,
        *,
        severity: str | None = None,
        step: int | None = None,
        epoch: int | None = None,
        status: str | None = None,
        payload: Mapping[str, Any] | None = None,
        timestamp: str | None = None,
    ) -> EventRecord:
        """Queue a structured event for batched persistence."""

        self._ensure_open()
        normalized_event_type = self._normalize_event_type(event_type)
        normalized_step = self._normalize_index("step", step)
        normalized_epoch = self._normalize_index("epoch", epoch)
        normalized_status = self._normalize_status(status)
        normalized_payload = self._normalize_payload(payload)

        record = EventRecord(
            id=self._new_id(),
            run_id=self.run_id,
            sequence_id=self._store.allocate_event_sequence(self.run_id),
            event_type=normalized_event_type,
            severity=self._normalize_severity(
                severity,
                event_type=normalized_event_type,
                status=normalized_status,
            ),
            step=normalized_step,
            epoch=normalized_epoch,
            timestamp=timestamp or utc_now(),
            status=normalized_status,
            payload=normalized_payload,
        )

        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 add_log(
        self,
        message: str,
        *,
        identifier: str | None = "Info",
        step: int | None = None,
        epoch: int | None = None,
        status: str | None = None,
        payload: Mapping[str, Any] | None = None,
        timestamp: str | None = None,
        event_type: str = LOG_EVENT_TYPE,
    ) -> EventRecord:
        """Queue a structured run log with an Info, Warning, or Error identifier."""

        normalized_identifier, log_payload = build_log_payload(
            message,
            identifier=identifier,
            payload=payload,
        )
        return self.log(
            event_type,
            severity=severity_for_log_identifier(normalized_identifier),
            step=step,
            epoch=epoch,
            status=status,
            payload=log_payload,
            timestamp=timestamp,
        )

    def flush(self) -> tuple[EventRecord, ...]:
        """Persist any buffered events 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_events(self.run_id, batch)
        except Exception:
            with self._lock:
                self._buffer = list(batch) + self._buffer
            raise

    def read(
        self,
        *,
        event_type: str | None = None,
        event_types: Sequence[str] | None = None,
        severity: 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[EventRecord, ...]:
        """Flush pending events and return ordered persisted records for the run."""

        if not self._closed:
            self.flush()
        return self._store.read_events(
            self.run_id,
            event_type=event_type,
            event_types=event_types,
            severity=severity,
            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) -> EventLogger:
        return self

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

    def _ensure_open(self) -> None:
        if self._closed:
            raise RuntimeError("EventLogger 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("EventLogger expects a TrainingSession or RunContext created by the Quark SDK.")

    @staticmethod
    def _normalize_event_type(event_type: str) -> str:
        return normalize_event_type(event_type)

    @staticmethod
    def _normalize_severity(
        severity: str | None,
        *,
        event_type: str,
        status: str | None,
    ) -> str:
        if severity is None:
            if event_type == "run_ended" and status == "failed":
                return "error"
            if event_type == "run_ended" and status == "interrupted":
                return "warning"
            return _DEFAULT_SEVERITY_BY_TYPE.get(event_type, "info")

        candidate = severity.strip().lower()
        if candidate not in _EVENT_SEVERITIES:
            raise ValueError(
                f"severity '{severity}' is not supported. Expected one of {sorted(_EVENT_SEVERITIES)}."
            )
        return candidate

    @staticmethod
    def _normalize_status(status: str | None) -> str | None:
        return normalize_status(status)

    @staticmethod
    def _normalize_index(label: str, value: int | None) -> int | None:
        return normalize_index(label, value)

    @staticmethod
    def _normalize_payload(payload: Mapping[str, Any] | None) -> dict[str, Any]:
        return normalize_json_payload(payload)

    @staticmethod
    def _new_id() -> str:
        return uuid.uuid4().hex
