"""
Utilities for logging.
"""

# Needed for loguru annotations
from __future__ import annotations

import inspect
import logging
import sys
from functools import wraps
from pathlib import Path
from typing import Any, Callable

import loguru
from flask import g, has_app_context
from loguru import logger
from sqlalchemy.exc import OperationalError

from ..web.database.models import ResearchLog
from ..web.services.socket_service import SocketIOService
from .db_utils import get_db_session

_LOG_DIR = Path(__file__).parents[2] / "data" / "logs"
_LOG_DIR.mkdir(exist_ok=True)
"""
Default log directory to use.
"""


class InterceptHandler(logging.Handler):
    """
    Intercepts logging messages and forwards them to Loguru's logger.
    """

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

        # Find caller from where originated the logged message.
        frame, depth = inspect.currentframe(), 0
        while frame:
            filename = frame.f_code.co_filename
            is_logging = filename == logging.__file__
            is_frozen = "importlib" in filename and "_bootstrap" in filename
            if depth > 0 and not (is_logging or is_frozen):
                break
            frame = frame.f_back
            depth += 1

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


def log_for_research(
    to_wrap: Callable[[int, ...], Any],
) -> Callable[[int, ...], Any]:
    """
    Decorator for a function that's part of the research process. It expects the function to
    take the research ID as the first parameter, and configures all log
    messages made during this request to include the research ID.

    Args:
        to_wrap: The function to wrap. Should take the research ID as the first parameter.

    Returns:
        The wrapped function.

    """

    @wraps(to_wrap)
    def wrapped(research_id: int, *args: Any, **kwargs: Any) -> Any:
        g.research_id = research_id
        to_wrap(research_id, *args, **kwargs)
        g.pop("research_id")

    return wrapped


def _get_research_id() -> int | None:
    """
    Gets the current research ID, if present.

    Returns:
        The current research ID, or None if it does not exist.

    """
    research_id = None
    if has_app_context():
        research_id = g.get("research_id")
    return research_id


def database_sink(message: loguru.Message) -> None:
    """
    Sink that saves messages to the database.

    Args:
        message: The log message to save.

    """
    record = message.record
    research_id = _get_research_id()

    # Create a new database entry.
    db_log = ResearchLog(
        timestamp=record["time"],
        message=str(message),
        module=record["name"],
        function=record["function"],
        line_no=int(record["line"]),
        level=record["level"].name,
        research_id=research_id,
    )

    # Save the entry to the database.
    db_session = get_db_session()
    try:
        db_session.add(db_log)
        db_session.commit()
    except OperationalError:
        # Something else is probably using the DB and we can't write to it
        # right now. Ignore this.
        db_session.rollback()
        return


def frontend_progress_sink(message: loguru.Message) -> None:
    """
    Sink that sends messages to the frontend.

    Args:
        message: The log message to send.

    """
    research_id = _get_research_id()
    if research_id is None:
        # If we don't have a research ID, don't send anything.
        return

    record = message.record
    frontend_log = dict(
        log_entry=dict(
            message=record["message"],
            type=record["level"].name,
            time=record["time"].isoformat(),
        ),
    )
    SocketIOService().emit_to_subscribers(
        "progress", research_id, frontend_log, enable_logging=False
    )


def config_logger(name: str) -> None:
    """
    Configures the default logger.

    Args:
        name: The name to use for the log file.

    """
    logger.enable("local_deep_research")
    logger.remove()

    # Log more important stuff to the console.
    logger.add(sys.stderr, level="INFO")
    logger.add(
        _LOG_DIR / f"{name}.log",
        level="DEBUG",
        enqueue=True,
        rotation="00:00",
        retention="30 days",
        compression="zip",
    )
    logger.add(database_sink)
    logger.add(frontend_progress_sink)

    # Add a special log level for milestones.
    try:
        logger.level("milestone", no=26, color="<magenta><bold>")
    except ValueError:
        # Level already exists, that's fine
        pass
