from __future__ import annotations

import logging
import logging.handlers
import random
import time
from contextlib import contextmanager
from pathlib import Path
from typing import Iterator

from ._compat import optional_import

DEFAULT_LOGGER_NAME = "aniket"
LOG_FORMAT = "%(asctime)s | %(name)s | %(levelname)s | %(message)s"
_FILE_HANDLERS: dict[str, logging.Handler] = {}
SUPPORTED_LOG_TYPES = (
    "debug", "info", "warning", "error", "critical", "notset",
    "query", "trace", "success", "exception", "table", "divider", "audit", "timer", "json",
)

# Custom log levels
TRACE = 5
SUCCESS = 25
AUDIT = 45

logging.addLevelName(TRACE, "TRACE")
logging.addLevelName(SUCCESS, "SUCCESS")
logging.addLevelName(AUDIT, "AUDIT")

# ANSI color codes
_LEVEL_COLORS: dict[str, str] = {
    "TRACE":    "\033[2m",        # dim
    "DEBUG":    "\033[36m",       # cyan
    "INFO":     "\033[0m",        # default
    "SUCCESS":  "\033[32m",       # green
    "WARNING":  "\033[33m",       # yellow
    "ERROR":    "\033[31m",       # red
    "CRITICAL": "\033[1;31m",     # bold red
    "AUDIT":    "\033[35m",       # magenta
}
_RESET = "\033[0m"


class _ColorFormatter(logging.Formatter):
    def format(self, record: logging.LogRecord) -> str:
        color = _LEVEL_COLORS.get(record.levelname, "")
        return f"{color}{super().format(record)}{_RESET}"


class _JsonFormatter(logging.Formatter):
    def format(self, record: logging.LogRecord) -> str:
        import json as _json

        data: dict = {
            "time": self.formatTime(record),
            "name": record.name,
            "level": record.levelname,
            "msg": record.getMessage(),
        }
        context = getattr(record, "_context", None)
        if context:
            data.update(context)
        if record.exc_info:
            data["exc"] = self.formatException(record.exc_info)
        return _json.dumps(data)


class _SamplingFilter(logging.Filter):
    def __init__(self, rate: float) -> None:
        super().__init__()
        self.rate = rate

    def filter(self, _record: logging.LogRecord) -> bool:
        return random.random() < self.rate


def _make_formatter(*, json: bool, colored: bool) -> logging.Formatter:
    if json:
        return _JsonFormatter()
    if colored:
        return _ColorFormatter(LOG_FORMAT)
    return logging.Formatter(LOG_FORMAT)


def _ensure_stream_handler(logger: logging.Logger, *, json: bool = False) -> None:
    formatter = _make_formatter(json=json, colored=not json)
    for handler in logger.handlers:
        if isinstance(handler, logging.StreamHandler) and not isinstance(handler, logging.FileHandler):
            handler.setFormatter(formatter)
            return
    stream_handler = logging.StreamHandler()
    stream_handler.setFormatter(formatter)
    logger.addHandler(stream_handler)


def _resolve_log_path(file_name: str | Path) -> Path:
    path = Path(file_name)
    if not path.suffix:
        path = path.with_suffix(".log")
    path.parent.mkdir(parents=True, exist_ok=True)
    return path


def get_logger(
    name: str = DEFAULT_LOGGER_NAME,
    *,
    file_name: str | Path | None = None,
    json: bool = False,
    rotate: bool = False,
    sample_rate: float | None = None,
) -> logging.Logger:
    logger = logging.getLogger(name)
    logger.setLevel(TRACE)
    _ensure_stream_handler(logger, json=json)

    if sample_rate is not None:
        logger.filters = [f for f in logger.filters if not isinstance(f, _SamplingFilter)]
        logger.addFilter(_SamplingFilter(sample_rate))

    if file_name:
        path = _resolve_log_path(file_name)
        handler_key = str(path.resolve())
        file_handler = _FILE_HANDLERS.get(handler_key)

        if file_handler is None:
            if rotate:
                file_handler = logging.handlers.RotatingFileHandler(
                    path, maxBytes=10 * 1024 * 1024, backupCount=5
                )
            else:
                file_handler = logging.FileHandler(path)
            file_handler.setFormatter(_make_formatter(json=json, colored=False))
            _FILE_HANDLERS[handler_key] = file_handler

        if file_handler not in logger.handlers:
            logger.addHandler(file_handler)

    return logger


