"""Logging configuration and context management."""

import contextvars
import json
import logging
import sys
from collections.abc import Iterable
from datetime import UTC, datetime
from typing import Any, ClassVar

# =============================================================================
# LOGGING CONTEXT VARS
# =============================================================================

_request_id_var: contextvars.ContextVar[str | None] = contextvars.ContextVar(
    "albus_request_id",
    default=None,
)
_agent_id_var: contextvars.ContextVar[str | None] = contextvars.ContextVar(
    "albus_agent_id",
    default=None,
)
_thread_id_var: contextvars.ContextVar[str | None] = contextvars.ContextVar(
    "albus_thread_id",
    default=None,
)
_pathway_id_var: contextvars.ContextVar[str | None] = contextvars.ContextVar(
    "albus_pathway_id",
    default=None,
)
_execution_id_var: contextvars.ContextVar[str | None] = contextvars.ContextVar(
    "albus_execution_id",
    default=None,
)
_trace_depth_var: contextvars.ContextVar[int] = contextvars.ContextVar(
    "albus_trace_depth",
    default=0,
)


def set_log_request_id(value: str | None) -> contextvars.Token:
    return _request_id_var.set(value)


def reset_log_request_id(token: contextvars.Token) -> None:
    _request_id_var.reset(token)


def set_log_agent_id(value: str | None) -> contextvars.Token:
    return _agent_id_var.set(value)


def reset_log_agent_id(token: contextvars.Token) -> None:
    _agent_id_var.reset(token)


def set_log_thread_id(value: str | None) -> contextvars.Token:
    return _thread_id_var.set(value)


def reset_log_thread_id(token: contextvars.Token) -> None:
    _thread_id_var.reset(token)


def set_log_pathway_id(value: str | None) -> contextvars.Token:
    return _pathway_id_var.set(value)


def reset_log_pathway_id(token: contextvars.Token) -> None:
    _pathway_id_var.reset(token)


def set_log_execution_id(value: str | None) -> contextvars.Token:
    return _execution_id_var.set(value)


def reset_log_execution_id(token: contextvars.Token) -> None:
    _execution_id_var.reset(token)


def push_trace_depth() -> contextvars.Token:
    """Increase trace depth (for nested indentation in logs)."""
    return _trace_depth_var.set(_trace_depth_var.get() + 1)


def pop_trace_depth(token: contextvars.Token) -> None:
    """Restore trace depth."""
    _trace_depth_var.reset(token)


# =============================================================================
# FORMATTERS
# =============================================================================

# ANSI codes for pretty logs
_RESET = "\033[0m"
_DIM = "\033[2m"
_BOLD = "\033[1m"
_RED = "\033[31m"
_YELLOW = "\033[33m"
_GREEN = "\033[32m"
_CYAN = "\033[36m"
_MAGENTA = "\033[35m"

_LEVEL_COLORS = {
    logging.DEBUG: _MAGENTA,
    logging.INFO: _GREEN,
    logging.WARNING: _YELLOW,
    logging.ERROR: _RED,
    logging.CRITICAL: _BOLD + _RED,
}

# Tree characters for trace view
_TREE_PIPE = "│"
_TREE_SPACE = "  "


class _ContextFilter(logging.Filter):
    """Injects contextvars into every LogRecord."""

    def filter(self, record: logging.LogRecord) -> bool:
        record.request_id = _request_id_var.get()
        record.agent_id = _agent_id_var.get()
        record.thread_id = _thread_id_var.get()
        record.pathway_id = _pathway_id_var.get()
        record.execution_id = _execution_id_var.get()
        record.trace_depth = _trace_depth_var.get()
        return True


class _JsonLogFormatter(logging.Formatter):
    """JSON formatter for production log ingestion (Datadog/Loki/Splunk)."""

    LEVEL_MAP: ClassVar[dict[int, str]] = {
        logging.DEBUG: "debug",
        logging.INFO: "info",
        logging.WARNING: "warn",
        logging.ERROR: "error",
        logging.CRITICAL: "critical",
    }

    def format(self, record: logging.LogRecord) -> str:
        ts = datetime.now(UTC).isoformat()
        obj: dict[str, Any] = {
            "ts": ts,
            "level": self.LEVEL_MAP.get(record.levelno, "info"),
            "logger": record.name,
            "msg": record.getMessage(),
        }
        # Add context if present
        for attr in ("request_id", "agent_id", "thread_id", "pathway_id", "execution_id"):
            val = getattr(record, attr, None)
            if val:
                obj[attr] = val
        # Add exception if present
        if record.exc_info:
            obj["exc"] = self.formatException(record.exc_info)
        return json.dumps(obj, default=str)


