"""
Centralized logging configuration for TASC-Stack services.

This module configures loguru to output JSON-formatted logs, which ensures that
multi-line log entries (like tracebacks) are captured as single log entries by
the Docker Loki logging driver.

Usage:
    Call `configure_logging()` early in your application startup:

    ```python
    from common.log_config import configure_logging
    configure_logging(service_name="actx-server")
    ```

    For services with a config that inherits from CommonSettings, you can pass
    the settings object to automatically extract the service name and environment:

    ```python
    from common.log_config import configure_logging
    from my_service.config import settings
    configure_logging(settings=settings)
    ```

Exception Handling:
    This module also installs global exception hooks to ensure ALL Python exceptions
    (including those that occur before logging is configured, or in asyncio tasks)
    are formatted as single-line JSON for proper Loki ingestion.

Warning Handling:
    Python warnings (e.g., DeprecationWarning, FutureWarning, UserWarning from
    third-party libraries) are also captured and formatted as single-line JSON.
    This ensures warnings from libraries like google-api-core or pydantic appear
    as single log entries in Loki rather than multi-line text.

    For very early exception catching (e.g., import errors), you can call:

    ```python
    from common.log_config import install_exception_hooks
    install_exception_hooks()
    ```

    This is automatically called by `configure_logging()`, but can be called earlier
    if needed.
"""

import sys
import json
import logging
import traceback
import warnings
from datetime import datetime, timezone
from typing import TYPE_CHECKING

from loguru import logger


class InterceptHandler(logging.Handler):
    """
    Intercept standard library logging and route to loguru.

    This ensures logs from uvicorn, sqlalchemy, and other libraries
    using standard logging are formatted as JSON.
    """

    def emit(self, record: logging.LogRecord) -> None:
        # Get corresponding Loguru level if it exists
        try:
            level = logger.level(record.levelname).name
        except ValueError:
            level = record.levelno

        # Find caller from where the logged message originated
        frame, depth = logging.currentframe(), 2
        while frame and frame.f_code.co_filename == logging.__file__:
            frame = frame.f_back
            depth += 1

        logger.opt(depth=depth, exception=record.exc_info).log(
            level, record.getMessage()
        )


if TYPE_CHECKING:
    from common.config import CommonSettings


# Track whether exception hooks have been installed
_exception_hooks_installed = False

# Store service context for exception hooks
_service_context: dict = {}


def _format_exception_as_json(
    exc_type: type,
    exc_value: BaseException,
    exc_tb,
    level: str = "ERROR",
    message: str | None = None,
    context: dict | None = None,
) -> str:
    """
    Format an exception as a single-line JSON string.

    This is used by exception hooks to ensure all Python errors
    (including those occurring before loguru is configured) are
    formatted consistently for Loki ingestion.
    """
    # Format the traceback as a single string
    tb_lines = traceback.format_exception(exc_type, exc_value, exc_tb)
    tb_str = "".join(tb_lines)

    log_entry = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "level": level,
        "message": message or str(exc_value),
        "logger": "exception_hook",
        "function": None,
        "line": None,
        "file": None,
        "exception": {
            "type": exc_type.__name__ if exc_type else None,
            "value": str(exc_value) if exc_value else None,
            "traceback": tb_str,
        },
    }

    # Add service context if available
    if _service_context:
        log_entry["extra"] = _service_context
    if context:
        log_entry.setdefault("extra", {}).update(context)

    return json.dumps(log_entry, default=str)


def _json_excepthook(exc_type: type, exc_value: BaseException, exc_tb):
    """
    Global exception hook that formats unhandled exceptions as JSON.

    This catches exceptions that escape to the top level of the program,
    ensuring they are logged as single-line JSON for Loki.
    """
    # Don't interfere with KeyboardInterrupt
    if issubclass(exc_type, KeyboardInterrupt):
        sys.__excepthook__(exc_type, exc_value, exc_tb)
        return

    json_str = _format_exception_as_json(
        exc_type, exc_value, exc_tb, message=f"Unhandled exception: {exc_value}"
    )
    sys.stderr.write(json_str + "\n")
    sys.stderr.flush()


