from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any
import json
import sys
import time
import traceback
import logging


_LEVELS = {
    "DEBUG": logging.DEBUG,
    "INFO": logging.INFO,
    "WARNING": logging.WARNING,
    "WARN": logging.WARNING,
    "ERROR": logging.ERROR,
    "CRITICAL": logging.CRITICAL,
}


@dataclass(frozen=True)
class LogEvent:
    ts: float
    level: str
    logger: str
    message: str
    fields: dict[str, Any] = field(default_factory=dict)
    exc: str | None = None


class LogSink:
    def emit(self, event: LogEvent) -> None:
        raise NotImplementedError


class ConsoleSink(LogSink):
    def __init__(self, stream=None, json_output: bool = True):
        self._stream = stream if stream is not None else sys.stdout
        self._json_output = json_output

    def emit(self, event: LogEvent) -> None:
        if self._json_output:
            payload = {
                "ts": event.ts,
                "level": event.level,
                "logger": event.logger,
                "message": event.message,
                "fields": event.fields,
            }
            if event.exc:
                payload["exc"] = event.exc
            self._stream.write(json.dumps(payload, ensure_ascii=True) + "\n")
            self._stream.flush()
            return
        line = f"[{event.level}] {event.logger}: {event.message}"
        if event.fields:
            line = f"{line} | {event.fields}"
        if event.exc:
            line = f"{line}\n{event.exc}"
        self._stream.write(line + "\n")
        self._stream.flush()


class CustomLogger:
    """Lightweight structured logger.

    This logger is intentionally minimal and pluggable. It emits LogEvent objects
    to sinks (ConsoleSink by default). Future LogAgent/EventStream integration
    can be implemented as additional sinks.
    """

    def __init__(
        self,
        name: str | None = None,
        *,
        level: str | int = "INFO",
        sinks: list[LogSink] | None = None,
        context: dict[str, Any] | None = None,
    ):
        self.name = name or "Logger"
        self._level = self._coerce_level(level)
        self._sinks = sinks if sinks is not None else [ConsoleSink()]
        self._context = dict(context or {})

    def bind(self, **fields: Any) -> "CustomLogger":
        merged = dict(self._context)
        merged.update(fields)
        return CustomLogger(
            self.name,
            level=self._level,
            sinks=self._sinks,
            context=merged,
        )

    def set_level(self, level: str | int) -> None:
        self._level = self._coerce_level(level)

    def log(self, level: str | int, message: str, *args, exc_info: bool = False, **fields):
        level_int = self._coerce_level(level)
        if level_int < self._level:
            return
        if args:
            message = message % args
        merged_fields = dict(self._context)
        merged_fields.update(fields)
        exc = None
        if exc_info:
            exc = "".join(traceback.format_exception(*sys.exc_info()))
        event = LogEvent(
            ts=time.time(),
            level=self._level_name(level_int),
            logger=self.name,
            message=message,
            fields=merged_fields,
            exc=exc,
        )
        for sink in self._sinks:
            sink.emit(event)

    def debug(self, message: str, *args, **fields):
        self.log("DEBUG", message, *args, **fields)

    def info(self, message: str, *args, **fields):
        self.log("INFO", message, *args, **fields)

    def warning(self, message: str, *args, **fields):
        self.log("WARNING", message, *args, **fields)

    def error(self, message: str, *args, exc_info: bool = False, **fields):
        self.log("ERROR", message, *args, exc_info=exc_info, **fields)

    def critical(self, message: str, *args, **fields):
        self.log("CRITICAL", message, *args, **fields)

    def exception(self, message: str, *args, **fields):
        self.error(message, *args, exc_info=True, **fields)

    @staticmethod
    def _coerce_level(level: str | int) -> int:
        if isinstance(level, int):
            return level
        level_str = level.upper()
        if level_str not in _LEVELS:
            raise ValueError(f"Invalid log level: {level}")
        return _LEVELS[level_str]

    @staticmethod
    def _level_name(level_int: int) -> str:
        for name, value in _LEVELS.items():
            if value == level_int:
                return name
        return str(level_int)

    @classmethod
    def get_logger(
        cls,
        name: str | None = None,
        *,
        level: str | int = "INFO",
    ) -> logging.Logger:
        """Compatibility helper that returns std logging.Logger.

        This keeps old call sites working while CustomLogger moves to structured sinks.
        """
        logger = logging.getLogger(name or "Logger")
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter("%(asctime)s [%(levelname)s] %(name)s: %(message)s")
            handler.setFormatter(formatter)
            logger.addHandler(handler)
        logger.setLevel(cls._coerce_level(level))
        logger.propagate = False
        return logger
