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 LogQueue, 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 50 ms (configurable).

    Usage::

        logger = Logger("my-app")
        logger.info("hello", user_id=42)
        logger.save("step1", {"acc": 0.9}, status="DONE", progress=100)
        logger.shutdown()

    Multiple Logger instances with different names may coexist in the same
    process — they do *not* share a queue, worker, or checkpoint store.
    """

    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,           # FIX #3: expose queue size
        overflow: str = "drop",                      # FIX #9: 'drop' | 'block' | 'raise'
    ) -> None:
        self._start_time = time.time()
        self._records_logged = 0
        self._checkpoints_saved = 0
        self._overflow = overflow

        # --- 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,
        )

        # FIX #18: use threading.local() so each thread has its own tag.
        # Previously self._current_tag was a plain instance attr — a race
        # condition where Thread A's log.tag() silently affected Thread B.
        self._local = threading.local()

        # --- Core pipeline ---
        effective_queue_size = queue_size or QUEUE_MAX_SIZE
        self._queue = LogQueue(maxsize=effective_queue_size)

        # --- 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 hooks ---
        self._shutdown_called = False  # guard against double-shutdown
        atexit.register(self._atexit_handler)
        try:
            signal.signal(signal.SIGTERM, lambda *_: self._atexit_handler())
        except (OSError, ValueError):
            pass

    # ------------------------------------------------------------------
    # Internal tag property (per-thread)
    # ------------------------------------------------------------------

    @property
    def _current_tag(self) -> Optional[str]:
        """FIX #18: per-thread tag — safe for multi-threaded use."""
        return getattr(self._local, "tag", None)

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

    # ------------------------------------------------------------------
    # Core logging methods
    # ------------------------------------------------------------------

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

        Level filtering happens BEFORE LogRecord construction to avoid
        unnecessary allocations.

        Args:
            level:   Severity Level enum value.
            msg:     Log message string.
            **extra: Arbitrary key/value metadata.
        """
        if level.value < self._config.level_value:
            return

        record = LogRecord(
            level=level,
            msg=str(msg),
            timestamp=time.time(),
            extra=extra,
            tag=self._current_tag,
        )

        # FIX #9: overflow modes
        if self._overflow == "block":
            # Block caller until there is space in the queue
            while len(self._queue) >= self._queue._deque.maxlen:
                self._queue._event.wait(timeout=0.01)
            self._queue.put(record)
        elif self._overflow == "raise":
            if len(self._queue) >= self._queue._deque.maxlen:
                raise RuntimeError(
                    f"[pulselog] Queue full ({self._queue._deque.maxlen:,} records). "
                    "Use overflow='drop' or increase queue_size."
                )
            self._queue.put(record)
        else:
            # default: drop oldest (original behaviour, now with warning)
            self._queue.put(record)

        self._records_logged += 1

    def debug(self, msg: str, **extra: Any) -> None:
        """Log a DEBUG-level message."""
        self._log(Level.DEBUG, msg, **extra)

    def info(self, msg: str, **extra: Any) -> None:
        """Log an INFO-level message."""
        self._log(Level.INFO, msg, **extra)

    def warning(self, msg: str, **extra: Any) -> None:
        """Log a WARNING-level message."""
        self._log(Level.WARNING, msg, **extra)

    def error(self, msg: str, **extra: Any) -> None:
        """Log an ERROR-level message."""
        self._log(Level.ERROR, msg, **extra)

    def critical(self, msg: str, **extra: Any) -> None:
        """Log a CRITICAL-level message."""
        self._log(Level.CRITICAL, msg, **extra)

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

        Call inside an ``except`` block to capture the full traceback
        automatically.

        Usage::

            try:
                result = risky_call()
            except ValueError as e:
                log.exception("risky_call failed", context="training")

        Args:
            msg:      Error message.
            exc_info: If True (default), captures the current exception.
            **extra:  Additional key/value context.
        """
        if exc_info:
            tb = traceback.format_exc()
            if tb and tb.strip() != "NoneType: None":
                extra["traceback"] = tb.strip()
        self._log(Level.ERROR, msg, **extra)

    # ------------------------------------------------------------------
    # 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 identifier for this checkpoint.
            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 one of the allowed values.
        """
        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:
            Dict with data and ``_meta``, or None if not found (never raises).
        """
        return self._checkpoint_store.load(name)

    def checkpoints(self) -> List[str]:
        """Return all checkpoint names ordered by 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).

        FIX #18: now uses threading.local() — setting a tag on Thread A
        no longer affects Thread B.

        Args:
            label: Short descriptive label shown in the dashboard.
        """
        self._current_tag = label
        divider_record = LogRecord(
            level=Level.INFO,
            msg=label,
            tag=label,
            extra={"_divider": True},
        )
        self._queue.put(divider_record)

    @contextlib.contextmanager
    def context(self, tag: str) -> Generator[None, None, None]:
        """FIX #19: context manager that scopes a tag to a block of code.

        Restores the previous tag (or None) on exit — even if an exception
        is raised.  Safe for multi-threaded use (each thread has its own tag).

        Usage::

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

    def divider(self, label: str = "") -> None:
        """Insert a visual divider line in the dashboard log stream.

        Args:
            label: Optional text label on the divider.
        """
        record = LogRecord(
            level=Level.INFO,
            msg=label or "─" * 40,
            tag=self._current_tag,
            extra={"_divider": True},
        )
        self._queue.put(record)

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

        FIX #13: timeout was hardcoded at 0.5s — too short for large queues.
        Now configurable and returns False if the queue wasn't fully drained
        within the timeout so callers know if logs were potentially lost.

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

        Returns:
            True if queue was 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.

        FIX #2: _atexit_handler previously called flush() then shutdown(),
        but shutdown() also calls flush() — two drain cycles on every exit.
        Now shutdown() is the single authoritative shutdown path.

        After this call the Logger instance should not be used.
        """
        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 / alerting.

        Returns:
            Dict with counters and current queue state.
        """
        dropped = self._queue.dropped_count()
        logged  = self._records_logged
        drop_rate = (dropped / logged) if logged > 0 else 0.0

        return {
            "records_logged":    logged,
            "records_dropped":   dropped,
            "drop_rate":         round(drop_rate, 4),      # e.g. 0.04 = 4%
            "queue_size":        len(self._queue),
            "queue_capacity":    self._queue._deque.maxlen,
            "queue_fill_pct":    round(
                len(self._queue) / (self._queue._deque.maxlen 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:
        """FIX #2: call shutdown() only — no longer calls flush() separately.

        Previously: flush() → shutdown() → flush() (double drain).
        Now:        shutdown() → flush() (single drain, correct).
        """
        try:
            self.shutdown()
        except Exception:
            pass  # must never raise in atexit