from __future__ import annotations

import atexit
import contextlib
import signal
import sys
import threading
import time
import traceback
from typing import Any, Dict, Generator, List, Optional

from pulselog.config import Config
from pulselog.checkpoint import CheckpointStore
from pulselog.exceptions import InvalidStatusError
from pulselog.levels import Level
from pulselog.queue import ShardedLogQueue, QUEUE_MAX_SIZE
from pulselog.record import LogRecord
from pulselog.worker import BackgroundWorker


class Logger:
    """Non-blocking real-time logger with an embedded browser dashboard.

    Every logging call is O(1) and never blocks the calling thread. A daemon
    background worker drains the queue and pushes batches to the WebSocket
    server every 10ms (configurable).

    Performance design:
        - ShardedLogQueue: each thread writes to its own private deque —
          zero cross-thread contention on the hot path.
        - _level_value cached as a flat attr — eliminates double lookup of
          self._config.level_value on every call.
        - Overflow mode resolved to a bound method at __init__ — eliminates
          two string comparisons from every _log() call.
        - LogRecord.timestamp=None by default — worker stamps the batch once
          at drain time instead of calling time.time() per record.
        - info_fast() / error_fast() etc. — zero-allocation paths that skip
          **kwargs entirely for calls with no metadata.
    """

    def __init__(
        self,
        name: str = "pulselog",
        host: Optional[str] = None,
        port: Optional[int] = None,
        auto_open: Optional[bool] = None,
        dashboard: Optional[bool] = None,
        checkpoint_path: Optional[str] = None,
        level: Optional[str] = None,
        worker_interval: Optional[float] = None,
        queue_size: Optional[int] = None,
        overflow: str = "drop",          # 'drop' | 'block' | 'raise'
    ) -> None:
        self._start_time = time.time()
        self._checkpoints_saved = 0

        # --- Config ---
        self._config = Config.load(
            name=name,
            host=host,
            port=port,
            auto_open=auto_open,
            dashboard=dashboard,
            checkpoint_path=checkpoint_path,
            level=level,
            worker_interval=worker_interval,
        )

        # Cache level threshold as a flat attribute.
        # self._config.level_value is two attr lookups per call; this is one.
        self._level_value: int = self._config.level_value

        # Per-thread tag — threading.local so threads don't clobber each other.
        self._local = threading.local()

        # --- Queue ---
        effective_queue_size = queue_size or QUEUE_MAX_SIZE
        self._queue = ShardedLogQueue(maxsize=effective_queue_size)

        # Resolve overflow strategy to a bound method ONCE at init.
        # Eliminates two string comparisons from every _log() call.
        if overflow == "block":
            self._put = self._put_block
        elif overflow == "raise":
            self._put = self._put_raise
        else:
            self._put = self._queue.put   # default: drop oldest

        # --- Checkpoint store ---
        self._checkpoint_store = CheckpointStore(self._config.checkpoint_path)

        # --- Dashboard server ---
        self._server = None
        handlers: List[Any] = []
        if self._config.dashboard:
            from pulselog.server.websocket import DashboardServer
            self._server = DashboardServer(self._config)
            self._server.start()
            handlers.append(self._server.broadcast)

        # --- Background worker ---
        self._worker = BackgroundWorker(
            queue=self._queue,
            handlers=handlers,
            interval=self._config.worker_interval,
        )
        self._worker.start()

        # --- Graceful shutdown ---
        self._shutdown_called = False
        atexit.register(self._atexit_handler)
        try:
            signal.signal(signal.SIGTERM, lambda *_: self._atexit_handler())
        except (OSError, ValueError):
            pass

    # ------------------------------------------------------------------
    # Per-thread tag
    # ------------------------------------------------------------------

    @property
    def _current_tag(self) -> Optional[str]:
        return getattr(self._local, "tag", None)

    @_current_tag.setter
    def _current_tag(self, value: Optional[str]) -> None:
        self._local.tag = value

    # ------------------------------------------------------------------
    # Overflow strategy helpers — bound at __init__, never called directly
    # ------------------------------------------------------------------

    def _put_block(self, record: LogRecord) -> None:
        """Block caller until space is available — overflow='block'."""
        while len(self._queue) >= self._queue._maxsize:
            self._queue._event.wait(timeout=0.01)
        self._queue.put(record)

    def _put_raise(self, record: LogRecord) -> None:
        """Raise immediately if queue is full — overflow='raise'."""
        if len(self._queue) >= self._queue._maxsize:
            raise RuntimeError(
                f"[pulselog] Queue full ({self._queue._maxsize:,} records). "
                "Use overflow='drop' or increase queue_size."
            )
        self._queue.put(record)

    # ------------------------------------------------------------------
    # Core hot path
    # ------------------------------------------------------------------

    def _log(self, level: Level, msg: str, extra: Optional[Dict[str, Any]]) -> None:
        """Internal dispatch — O(1), never blocks.

        extra is a plain Optional[dict] (not **kwargs) so callers control
        whether a dict is allocated. Public methods pass (extra if extra else None)
        so no-kwarg calls allocate nothing beyond the LogRecord itself.
        """
        if level.value < self._level_value:
            return
        self._put(LogRecord(
            level=level,
            msg=msg,
            timestamp=None,          # worker stamps the batch at drain time
            extra=extra,             # None when no kwargs — no empty dict alloc
            tag=self._current_tag,
        ))

    # ------------------------------------------------------------------
    # Standard API — **kwargs surface (fully backwards compatible)
    # ------------------------------------------------------------------

    def debug(self, msg: str, **extra: Any) -> None:
        """Log a DEBUG message with optional metadata."""
        self._log(Level.DEBUG, msg, extra if extra else None)

    def info(self, msg: str, **extra: Any) -> None:
        """Log an INFO message with optional metadata."""
        self._log(Level.INFO, msg, extra if extra else None)

    def warning(self, msg: str, **extra: Any) -> None:
        """Log a WARNING message with optional metadata."""
        self._log(Level.WARNING, msg, extra if extra else None)

    def error(self, msg: str, **extra: Any) -> None:
        """Log an ERROR message with optional metadata."""
        self._log(Level.ERROR, msg, extra if extra else None)

    def critical(self, msg: str, **extra: Any) -> None:
        """Log a CRITICAL message with optional metadata."""
        self._log(Level.CRITICAL, msg, extra if extra else None)

    def exception(self, msg: str, exc_info: bool = True, **extra: Any) -> None:
        """Log an ERROR with the current exception traceback attached.

        Call inside an except block to capture the full traceback:

            try:
                risky()
            except ValueError:
                log.exception("risky failed", context="training")
        """
        if exc_info:
            tb = traceback.format_exc()
            if tb and tb.strip() != "NoneType: None":
                extra["traceback"] = tb.strip()
        self._log(Level.ERROR, msg, extra if extra else None)

    # ------------------------------------------------------------------
    # Zero-allocation fast paths — no **kwargs dict constructed
    # ------------------------------------------------------------------
    #
    # When you call logger.info("msg", loss=0.4), Python builds the dict
    # {"loss": 0.4} in C BEFORE entering info(). There is no way to avoid
    # this with **kwargs syntax — it happens before our code runs.
    #
    # These methods accept only a string. No dict is constructed anywhere
    # in the call chain. Use them in tight loops where metadata is not
    # needed per-record:
    #
    #   for step in range(100_000):
    #       logger.info_fast("step")        # ~400-500k/sec
    #       # vs:
    #       logger.info("step", n=step)     # ~161k/sec (**kwargs tax)
    #
    # For calls that DO need metadata, keep using the standard API.
    # ------------------------------------------------------------------

    def debug_fast(self, msg: str) -> None:
        """Zero-allocation DEBUG — no kwargs, no dict created."""
        if Level.DEBUG.value < self._level_value:
            return
        self._put(LogRecord(
            level=Level.DEBUG,
            msg=msg,
            timestamp=None,
            tag=self._current_tag,
        ))

    def info_fast(self, msg: str) -> None:
        """Zero-allocation INFO — no kwargs, no dict created.

        Use in tight loops where no per-record metadata is needed:

            for step in range(100_000):
                logger.info_fast("step done")    # ~400-500k/sec
        """
        if Level.INFO.value < self._level_value:
            return
        self._put(LogRecord(
            level=Level.INFO,
            msg=msg,
            timestamp=None,
            tag=self._current_tag,
        ))

    def warning_fast(self, msg: str) -> None:
        """Zero-allocation WARNING — no kwargs, no dict created."""
        if Level.WARNING.value < self._level_value:
            return
        self._put(LogRecord(
            level=Level.WARNING,
            msg=msg,
            timestamp=None,
            tag=self._current_tag,
        ))

    def error_fast(self, msg: str) -> None:
        """Zero-allocation ERROR — no kwargs, no dict created."""
        if Level.ERROR.value < self._level_value:
            return
        self._put(LogRecord(
            level=Level.ERROR,
            msg=msg,
            timestamp=None,
            tag=self._current_tag,
        ))

    def critical_fast(self, msg: str) -> None:
        """Zero-allocation CRITICAL — no kwargs, no dict created."""
        if Level.CRITICAL.value < self._level_value:
            return
        self._put(LogRecord(
            level=Level.CRITICAL,
            msg=msg,
            timestamp=None,
            tag=self._current_tag,
        ))

    # ------------------------------------------------------------------
    # Checkpoints
    # ------------------------------------------------------------------

    def save(
        self,
        name: str,
        data: Dict[str, Any],
        status: str = "DONE",
        note: str = "",
        progress: Optional[float] = None,
    ) -> Dict[str, Any]:
        """Persist a checkpoint and emit it as a log event.

        Args:
            name:     Unique checkpoint identifier.
            data:     JSON-serialisable dict to store.
            status:   One of DONE | IN_PROGRESS | FAILED | SKIPPED.
            note:     Human-readable description.
            progress: Optional 0–100 progress value.

        Returns:
            Full record dict including _meta.

        Raises:
            InvalidStatusError: If status is not an allowed value.
        """
        result = self._checkpoint_store.save(
            name=name, data=data, status=status, note=note, progress=progress
        )
        self._checkpoints_saved += 1
        self._log(
            Level.INFO,
            f"[checkpoint] {name} → {status.upper()}"
            + (f" ({progress:.0f}%)" if progress is not None else ""),
            {"_checkpoint": {"name": name, "status": status.upper(),
                             "note": note, "progress": progress, **data}},
        )
        return result

    def load(self, name: str) -> Optional[Dict[str, Any]]:
        """Load a checkpoint by name. Returns None if not found."""
        return self._checkpoint_store.load(name)

    def checkpoints(self) -> List[str]:
        """Return all checkpoint names, most-recently saved first."""
        return self._checkpoint_store.checkpoints()

    def delete_checkpoint(self, name: str) -> None:
        """Delete a checkpoint. Silent if it doesn't exist."""
        self._checkpoint_store.delete_checkpoint(name)

    # ------------------------------------------------------------------
    # Utilities
    # ------------------------------------------------------------------

    def tag(self, label: str) -> None:
        """Group subsequent log entries under label (current thread only)."""
        self._current_tag = label
        self._queue.put(LogRecord(
            level=Level.INFO,
            msg=label,
            tag=label,
            extra={"_divider": True},
        ))

    @contextlib.contextmanager
    def context(self, tag: str) -> Generator[None, None, None]:
        """Scope a tag to a block of code. Restores the previous tag on exit.

        Safe for multi-threaded use — each thread has its own tag via
        threading.local.

            with log.context("training"):
                for epoch in range(10):
                    log.info("epoch done", loss=0.4)
            # tag restored here
        """
        previous = self._current_tag
        self.tag(tag)
        try:
            yield
        finally:
            self._current_tag = previous

    def divider(self, label: str = "") -> None:
        """Insert a visual divider in the dashboard log stream."""
        self._queue.put(LogRecord(
            level=Level.INFO,
            msg=label or "─" * 40,
            tag=self._current_tag,
            extra={"_divider": True},
        ))

    def flush(self, timeout: float = 2.0) -> bool:
        """Force-drain the queue synchronously.

        Args:
            timeout: Max seconds to wait. Default 2.0s.

        Returns:
            True if fully drained, False if timeout was hit.
        """
        self._queue._event.set()
        deadline = time.monotonic() + timeout
        while not self._queue.is_empty() and time.monotonic() < deadline:
            time.sleep(0.005)
        drained = self._queue.is_empty()
        if not drained:
            sys.stderr.write(
                f"[pulselog] flush() timed out after {timeout}s — "
                f"{len(self._queue):,} records still in queue.\n"
            )
        return drained

    def shutdown(self) -> None:
        """Gracefully shut down the worker and dashboard server."""
        if self._shutdown_called:
            return
        self._shutdown_called = True
        self.flush()
        self._worker.shutdown()
        if self._server:
            self._server.shutdown()

    def stats(self) -> Dict[str, Any]:
        """Return operational metrics for monitoring."""
        dropped  = self._queue.dropped_count()
        q_len    = len(self._queue)
        return {
            "records_dropped":   dropped,
            "queue_size":        q_len,
            "queue_capacity":    self._queue._maxsize,
            "queue_fill_pct":    round(q_len / (self._queue._maxsize or 1) * 100, 1),
            "checkpoints_saved": self._checkpoints_saved,
            "dashboard_clients": self._server.client_count if self._server else 0,
            "uptime_seconds":    round(time.time() - self._start_time, 2),
        }

    # ------------------------------------------------------------------
    # Private
    # ------------------------------------------------------------------

    def _atexit_handler(self) -> None:
        try:
            self.shutdown()
        except Exception:
            pass