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:
  _instance_count = 0

  def __init__(self):
    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 = 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: Dict[str, Any] = None, error: Optional[Exception] = None, document_id: str = "") -> None:
    self._emit(event_type="LOG", level=level, message=message, ctx=context or {}, document_id=document_id, error=error)

  def recovery(self, message: str, context: Dict[str, Any] = None, document_id: str = "") -> None:
    self._emit(event_type="RECOVERY", level="INFO", message=message, ctx=context or {}, document_id=document_id)

  def start_heartbeat(self) -> None:
    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 for_document(self, document_id: str) -> "TinyDocumentLogger":
    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:
  def __init__(self, logger: TinyLogger, document_id: str):
    self._logger = logger
    self._document_id = document_id

  def log(self, message: str, level: str = "INFO", **kwargs) -> None:
    self._logger.log(message, level=level, document_id=self._document_id, **kwargs)

  def recovery(self, message: str, **kwargs) -> None:
    self._logger.recovery(message, document_id=self._document_id, **kwargs)


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