"""Structured logging for NSCT — JSONLines formatter, module loggers, env-config.

Environment variables
---------------------
NSCT_LOG_LEVEL   — Log level string (DEBUG, INFO, WARNING, ERROR, CRITICAL).
                   Default: INFO.
NSCT_LOG_FILE    — If set, write JSON logs to this file instead of stdout.
                   When unset, logs go to stdout.

Usage
-----
    from nsct.logging_config import get_logger, setup_logging

    setup_logging()                       # call once at app bootstrap
    logger = get_logger(__name__)         # any module
    logger.info("started", research_run_id="abc")
"""

from __future__ import annotations

import json
import logging
import os
import sys
import time
import uuid
from contextvars import ContextVar
from typing import Any

from nsct import __version__

# ---------------------------------------------------------------------------
# Per-request / per-run context (ContextVar for async safety)
# ---------------------------------------------------------------------------

# Legacy dict-based context (kept for backward compat with existing callers)
_request_ctx: dict[str, str] = {}
_request_ctx_token: ContextVar[dict[str, Any]] = ContextVar("_request_ctx", default={})


def set_request_ctx(**kwargs: str) -> None:
    """Set (merge) key-value pairs into the per-request context dict."""
    _request_ctx.update(kwargs)
    ctx = _request_ctx_token.get()
    ctx.update(kwargs)
    _request_ctx_token.set(ctx)


def clear_request_ctx() -> None:
    """Clear all per-request context entries."""
    _request_ctx.clear()
    _request_ctx_token.set({})


# ---------------------------------------------------------------------------
# Helpers: research_run_id convenience
# ---------------------------------------------------------------------------

_research_run_id: ContextVar[str | None] = ContextVar("_research_run_id", default=None)


def set_research_run_id(run_id: str) -> None:
    _research_run_id.set(run_id)


def get_research_run_id() -> str | None:
    return _research_run_id.get()


def clear_research_run_id() -> None:
    _research_run_id.set(None)


def get_stage_ctx() -> str | None:
    """Read the 'stage' value from the active context dict."""
    return _request_ctx.get("stage")


# ---------------------------------------------------------------------------
# JSON Formatter — JSONLines with research_run_id, stage, etc.
# ---------------------------------------------------------------------------


class JSONFormatter(logging.Formatter):
    """Emit log records as one JSON object per line (JSONLines)."""

    default_fmt: str = ""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self._start_ts: float = time.monotonic()

    def format(self, record: logging.LogRecord) -> str:
        ts = time.strftime("%Y-%m-%dT%H:%M:%S", time.gmtime(record.created))
        elapsed_ms = round((time.monotonic() - self._start_ts) * 1000, 2)

        # Core fields
        data: dict[str, Any] = {
            "timestamp": ts,
            "level": record.levelname,
            "module": record.name,
            "message": record.getMessage(),
            "elapsed_ms": elapsed_ms,
            "pid": record.process,
            "thread": record.thread,
            "function": record.funcName,
            "line": record.lineno,
            "version": __version__,
        }

        # Pull in the ContextVar-based request context
        ctx = _request_ctx_token.get()
        if ctx:
            for k, v in ctx.items():
                data[k] = v

        # Legacy _request_ctx (for callers who only use set_request_ctx)
        for k, v in _request_ctx.items():
            if k not in data:
                data[k] = v

        # research_run_id and stage convenience
        rid = get_research_run_id()
        if rid and "research_run_id" not in data:
            data["research_run_id"] = rid

        # exception info
        if record.exc_info and record.exc_info[0] is not None:
            data["exc_info"] = self.formatException(record.exc_info)

        return json.dumps(data, default=str, ensure_ascii=False)


# ---------------------------------------------------------------------------
# Bootstrap
# ---------------------------------------------------------------------------

_MODULE_LOGGERS: dict[str, logging.Logger] = {}


def _make_module_logger(name: str) -> logging.Logger:
    """Return (and cache) a pre-configured child logger."""
    if name in _MODULE_LOGGERS:
        return _MODULE_LOGGERS[name]
    # Only create loggers for known top-level modules
    known_prefixes = ("nsct.api", "nsct.orchestration", "nsct.stages", "nsct.cli")
    is_known = any(name == p or name.startswith(p + ".") for p in known_prefixes)
    logger = logging.getLogger(name)
    if is_known:
        _MODULE_LOGGERS[name] = logger
    return logger


# Map short names → full module names
_MODULE_ALIASES = {
    "api": "nsct.api",
    "rest_research": "nsct.api.rest_research",
    "orchestration": "nsct.orchestration",
    "stages": "nsct.stages",
    "cli": "nsct.cli",
    "metrics": "nsct.metrics",
    "logging_config": "nsct.logging_config",
}


def get_logger(name: str | None = None) -> logging.Logger:
    """Return a pre-configured child logger.

    Short names (``api``, ``orchestration``, ``stages``, ``cli``) are
    expanded to their full module path so they hit the known-prefix gate.
    ``None`` → ``nsct`` root.
    """
    if name is None or name in _MODULE_ALIASES:
        name = _MODULE_ALIASES.get(name, "nsct") if name else "nsct"
    return _make_module_logger(name)


def setup_logging(
    level: str | None = None,
    json_format: bool = True,
) -> None:
    """Configure logging.

    Reads *NSCT_LOG_LEVEL* and *NSCT_LOG_FILE* from the environment so
    callers don't have to pass arguments.

    Parameters
    ----------
    level : str | None
        Log level string. Defaults to ``NSCT_LOG_LEVEL`` env, then ``INFO``.
    json_format : bool
        If True, use JSONFormatter; else use a human-readable format.
    """
    import os

    effective_level = level or os.environ.get("NSCT_LOG_LEVEL", "INFO")
    log_file = os.environ.get("NSCT_LOG_FILE")

    log_level = getattr(logging, effective_level.upper(), logging.INFO)

    root = logging.getLogger()
    root.setLevel(log_level)
    root.handlers.clear()

    if log_file:
        # File handler — JSONLines
        fh = logging.FileHandler(log_file, encoding="utf-8")
        fh.setLevel(log_level)
        if json_format:
            fh.setFormatter(JSONFormatter())
        else:
            fh.setFormatter(
                logging.Formatter(
                    "%(asctime)s [%(levelname)-8s] %(name)s:%(funcName)s:%(lineno)d — %(message)s",
                    datefmt="%Y-%m-%d %H:%M:%S",
                )
            )
        root.addHandler(fh)
    else:
        # Console handler — JSONLines
        ch = logging.StreamHandler(sys.stdout)
        ch.setLevel(log_level)
        if json_format:
            ch.setFormatter(JSONFormatter())
        else:
            ch.setFormatter(
                logging.Formatter(
                    "%(asctime)s [%(levelname)-8s] %(name)s:%(funcName)s:%(lineno)d — %(message)s",
                    datefmt="%Y-%m-%d %H:%M:%S",
                )
            )
        root.addHandler(ch)


# Quick-export for convenience
logger = get_logger()