# ── SQL helpers ───────────────────────────────────────────────────────────────

def _coerce_sqlalchemy_statement(msg: object) -> object:
    statement = getattr(msg, "statement", None)
    if statement is not None:
        return statement
    return msg


def _resolve_sqlalchemy_dialect(dialect: object | None, statement: object) -> object | None:
    candidate = dialect

    if candidate is None:
        candidate = getattr(statement, "bind", None)

    if candidate is None:
        session = getattr(statement, "session", None)
        if session is not None:
            get_bind = getattr(session, "get_bind", None)
            if callable(get_bind):
                try:
                    candidate = get_bind()
                except Exception:
                    candidate = None
            if candidate is None:
                candidate = getattr(session, "bind", None)

    if candidate is None:
        return None

    if isinstance(candidate, str):
        dialect_module = optional_import(f"sqlalchemy.dialects.{candidate}")
        if dialect_module is None:
            return None
        dialect_factory = getattr(dialect_module, "dialect", None)
        return dialect_factory() if callable(dialect_factory) else None

    sync_engine = getattr(candidate, "sync_engine", None)
    if sync_engine is not None:
        candidate = sync_engine

    sync_connection = getattr(candidate, "sync_connection", None)
    if sync_connection is not None:
        candidate = sync_connection

    return getattr(candidate, "dialect", candidate)


def _render_query_message(msg: object, dialect: object | None = None) -> str:
    statement = _coerce_sqlalchemy_statement(msg)
    compile_fn = getattr(statement, "compile", None)
    if not callable(compile_fn):
        return str(msg)

    compile_kwargs = {"literal_binds": True}
    resolved_dialect = _resolve_sqlalchemy_dialect(dialect, statement)

    try:
        if resolved_dialect is not None:
            compiled = compile_fn(dialect=resolved_dialect, compile_kwargs=compile_kwargs)
        else:
            compiled = compile_fn(compile_kwargs=compile_kwargs)
    except TypeError:
        try:
            compiled = compile_fn()
        except Exception:
            return str(msg)
    except Exception:
        return str(msg)

    return f"SQL: {str(compiled)}"


# ── Table renderer ────────────────────────────────────────────────────────────

def _render_table(data: object) -> str:
    if not isinstance(data, list) or not data:
        return f"(table: expected list of dicts, got {type(data).__name__})"
    headers = list(data[0].keys())
    col_widths = {h: len(str(h)) for h in headers}
    for row in data:
        for h in headers:
            col_widths[h] = max(col_widths[h], len(str(row.get(h, ""))))
    sep = "+" + "+".join("-" * (w + 2) for w in col_widths.values()) + "+"
    header_row = "|" + "|".join(f" {str(h):<{col_widths[h]}} " for h in headers) + "|"
    rows = [
        "|" + "|".join(f" {str(row.get(h, '')):<{col_widths[h]}} " for h in headers) + "|"
        for row in data
    ]
    return "\n".join([sep, header_row, sep] + rows + [sep])


# ── Timer context manager ─────────────────────────────────────────────────────

@contextmanager
def _timer_context(label: str, target_logger: logging.Logger) -> Iterator[None]:
    start = time.perf_counter()
    try:
        yield
    finally:
        elapsed = time.perf_counter() - start
        target_logger.info(f"TIMER[{label}]: {elapsed:.4f}s")


# ── Internal log dispatcher ───────────────────────────────────────────────────

def _log_with_context(
    logger: logging.Logger,
    level: int,
    msg: object,
    context: dict | None,
    exc_info: bool | BaseException,
) -> None:
    extra = {"_context": context} if context else {}
    has_json = any(
        isinstance(getattr(h, "formatter", None), _JsonFormatter)
        for h in logger.handlers
    )

    display_msg = msg
    if context and not has_json:
        ctx_str = " | " + " ".join(f"{k}={v}" for k, v in context.items())
        display_msg = f"{msg}{ctx_str}"

    ei: bool | tuple | None = None
    if isinstance(exc_info, BaseException):
        ei = (type(exc_info), exc_info, exc_info.__traceback__)
    elif exc_info:
        ei = True

    logger.log(level, display_msg, exc_info=ei, extra=extra)


