import json, logging, os, socket, sys, threading, time
from typing import Any, Dict, Optional

from tinysdk.log_event import LogEvent
from tinysdk.utils import suppress_raise


class TinyLogger:
  """Structured JSON logger for Kubernetes microservices.

  Emits newline-delimited JSON to stdout, formatted for log aggregators
  (Datadog, Loki, CloudWatch, etc.). Each log line is a self-contained JSON
  object that includes service metadata, log level, event type, message, and
  optional context/error fields — no multi-line logs, no plain text.

  Service metadata (service name, cluster, namespace, instance ID) is read
  from environment variables at construction time, so a single ``TinyLogger()``
  call picks up the pod's identity automatically in Kubernetes.

  Attributes:
    service_name (str): Value of SERVICE_NAME env var (default: "unknown-service")
    cluster (str): Value of CLUSTER env var (default: "unknown-cluster")
    namespace (str): Value of NAMESPACE env var (default: "default")
    instance_id (str): Value of INSTANCE_ID env var (default: hostname)
    heartbeat_interval (int): Seconds between heartbeat events (HEARTBEAT_INTERVAL, default: 30)
    log_level (str): Minimum log level (LOG_LEVEL, default: "INFO")

  Environment Variables:
    SERVICE_NAME: Name of the microservice
    CLUSTER: Kubernetes cluster name
    NAMESPACE: Kubernetes namespace
    INSTANCE_ID: Pod/instance identifier (falls back to hostname)
    HEARTBEAT_INTERVAL: Seconds between heartbeat log events (default: 30)
    LOG_LEVEL: Minimum log level to emit — DEBUG/INFO/WARNING/ERROR/CRITICAL (default: INFO)

  Example:
    >>> logger = TinyLogger()
    >>> logger.log("Processing started", level="INFO", context={"job_id": "abc123"})
    >>> logger.log("Something went wrong", level="ERROR", error=exc)
    >>> logger.recovery("Reconnected to RabbitMQ")
    >>>
    >>> # Scoped to a specific document for all subsequent calls
    >>> doc_logger = logger.for_document("doc-42")
    >>> doc_logger.log("OCR complete")
    >>>
    >>> # Periodic heartbeat (daemon thread, safe to leave running)
    >>> logger.start_heartbeat()
  """

  _instance_count = 0

  def __init__(self) -> None:
    TinyLogger._instance_count += 1
    self._instance_number = TinyLogger._instance_count
    self.service_name = os.getenv("SERVICE_NAME", "unknown-service")
    self.cluster = os.getenv("CLUSTER", "unknown-cluster")
    self.namespace = os.getenv("NAMESPACE", "default")
    self.instance_id = os.getenv("INSTANCE_ID", socket.gethostname())
    self.heartbeat_interval = int(os.getenv("HEARTBEAT_INTERVAL", "30"))
    self.log_level = os.getenv("LOG_LEVEL", "INFO").upper()
    self.start_time = time.time()
    self._heartbeat_seq = 0
    self._heartbeat_thread: Optional[threading.Thread] = None
    self._heartbeat_stop_event = threading.Event()
    self._logger = self._setup_logging()

  def _setup_logging(self) -> logging.Logger:
    handler = logging.StreamHandler(sys.stdout)
    handler.setLevel(self.log_level)
    handler.setFormatter(_TinyJsonFormatter())
    logger = logging.getLogger(f"tinysdk.logger.{self._instance_number}")
    logger.setLevel(self.log_level)
    logger.addHandler(handler)
    logger.propagate = False
    return logger

  def log(
    self, message: str, level: str = "INFO", context: Optional[Dict[str, Any]] = None, error: Optional[Exception] = None, document_id: str = ""
  ) -> None:
    """Emit a structured log event.

    Args:
      message (str): Human-readable log message (truncated to 500 chars).
      level (str): Log level — DEBUG, INFO, WARNING, ERROR, or CRITICAL. Defaults to INFO.
      context (Dict[str, Any], optional): Arbitrary key-value pairs included in the JSON output.
      error (Exception, optional): Exception to serialize into the ``error`` field.
      document_id (str): Optional document identifier for tracing a specific document through the pipeline.
    """
    self._emit(event_type="LOG", level=level, message=message, ctx=context or {}, document_id=document_id, error=error)

  def recovery(self, message: str, context: Optional[Dict[str, Any]] = None, document_id: str = "") -> None:
    """Emit a RECOVERY event indicating the service has recovered from an error state.

    Logged at INFO level with event_type "RECOVERY" so log aggregators can
    distinguish recoveries from regular info messages.

    Args:
      message (str): Description of what was recovered (e.g. "Reconnected to RabbitMQ").
      context (Dict[str, Any], optional): Arbitrary key-value pairs included in the JSON output.
      document_id (str): Optional document identifier.
    """
    self._emit(event_type="RECOVERY", level="INFO", message=message, ctx=context or {}, document_id=document_id)

  def start_heartbeat(self) -> None:
    """Start a background daemon thread that emits a HEARTBEAT log event at regular intervals.

    The heartbeat includes ``uptime_seconds`` and a monotonically increasing
    ``heartbeat_seq`` counter, making it easy to detect pod restarts or log gaps.
    Safe to call multiple times — only one thread will run at a time.
    Interval is controlled by the HEARTBEAT_INTERVAL environment variable (default: 30s).
    """
    if self._heartbeat_thread is not None and self._heartbeat_thread.is_alive():
      return
    self._heartbeat_stop_event.clear()
    self._heartbeat_thread = threading.Thread(target=self._heartbeat_loop, daemon=True)
    self._heartbeat_thread.start()

  def stop_heartbeat(self) -> None:
    """Stop the background heartbeat thread and wait for it to exit.

    Safe to call even if the heartbeat was never started.
    After stopping, :meth:`start_heartbeat` can be called again to restart.
    """
    self._heartbeat_stop_event.set()
    if self._heartbeat_thread is not None and self._heartbeat_thread.is_alive():
      self._heartbeat_thread.join(timeout=5)
    self._heartbeat_thread = None

  def for_document(self, document_id: str) -> "TinyDocumentLogger":
    """Return a document-scoped logger that automatically attaches a document_id to every call.

    Args:
      document_id (str): The document identifier to attach to all log events.

    Returns:
      TinyDocumentLogger: A thin wrapper around this logger with document_id pre-filled.

    Example:
      >>> doc_logger = logger.for_document("doc-42")
      >>> doc_logger.log("OCR complete")        # document_id="doc-42" added automatically
      >>> doc_logger.log("Extraction failed", level="ERROR", error=exc)
    """
    return TinyDocumentLogger(self, document_id)

  @suppress_raise()
  def _emit(self, event_type: str, level: str, message: str, ctx: Dict[str, Any], document_id: str = "", error: Optional[Exception] = None) -> None:
    level = level.upper()
    if level not in ("DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"):
      level = "INFO"
    error_str = self._format_error(error) if error else ""
    event = LogEvent(self.service_name, self.instance_id, self.cluster, self.namespace, level, event_type, message[:500], ctx, document_id, error_str)
    log_func = getattr(self._logger, level.lower(), self._logger.info)
    log_func(json.dumps(event.to_dict(), separators=(",", ":"), ensure_ascii=True))

  @suppress_raise(default="")
  def _format_error(self, error: Optional[Exception]) -> str:
    if error is None:
      return ""
    name = type(error).__name__
    return f"{name}: {error}" if str(error) else name

  @suppress_raise()
  def _heartbeat_loop(self) -> None:
    while not self._heartbeat_stop_event.is_set():
      self._heartbeat_stop_event.wait(self.heartbeat_interval)
      if self._heartbeat_stop_event.is_set():
        break
      self._heartbeat_seq += 1
      context = {"uptime_seconds": int(time.time() - self.start_time), "heartbeat_seq": self._heartbeat_seq}
      self._emit(event_type="HEARTBEAT", level="INFO", message="service alive", ctx=context)


class TinyDocumentLogger:
  """A document-scoped view of TinyLogger that pre-fills ``document_id`` on every call.

  Obtain an instance via :meth:`TinyLogger.for_document` rather than constructing directly.

  Example:
    >>> doc_logger = logger.for_document("doc-42")
    >>> doc_logger.log("Processing started")
    >>> doc_logger.recovery("Retried successfully")
  """

  def __init__(self, logger: TinyLogger, document_id: str) -> None:
    self._logger = logger
    self._document_id = document_id

  def log(self, message: str, level: str = "INFO", **kwargs: Any) -> None:
    """Emit a log event with this logger's document_id pre-filled. Accepts the same kwargs as TinyLogger.log."""
    self._logger.log(message, level=level, document_id=self._document_id, **kwargs)

  def recovery(self, message: str, **kwargs: Any) -> None:
    """Emit a RECOVERY event with this logger's document_id pre-filled. Accepts the same kwargs as TinyLogger.recovery."""
    self._logger.recovery(message, document_id=self._document_id, **kwargs)


class _TinyJsonFormatter(logging.Formatter):
  def format(self, record: logging.LogRecord) -> str:
    return record.getMessage()