class _PrettyLogFormatter(logging.Formatter):
    """Human-readable trace-tree formatter for dev.

    Features:
    - Indentation based on trace depth (nested context)
    - Context tags (req/agent/thread/pathway/execution)
    - Color-coded levels
    - Tree-style visual hierarchy
    """

    LEVEL_NAMES: ClassVar[dict[int, str]] = {
        logging.DEBUG: "DEBUG",
        logging.INFO: "INFO ",
        logging.WARNING: "WARN ",
        logging.ERROR: "ERROR",
        logging.CRITICAL: "CRIT ",
    }

    LEVEL_ICONS: ClassVar[dict[int, str]] = {
        logging.DEBUG: "·",
        logging.INFO: "→",
        logging.WARNING: "⚠",
        logging.ERROR: "✗",
        logging.CRITICAL: "‼",
    }

    def __init__(self, use_color: bool = False, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self._use_color = use_color

    def _short_name(self, name: str) -> str:
        """Use last segment of logger name."""
        if not name:
            return "root"
        return name.rsplit(".", maxsplit=1)[-1] if "." in name else name

    def format(self, record: logging.LogRecord) -> str:
        levelno = record.levelno
        level_label = self.LEVEL_NAMES.get(levelno, record.levelname)
        short_name = self._short_name(record.name)
        timestamp = self.formatTime(record, datefmt="%H:%M:%S")
        icon = self.LEVEL_ICONS.get(levelno, "→")
        depth = getattr(record, "trace_depth", 0) or 0

        # Context tags (compact form)
        ctx_parts: list[str] = []
        req = getattr(record, "request_id", None)
        if req:
            ctx_parts.append(f"req={req}")
        agent = getattr(record, "agent_id", None)
        if agent:
            ctx_parts.append(f"a={agent}")
        thread = getattr(record, "thread_id", None)
        if thread:
            ctx_parts.append(f"t={thread}")
        pathway = getattr(record, "pathway_id", None)
        if pathway:
            ctx_parts.append(f"p={pathway}")
        exe = getattr(record, "execution_id", None)
        if exe:
            ctx_parts.append(f"x={exe}")

        ctx = ""
        if ctx_parts:
            ctx = f" [{' '.join(ctx_parts)}]"

        # Indentation for tree view
        indent = _TREE_PIPE + _TREE_SPACE if depth > 0 else ""
        indent = indent * depth

        if self._use_color:
            color = _LEVEL_COLORS.get(levelno, _RESET)
            timestamp = f"{_DIM}{timestamp}{_RESET}"
            level_label = f"{color}{level_label}{_RESET}"
            short_name = f"{_DIM}{short_name}{_RESET}"
            indent_colored = f"{_DIM}{indent}{_RESET}" if indent else ""
            icon_colored = f"{color}{icon}{_RESET}"
            ctx_colored = f"{_DIM}{ctx}{_RESET}" if ctx else ""
            line = f"{timestamp}  {level_label}  {short_name}{ctx_colored}  {indent_colored}{icon_colored} {record.getMessage()}"
        else:
            line = f"{timestamp}  {level_label}  {short_name}{ctx}  {indent}{icon} {record.getMessage()}"

        if record.exc_info:
            line = f"{line}\n{self.formatException(record.exc_info)}"
        return line


# =============================================================================
# CONFIGURATION
# =============================================================================


def _parse_level(level: str) -> int:
    lvl = (level or "").strip().upper()
    return getattr(logging, lvl, logging.INFO) if lvl else logging.INFO


def _set_levels(names: Iterable[str], level: int) -> None:
    for n in names:
        try:
            logging.getLogger(n).setLevel(level)
        except Exception:
            logging.getLogger(__name__).debug("Failed to set log level for %s", n, exc_info=True)


def configure_logging(*, log_level: str | None = None) -> None:
    """Configure Python logging for the runtime.

    Supports two output formats (via ALBUS_LOG_FORMAT env var):
    - "pretty" (default for TTY): Human-readable trace-tree format with colors
    - "json": Structured JSON lines for production log aggregation

    Environment variables:
    - ALBUS_LOG_LEVEL: Log level (DEBUG, INFO, WARN, ERROR)
    - ALBUS_LOG_FORMAT: "pretty" or "json" (default: pretty if TTY, json otherwise)
    - ALBUS_NO_COLOR: Disable colors in pretty mode
    - ALBUS_HTTP_TRACE: Enable HTTP client debug logs
    - ALBUS_ACCESS_LOG: Enable aiohttp access logs
    """
    from core.config import get_settings

    settings = get_settings()

    env_level = settings.albus_log_level
    root_level = _parse_level(env_level or (log_level or ""))
    if env_level is None and not (log_level or "").strip():
        root_level = logging.INFO

    # Determine format: explicit env var, or auto-detect from TTY
    is_tty = getattr(sys.stderr, "isatty", lambda: False)()
    log_format = settings.albus_log_format.strip().lower()
    if log_format not in ("pretty", "json"):
        log_format = "pretty" if is_tty else "json"

    use_color = (
        log_format == "pretty"
        and is_tty
        and settings.albus_no_color.lower() not in ("1", "true", "yes")
    )

    root = logging.getLogger()
    root.setLevel(root_level)
    for h in root.handlers[:]:
        root.removeHandler(h)

    handler = logging.StreamHandler(sys.stderr)
    handler.addFilter(_ContextFilter())

    if log_format == "json":
        handler.setFormatter(_JsonLogFormatter())
    else:
        handler.setFormatter(_PrettyLogFormatter(use_color=use_color))

    root.addHandler(handler)

    http_trace = settings.albus_http_trace.lower() in ("1", "true")
    if not http_trace:
        _set_levels(["openai", "httpx", "httpcore"], logging.WARNING)

    access_log = settings.albus_access_log.lower() in ("1", "true")
    if not access_log:
        _set_levels(["aiohttp.access"], logging.WARNING)


__all__ = [
    # Configuration
    "configure_logging",
    "pop_trace_depth",
    "push_trace_depth",
    "reset_log_agent_id",
    "reset_log_execution_id",
    "reset_log_pathway_id",
    "reset_log_request_id",
    "reset_log_thread_id",
    "set_log_agent_id",
    "set_log_execution_id",
    "set_log_pathway_id",
    # Context setters
    "set_log_request_id",
    "set_log_thread_id",
]
