from __future__ import annotations

import logging
import os
import re
import sys
from dataclasses import dataclass, field
from datetime import datetime
from statistics import mean
from typing import Any, TYPE_CHECKING

import click
from relationalai.runtime.execution_summary import (
    SUMMARY_EVENT_GRAPH_INDEX_ENABLED,
    SUMMARY_EVENT_SERVICE_ENDPOINT_CACHE_HIT,
    SUMMARY_EVENT_TOKEN_CACHE_HIT,
)

if TYPE_CHECKING:
    from relationalai.config.config import Config

EXECUTION_LOGGER_NAME = "relationalai.client.execution"
_CACHE_EVENT_LABELS: dict[str, str] = {
    SUMMARY_EVENT_SERVICE_ENDPOINT_CACHE_HIT: "Direct Access endpoint cache",
    SUMMARY_EVENT_TOKEN_CACHE_HIT: "Direct Access token cache",
}

_KIND_RE = re.compile(r"(kind=)(\S+)")
_OP_RE = re.compile(r"(op=)(\S+)")


class ExecutionLogFormatter(logging.Formatter):
    """Formatter for buffered CLI execution logs with light field coloring."""

    def __init__(self, fmt: str, *, enable_color: bool) -> None:
        super().__init__(fmt)
        self._fmt_raw = fmt
        self._enable_color = enable_color

    def _format_kind(self, value: str) -> str:
        value = value.upper()
        if not self._enable_color:
            return value
        return click.style(value, fg="green", bold=True)

    def _format_op(self, value: str) -> str:
        if not self._enable_color:
            return value
        return click.style(value, fg="green", bold=True)

    def _format_level(self, value: str) -> str:
        level = str(value or "").upper()
        if not self._enable_color:
            return level
        style_by_level = {
            "DEBUG": {"fg": "blue", "dim": True},
            "INFO": {"fg": "cyan", "bold": True},
            "WARNING": {"fg": "yellow", "bold": True},
            "ERROR": {"fg": "red", "bold": True},
            "CRITICAL": {"fg": "magenta", "bold": True},
        }
        return click.style(level, **style_by_level.get(level, {"bold": True}))

    def _split_message(self, msg: str) -> tuple[str | None, str | None, str]:
        kind_match = _KIND_RE.search(msg)
        op_match = _OP_RE.search(msg)

        kind_part = None
        if kind_match is not None:
            kind_part = f"{kind_match.group(1)}{self._format_kind(kind_match.group(2))}"

        op_part = None
        if op_match is not None:
            op_part = f"{op_match.group(1)}{self._format_op(op_match.group(2))}"

        if kind_match is None and op_match is None:
            return None, None, msg

        remaining = msg
        for match in sorted([m for m in (kind_match, op_match) if m is not None], key=lambda m: m.start(), reverse=True):
            remaining = (remaining[:match.start()] + remaining[match.end():]).strip()
        remaining = re.sub(r"\s{2,}", " ", remaining)
        return kind_part, op_part, remaining

    def format(self, record: logging.LogRecord) -> str:
        kind_part, op_part, msg = self._split_message(record.getMessage())
        if self._fmt_raw == "%(message)s":
            pieces = [p for p in (kind_part, op_part, msg) if p]
            return " ".join(pieces)

        pieces = [self.formatTime(record, self.datefmt), self._format_level(record.levelname)]
        if kind_part is not None:
            pieces.append(kind_part)
        if op_part is not None:
            pieces.append(op_part)
        pieces.append(record.name)
        if msg:
            pieces.append(msg)
        formatted = " ".join(pieces)

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


def _format_local_timestamp(timestamp: float) -> str:
    return datetime.fromtimestamp(timestamp).astimezone().strftime("%Y-%m-%d %H:%M:%S")


def _stderr_supports_color() -> bool:
    if os.environ.get("NO_COLOR"):
        return False
    return bool(getattr(sys.stderr, "isatty", lambda: False)())


def _format_context_value(value: str, *, enable_color: bool) -> str:
    if not enable_color:
        return value
    return click.style(value, fg="green", bold=True)


def _format_context_emphasis(value: str, *, enable_color: bool) -> str:
    if not enable_color:
        return value
    return click.style(value, bold=True)


@dataclass
class ExecutionSummaryState:
    graph_index_enabled: bool = False
    cache_events: dict[str, str | None] = field(default_factory=dict)

    def clear(self) -> None:
        self.graph_index_enabled = False
        self.cache_events.clear()


class BufferedExecutionLogHandler(logging.Handler):
    def __init__(self) -> None:
        super().__init__(level=logging.INFO)
        self.lines: list[str] = []
        self.summary = ExecutionSummaryState()

    def emit(self, record: logging.LogRecord) -> None:
        try:
            summary_event = getattr(record, "summary_event", None)
            if summary_event == SUMMARY_EVENT_GRAPH_INDEX_ENABLED:
                self.summary.graph_index_enabled = True
            elif isinstance(summary_event, str) and summary_event in _CACHE_EVENT_LABELS:
                created_at_ts = getattr(record, "cache_created_at_ts", None)
                created_at = _format_local_timestamp(float(created_at_ts)) if isinstance(created_at_ts, (int, float)) else None
                existing = self.summary.cache_events.get(summary_event)
                if summary_event not in self.summary.cache_events or (existing is None and created_at is not None):
                    self.summary.cache_events[summary_event] = created_at
            msg = self.format(record)
        except Exception:
            # Keep diagnostics best-effort; avoid breaking CLI execution.
            msg = record.getMessage()
        self.lines.append(msg)