def _asyncio_exception_handler(loop, context: dict):
    """
    Asyncio exception handler that formats exceptions as JSON.

    This catches exceptions in asyncio tasks that aren't properly awaited
    or that escape task boundaries.
    """
    exception = context.get("exception")
    message = context.get("message", "Unhandled exception in asyncio task")

    if exception:
        json_str = _format_exception_as_json(
            type(exception),
            exception,
            exception.__traceback__,
            message=message,
            context={
                "asyncio_context": {
                    k: str(v) for k, v in context.items() if k != "exception"
                }
            },
        )
    else:
        # No exception object, just a message
        log_entry = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "level": "ERROR",
            "message": message,
            "logger": "asyncio_exception_handler",
            "function": None,
            "line": None,
            "file": None,
            "extra": {
                **_service_context,
                "asyncio_context": {k: str(v) for k, v in context.items()},
            },
        }
        json_str = json.dumps(log_entry, default=str)

    sys.stderr.write(json_str + "\n")
    sys.stderr.flush()


def install_exception_hooks(
    service_name: str | None = None,
    environment: str | None = None,
) -> None:
    """
    Install global exception hooks to ensure all Python exceptions are JSON-formatted.

    This can be called very early in application startup (even before configure_logging)
    to catch import-time errors. It's safe to call multiple times.

    Args:
        service_name: Name of the service (used in log context)
        environment: Environment name (e.g., "dev-local", "prod")
    """
    global _exception_hooks_installed, _service_context

    # Update service context if provided
    if service_name:
        _service_context["service"] = service_name
    if environment:
        _service_context["environment"] = environment

    if _exception_hooks_installed:
        return

    # Install sys.excepthook for unhandled exceptions
    sys.excepthook = _json_excepthook

    # Install asyncio exception handler
    # We need to be careful here because asyncio might not be fully initialized
    try:
        import asyncio

        # Try to get the running loop, but don't fail if there isn't one
        try:
            loop = asyncio.get_running_loop()
            loop.set_exception_handler(_asyncio_exception_handler)
        except RuntimeError:
            # No running loop yet - that's fine, we'll set it when one is created
            # by patching the default exception handler
            pass

        # Also set the handler on the default event loop policy
        # This ensures new loops get our handler
        _original_new_event_loop = asyncio.new_event_loop

        def _patched_new_event_loop():
            loop = _original_new_event_loop()
            loop.set_exception_handler(_asyncio_exception_handler)
            return loop

        asyncio.new_event_loop = _patched_new_event_loop

    except Exception:
        # If asyncio setup fails, continue anyway - sys.excepthook is the critical one
        pass

    _exception_hooks_installed = True


# Track whether warning hooks have been installed
_warning_hooks_installed = False


def _json_showwarning(
    message: Warning | str,
    category: type[Warning],
    filename: str,
    lineno: int,
    file=None,
    line: str | None = None,
) -> None:
    """
    Custom showwarning handler that formats warnings as single-line JSON.

    This replaces the default warnings.showwarning to ensure Python warnings
    (e.g., DeprecationWarning, FutureWarning, UserWarning) are captured as
    single log entries for proper Loki ingestion.
    """
    log_entry = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "level": "WARNING",
        "message": str(message),
        "logger": "warnings",
        "function": None,
        "line": lineno,
        "file": filename,
        "extra": {
            **_service_context,
            "warning_category": category.__name__,
            "source_line": line,
        },
    }

    json_str = json.dumps(log_entry, default=str)
    sys.stderr.write(json_str + "\n")
    sys.stderr.flush()


def install_warning_hooks() -> None:
    """
    Install custom warning handler to ensure Python warnings are JSON-formatted.

    This replaces the default warnings.showwarning function with one that
    outputs single-line JSON, ensuring warnings from third-party libraries
    are properly captured by Loki.
    """
    global _warning_hooks_installed

    if _warning_hooks_installed:
        return

    warnings.showwarning = _json_showwarning
    _warning_hooks_installed = True


# Track whether stdlib logging interception has been installed
_stdlib_logging_intercepted = False


def intercept_stdlib_logging(level: int = logging.DEBUG) -> None:
    """
    Intercept standard library logging and route to loguru.

    This captures logs from uvicorn, sqlalchemy, httpx, and other libraries
    that use standard library logging, ensuring they are formatted as JSON.

    Args:
        level: Minimum level to intercept (default: DEBUG to capture all)
    """
    global _stdlib_logging_intercepted

    if _stdlib_logging_intercepted:
        return

    # Install the intercept handler on the root logger
    logging.basicConfig(handlers=[InterceptHandler()], level=level, force=True)

    # Explicitly intercept loggers from libraries that configure their own handlers
    # at import time, bypassing our JSON setup:
    # - uvicorn: may be configured separately
    # - litellm: configures handlers with ANSI color formatting
    loggers_to_intercept = [
        "uvicorn",
        "uvicorn.error",
        "uvicorn.access",
        "LiteLLM",
        "LiteLLM Proxy",
        "LiteLLM Router",
        "litellm",
    ]
    for logger_name in loggers_to_intercept:
        lib_logger = logging.getLogger(logger_name)
        lib_logger.handlers = [InterceptHandler()]
        lib_logger.propagate = False

    _stdlib_logging_intercepted = True