# ── Public API ────────────────────────────────────────────────────────────────

def logs(
    msg: object = "",
    type: str = "info",
    file_name: str | Path | None = None,
    logger: logging.Logger | None = None,
    dialect: object | None = None,
    context: dict | None = None,
    exc_info: bool | BaseException = False,
    redact: list[str] | None = None,
    sample_rate: float | None = None,
    indent: int = 4,
):
    """
    Unified logging function.

    Parameters
    ----------
    msg         : the message or data to log
    type        : log level / mode — one of SUPPORTED_LOG_TYPES
    file_name   : optional file path to also write logs to
    logger      : use a specific logger instead of the default
    dialect     : SQLAlchemy dialect for query rendering (type="query" only)
    context     : dict of extra fields attached to the log line
    exc_info    : True to capture current exception, or pass an exception object
    redact      : list of context keys whose values should be masked as "***"
    sample_rate : 0.0–1.0 — probabilistically drop this fraction of log calls

    Special types
    -------------
    "trace"     : custom level below DEBUG (dim output)
    "success"   : custom level between INFO and WARNING (green output)
    "audit"     : custom level between ERROR and CRITICAL (magenta output)
    "exception" : logs at ERROR level with full traceback attached
    "query"     : renders a SQLAlchemy statement with bound parameters
    "table"     : msg must be list[dict], renders as ASCII table
    "divider"   : prints a visual section separator; msg is the label
    "timer"     : returns a context manager that logs elapsed time on exit

    Examples
    --------
    logs("hello")
    logs("watch out", type="warning")
    logs("secret", context={"token": "abc"}, redact=["token"])
    logs([{"id": 1, "name": "Alice"}], type="table")
    with logs("db query", type="timer"):
        time.sleep(0.1)
    """
    if sample_rate is not None and random.random() >= sample_rate:
        return

    target_logger = logger or get_logger(file_name=file_name)
    normalized_type = str(type).lower()

    safe_context: dict | None = None
    if context:
        safe_context = {
            k: "***" if (redact and k in redact) else v
            for k, v in context.items()
        }

    if normalized_type == "json":
        import json as _json
        try:
            rendered = _json.dumps(msg, indent=indent, default=str)
        except Exception:
            rendered = str(msg)
        _log_with_context(target_logger, logging.INFO, rendered, safe_context, exc_info)
        return

    if normalized_type == "query":
        _log_with_context(
            target_logger, logging.INFO,
            _render_query_message(msg, dialect=dialect),
            safe_context, exc_info,
        )
        return

    if normalized_type == "timer":
        return _timer_context(str(msg) or "block", target_logger)

    if normalized_type == "table":
        _log_with_context(target_logger, logging.INFO, _render_table(msg), safe_context, exc_info)
        return

    if normalized_type == "divider":
        label = str(msg)
        width = 60
        line = (
            f"{'─' * 4} {label} {'─' * max(0, width - len(label) - 6)}"
            if label else "─" * width
        )
        _log_with_context(target_logger, logging.INFO, line, safe_context, exc_info)
        return

    if normalized_type == "exception":
        _log_with_context(
            target_logger, logging.ERROR, msg,
            safe_context, exc_info if exc_info else True,
        )
        return

    if normalized_type == "trace":
        _log_with_context(target_logger, TRACE, msg, safe_context, exc_info)
        return

    if normalized_type == "success":
        _log_with_context(target_logger, SUCCESS, msg, safe_context, exc_info)
        return

    if normalized_type == "audit":
        _log_with_context(target_logger, AUDIT, msg, safe_context, exc_info)
        return

    level_name = normalized_type.upper()
    level = getattr(logging, level_name, logging.INFO)
    _log_with_context(target_logger, level, msg, safe_context, exc_info)