def enable_execution_logging_buffer(cfg: "Config") -> BufferedExecutionLogHandler | None:
    """Buffer execution middleware logs when enabled in config.

    This avoids interleaving logs with spinners and tables. Buffered logs should be
    flushed at the end of the CLI command.
    """
    if not cfg.execution.logging:
        return None

    log = logging.getLogger(EXECUTION_LOGGER_NAME)
    level = logging.INFO
    previous_propagate = log.propagate
    log.setLevel(level)

    existing = next((h for h in log.handlers if isinstance(h, BufferedExecutionLogHandler)), None)
    if existing is not None:
        return existing

    log.propagate = False
    h = BufferedExecutionLogHandler()
    h._previous_propagate = previous_propagate  # type: ignore[attr-defined]
    h.setLevel(level)
    h.setFormatter(
        ExecutionLogFormatter(
            "%(asctime)s %(levelname)s %(name)s %(message)s",
            enable_color=_stderr_supports_color(),
        )
    )
    log.addHandler(h)
    return h


def _effective_connection_name(cfg: "Config") -> str | None:
    default_conn = cfg.default_connection
    if isinstance(default_conn, str) and default_conn.strip() and default_conn in cfg.connections:
        return default_conn
    if len(cfg.connections) == 1:
        return next(iter(cfg.connections.keys()))
    return None


def _format_execution_context(
    cfg: "Config | None",
    handler: BufferedExecutionLogHandler | None,
    *,
    enable_color: bool = False,
) -> str | None:
    if cfg is None:
        return None

    lines: list[str] = []
    if cfg.direct_access:
        lines.append(f"Direct access: {_format_context_value('enabled', enable_color=enable_color)}")
    connection_name = _effective_connection_name(cfg)
    if connection_name is not None:
        conn = cfg.connections.get(connection_name)
        auth = getattr(conn, "authenticator", None) if conn is not None else None
        if isinstance(auth, str) and auth.strip():
            lines.append(
                "Authenticator: "
                f"{_format_context_value(auth, enable_color=enable_color)} "
                f"(connection: {_format_context_value(connection_name, enable_color=enable_color)})"
            )
        else:
            lines.append(f"Connection: {_format_context_value(connection_name, enable_color=enable_color)}")

    if handler is not None and handler.summary.graph_index_enabled:
        lines.append("Graph index: enabled")

    if handler is not None:
        cache_labels = sorted(
            (
                f"{_CACHE_EVENT_LABELS[event]} ({_format_context_emphasis(timestamp, enable_color=enable_color)})"
                if timestamp
                else _CACHE_EVENT_LABELS[event]
            )
            for event, timestamp in handler.summary.cache_events.items()
        )
        if cache_labels:
            lines.append(f"Caching used: {', '.join(cache_labels)}")

    if not lines:
        return None

    return "\n".join(["", "Execution context", "-----------------", *lines])


def flush_execution_logging_buffer(handler: BufferedExecutionLogHandler | None, *, cfg: "Config | None" = None) -> None:
    summary = _format_execution_context(cfg, handler, enable_color=_stderr_supports_color())
    if handler is None:
        return
    if not handler.lines and summary is None:
        return
    sections: list[str] = []
    if summary is not None:
        sections.append(summary.rstrip())
    if handler.lines:
        sections.append("Execution logs\n--------------\n" + "\n".join(handler.lines).rstrip())
    sys.stderr.write("\n" + "\n\n".join(section for section in sections if section).rstrip() + "\n")
    handler.lines.clear()
    handler.summary.clear()


def disable_execution_logging_buffer(handler: BufferedExecutionLogHandler | None) -> None:
    if handler is None:
        return
    log = logging.getLogger(EXECUTION_LOGGER_NAME)
    try:
        log.removeHandler(handler)
    except Exception:
        pass
    log.propagate = bool(getattr(handler, "_previous_propagate", True))


def _format_execution_metrics(metrics: Any | None) -> str | None:
    if metrics is None:
        return None

    counters: dict[str, Any] | None = getattr(metrics, "counters", None)
    timings: dict[str, Any] | None = getattr(metrics, "timings_ms", None)

    has_counters = isinstance(counters, dict) and any(v for v in counters.values())
    has_timings = isinstance(timings, dict) and any(bool(v) for v in timings.values())
    if not has_counters and not has_timings:
        return None

    lines: list[str] = []
    lines.append("")
    lines.append("Execution metrics")
    lines.append("-----------------")

    if has_counters and counters is not None:
        lines.append("Counters:")
        for k in sorted(counters.keys()):
            v = counters.get(k, 0)
            if v:
                lines.append(f"- {k}: {v}")

    if has_timings and timings is not None:
        lines.append("")
        lines.append("Timings (ms):")
        for k in sorted(timings.keys()):
            xs = timings.get(k) or []
            if not xs:
                continue
            xs_sorted = sorted(float(x) for x in xs)
            n = len(xs_sorted)
            p50 = xs_sorted[int(0.50 * (n - 1))]
            p95 = xs_sorted[int(0.95 * (n - 1))]
            lines.append(
                f"- {k}: n={n} min={xs_sorted[0]:.2f} avg={mean(xs_sorted):.2f} "
                f"p50={p50:.2f} p95={p95:.2f} max={xs_sorted[-1]:.2f}"
            )

    return "\n".join(lines)


def print_execution_metrics_if_enabled(cfg: "Config", metrics: Any | None) -> None:
    if not cfg.execution.metrics:
        return
    out = _format_execution_metrics(metrics)
    if out is None:
        return
    sys.stdout.write(out.rstrip() + "\n")