def _get_trace_context() -> dict | None:
    """
    Get the current OpenTelemetry trace context if available.

    Returns {"trace_id": ..., "span_id": ...} if a valid span is active,
    None otherwise. Fails silently if OTel is not installed.
    """
    try:
        from opentelemetry import trace

        span = trace.get_current_span()
        ctx = span.get_span_context()
        if ctx and ctx.trace_id != 0:
            return {
                "trace_id": format(ctx.trace_id, "032x"),
                "span_id": format(ctx.span_id, "016x"),
            }
    except Exception:
        pass
    return None


_PROMOTED_KEYS = {
    "agent_session_id", "agent_id", "task_run_id",
    "task_id", "user_id", "trigger_id",
}


def json_serializer(record: dict) -> str:
    """Serialize a loguru record to a single-line JSON string."""
    message = record["message"]

    # Include full traceback in message if exception is present
    if record["exception"] is not None:
        exc = record["exception"]
        if exc.type and exc.value and exc.traceback:
            tb_str = "".join(
                traceback.format_exception(exc.type, exc.value, exc.traceback)
            )
            message = f"{message}: {tb_str}" if message else tb_str

    subset = {
        "timestamp": record["time"].isoformat(),
        "level": record["level"].name,
        "message": message,
        "logger": record["name"],
        "function": record["function"],
        "line": record["line"],
        "file": record["file"].path if record["file"] else None,
    }

    # Inject trace context if OpenTelemetry is active
    trace_ctx = _get_trace_context()
    if trace_ctx:
        subset["trace_id"] = trace_ctx["trace_id"]
        subset["span_id"] = trace_ctx["span_id"]

    # Promote domain IDs from extra to top-level for Loki query ergonomics
    if record["extra"]:
        for key in _PROMOTED_KEYS:
            if key in record["extra"]:
                subset[key] = record["extra"][key]

    # Include any extra data bound to the logger
    if record["extra"]:
        subset["extra"] = record["extra"]

    return json.dumps(subset, default=str)


def json_sink(message):
    """Sink that writes JSON-formatted log records to stderr."""
    record = message.record
    json_str = json_serializer(record)
    sys.stderr.write(json_str + "\n")
    sys.stderr.flush()


def configure_logging(
    service_name: str | None = None,
    environment: str | None = None,
    settings: "CommonSettings | None" = None,
    level: str = "INFO",
    json_format: bool = True,
) -> None:
    """
    Configure loguru for JSON-formatted logging suitable for Loki ingestion.

    This should be called once at application startup, before any logging occurs.

    Args:
        service_name: Name of the service (used in log context)
        environment: Environment name (e.g., "dev-local", "prod")
        settings: CommonSettings instance to extract service_name and environment from
        level: Minimum log level to capture (default: INFO)
        json_format: If True, output JSON logs. If False, use human-readable format.
                     Set to False for local development if you prefer readable logs.
    """
    # Extract values from settings if provided
    if settings is not None:
        service_name = service_name or settings.SERVICE_NAME
        environment = environment or settings.ENVIRONMENT

    # Install global exception hooks to catch errors outside loguru
    # (e.g., import errors, asyncio task exceptions)
    install_exception_hooks(service_name=service_name, environment=environment)

    # Install warning hooks to capture Python warnings as JSON
    if json_format:
        install_warning_hooks()

    # Intercept standard library logging (uvicorn, sqlalchemy, etc.)
    if json_format:
        intercept_stdlib_logging()

    # Remove default handler
    logger.remove()

    if json_format:
        # Add JSON handler for Loki-compatible output
        logger.add(
            json_sink,
            level=level,
            serialize=False,  # We handle serialization ourselves
            backtrace=True,
            diagnose=True,
        )
    else:
        # Human-readable format for local development
        logger.add(
            sys.stderr,
            level=level,
            format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>",
            backtrace=True,
            diagnose=True,
        )

    # Bind service context if provided
    if service_name or environment:
        context = {}
        if service_name:
            context["service"] = service_name
        if environment:
            context["environment"] = environment
        logger.configure(extra=context)
