import asyncio
import json
import logging
import os
import socket
import threading
import uuid
from contextlib import contextmanager
from contextvars import ContextVar
from pathlib import Path
from typing import Literal

import hydra
import zmq
import zmq.asyncio
from rich.console import Console
from rich.logging import RichHandler

# Task context variable
TASK_CONTEXT_VAR: ContextVar[str | None] = ContextVar("CURRENT_TASK_ID", default=None)

# Maximum length for debug log content (to prevent log bloat)
DEBUG_LOG_MAX_LENGTH = 500


def truncate_for_log(content: str, max_length: int = DEBUG_LOG_MAX_LENGTH) -> str:
    """Truncate long content for logging to prevent log bloat.

    Args:
        content: The content to potentially truncate
        max_length: Maximum length before truncation (default: 500)

    Returns:
        Truncated content with length indicator if truncated
    """
    if not content:
        return content
    content_str = str(content)
    if len(content_str) <= max_length:
        return content_str
    return f"{content_str[:max_length]}... ({len(content_str)} chars total)"


# Context variables for logging identifiers
trace_id_var: ContextVar[str] = ContextVar("trace_id", default="")
chat_id_var: ContextVar[str] = ContextVar("chat_id", default="")
user_id_var: ContextVar[str] = ContextVar("user_id", default="")
message_id_var: ContextVar[str] = ContextVar("message_id", default="")

# Global variable to store the actual ZMQ address being used
_ZMQ_ADDRESS: str = "tcp://127.0.0.1:6000"


def generate_trace_id() -> str:
    """Generate a new trace ID"""
    return str(uuid.uuid4())


def get_trace_id() -> str:
    """Get current trace ID from context"""
    return trace_id_var.get()


def set_trace_id(trace_id: str = None) -> str:
    """Set trace ID in context, generate new one if not provided"""
    if trace_id is None:
        trace_id = generate_trace_id()
    trace_id_var.set(trace_id)
    return trace_id


def get_chat_id() -> str:
    """Get current chat ID from context"""
    return chat_id_var.get()


def set_chat_id(chat_id: str) -> str:
    """Set chat ID in context"""
    chat_id_var.set(chat_id)
    return chat_id


def get_user_id() -> str:
    """Get current user ID from context"""
    return user_id_var.get()


def set_user_id(user_id: str) -> str:
    """Set user ID in context"""
    user_id_var.set(user_id)
    return user_id


def get_message_id() -> str:
    """Get current message ID from context"""
    return message_id_var.get()


def set_message_id(message_id: str) -> str:
    """Set message ID in context"""
    message_id_var.set(message_id)
    return message_id


def find_available_port(start_port: int = 6000, max_attempts: int = 10) -> int:
    """Find an available port starting from start_port."""
    for port in range(start_port, start_port + max_attempts):
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.bind(("127.0.0.1", port))
                return port
        except OSError:
            continue
    raise RuntimeError(
        f"Could not find an available port in range {start_port}-{start_port + max_attempts - 1}"
    )


def get_zmq_address() -> str:
    """Get the current ZMQ address."""
    return _ZMQ_ADDRESS


def set_zmq_address(address: str) -> None:
    """Set the ZMQ address."""
    global _ZMQ_ADDRESS
    _ZMQ_ADDRESS = address


def _extract_port_from_address(addr: str) -> int:
    """Extract port number from ZMQ address."""
    try:
        return int(addr.split(":")[-1])
    except (ValueError, IndexError):
        return 6000


def _bind_zmq_socket(sock, bind_addr: str) -> str:
    """Bind ZMQ socket to an available port and return the actual address."""
    port = _extract_port_from_address(bind_addr)

    try:
        available_port = find_available_port(port)
        actual_addr = f"tcp://127.0.0.1:{available_port}"
        sock.bind(actual_addr)
        return actual_addr
    except RuntimeError:
        # Fallback to random port
        port = sock.bind_to_random_port("tcp://127.0.0.1")
        return f"tcp://127.0.0.1:{port}"


class ZMQLogHandler(logging.Handler):
    def __init__(self, addr=None, tool_name="unknown_tool"):
        super().__init__()
        ctx = zmq.Context()
        self.sock = ctx.socket(zmq.PUSH)

        # Use the global ZMQ address if no specific address is provided
        if addr is None:
            addr = get_zmq_address()

        # Try to connect to the address
        try:
            self.sock.connect(addr)
            logging.getLogger(__name__).info(f"ZMQ handler connected to: {addr}")
        except zmq.error.ZMQError as e:
            # If connection fails, disable the handler
            logging.getLogger(__name__).warning(f"Could not connect to ZMQ listener at {addr}: {e}")
            logging.getLogger(__name__).warning("Disabling ZMQ logging for this handler")
            self.sock = None

        self.task_id = os.environ.get("TASK_ID", "0")
        self.tool_name = tool_name

    def emit(self, record):
        if self.sock is None:
            return

        try:
            msg = f"{record.getMessage()}"
            self.sock.send_string(f"{self.task_id}||{self.tool_name}||{msg}")
        except Exception:
            self.handleError(record)


async def zmq_log_listener(bind_addr="tcp://127.0.0.1:6000"):
    ctx = zmq.asyncio.Context()
    sock = ctx.socket(zmq.PULL)

    # Bind to available port
    actual_addr = _bind_zmq_socket(sock, bind_addr)
    set_zmq_address(actual_addr)
    logging.getLogger(__name__).info(f"ZMQ listener bound to: {actual_addr}")

    root_logger = logging.getLogger()

    while True:
        raw = await sock.recv_string()
        if "||" in raw:
            task_id, tool_name, msg = raw.split("||", 2)

            record = root_logger.makeRecord(
                name=f"[TOOL] {tool_name}",
                level=logging.INFO,
                fn="",
                lno=0,
                msg=msg,
                args=(),
                exc_info=None,
            )
            record.task_id = task_id

            root_logger.handle(record)
        else:
            root_logger.info(raw)


def start_zmq_listener():
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.run_until_complete(zmq_log_listener())


def setup_mcp_logging(level="INFO", addr=None, tool_name="unknown_tool"):
    root = logging.getLogger()
    root.setLevel(level)

    # Remove root handlers
    for h in root.handlers[:]:
        root.removeHandler(h)
        h.close()

    # Remove all handlers from fastmcp child loggers
    for _name, logger in logging.Logger.manager.loggerDict.items():
        if isinstance(logger, logging.Logger):
            for h in logger.handlers[:]:
                logger.removeHandler(h)
                h.close()
            logger.propagate = True  # Ensure bubbling to root

    # Re-add the ZMQ handler (will use global address if addr is None)
    handler = ZMQLogHandler(addr=addr, tool_name=tool_name)
    handler.setFormatter(logging.Formatter("[TOOL] %(asctime)s %(levelname)s: %(message)s"))
    root.addHandler(handler)


def setup_log_record_factory():
    old_factory = logging.getLogRecordFactory()

    def record_factory(*args, **kwargs):
        record = old_factory(*args, **kwargs)
        record.task_id = TASK_CONTEXT_VAR.get()
        return record

    logging.setLogRecordFactory(record_factory)


class TaskFilter(logging.Filter):
    def __init__(self, task_id: str):
        super().__init__()
        self.task_id = task_id

    def filter(self, record: logging.LogRecord) -> bool:
        return getattr(record, "task_id", None) == self.task_id


def make_task_logger(task_id: str, log_dir: Path) -> logging.Handler:
    log_dir.mkdir(parents=True, exist_ok=True)
    file_path = log_dir / f"task_{task_id}.log"
    fh = logging.FileHandler(file_path, encoding="utf-8")
    fmt = logging.Formatter("%(asctime)s [%(levelname)s] %(name)s: %(message)s")
    fh.setFormatter(fmt)
    fh.addFilter(TaskFilter(task_id))
    logging.getLogger().addHandler(fh)
    return fh


class JSONFormatter(logging.Formatter):
    """Custom JSON formatter that includes trace_id, chat_id, user_id, and message_id"""

    def format(self, record):
        log_obj = {
            "timestamp": self.formatTime(record, datefmt="%Y-%m-%d %H:%M:%S"),
            "level": record.levelname,
            "logger": record.name,
            "file": record.filename,
            "line": record.lineno,
            "message": record.getMessage(),
            "trace_id": get_trace_id(),
            "chat_id": get_chat_id(),
            "user_id": get_user_id(),
            "message_id": get_message_id(),
        }

        # Add exception info if present
        if record.exc_info:
            log_obj["exception"] = self.formatException(record.exc_info)

        return json.dumps(log_obj, ensure_ascii=False)


def remove_all_console_handlers():
    """
    Remove all console handlers (StreamHandler/RichHandler) from all loggers in the current process.
    """
    for _name, logger in logging.Logger.manager.loggerDict.items():
        if isinstance(logger, logging.Logger):
            handlers_to_remove = []
            for h in logger.handlers:
                if isinstance(h, (logging.StreamHandler, RichHandler)):
                    handlers_to_remove.append(h)
            for h in handlers_to_remove:
                logger.removeHandler(h)
                h.close()

    root_logger = logging.getLogger()
    handlers_to_remove = []
    for h in root_logger.handlers:
        if isinstance(h, logging.StreamHandler):
            handlers_to_remove.append(h)
    for h in handlers_to_remove:
        root_logger.removeHandler(h)
        h.close()


@contextmanager
def task_logging_context(task_id: str, log_dir: Path):
    token = TASK_CONTEXT_VAR.set(task_id)
    handler = make_task_logger(task_id, log_dir / "task_logs")
    try:
        yield
    finally:
        TASK_CONTEXT_VAR.reset(token)
        logging.getLogger().removeHandler(handler)
        handler.close()


def init_logging_for_benchmark_evaluation(print_task_logs=False):
    threading.Thread(target=start_zmq_listener, daemon=True).start()  # monitoring tool logs
    logging.basicConfig(handlers=[])
    setup_log_record_factory()
    if not print_task_logs:
        remove_all_console_handlers()


def bootstrap_logger(
    level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] | int = "INFO",
    logger_name: str = "mem_deep_research",
    logger: logging.Logger | None = None,
    log_dir: str | Path | None = None,  # Log storage directory
    log_filename: str = "mem_deep_research.log",  # Default log filename
    to_console: bool = True,  # Whether to display to console
) -> logging.Logger:
    """Configure only this logger, not the root logger"""
    if logger is None:
        logger = logging.getLogger(logger_name)

    for handler in logger.handlers[:]:
        logger.removeHandler(handler)

    # Check if we should use JSON logging (production mode)
    is_prod = os.getenv("ENV") == "prod"

    if to_console:
        if is_prod:
            # Use JSON formatter in production
            handler = logging.StreamHandler()
            handler.setFormatter(JSONFormatter())
        else:
            # Use Rich handler in development
            handler = RichHandler(
                console=Console(
                    stderr=True,
                    width=200,
                    color_system=None,
                    force_terminal=False,
                    legacy_windows=False,
                ),
                rich_tracebacks=True,
                tracebacks_suppress=[hydra],
                tracebacks_show_locals=True,
                show_level=False,
            )
            formatter = logging.Formatter("[%(levelname)s] %(message)s")
            handler.setFormatter(formatter)
        logger.addHandler(handler)

    if log_dir is not None:
        log_dir = Path(log_dir)
        log_dir.mkdir(parents=True, exist_ok=True)
        file_path = log_dir / log_filename
        file_handler = logging.FileHandler(file_path, encoding="utf-8")
        if is_prod:
            # Use JSON formatter for file logging in production
            file_handler.setFormatter(JSONFormatter())
        else:
            # Use standard formatter in development
            file_handler.setFormatter(
                logging.Formatter("%(asctime)s [%(levelname)s] %(name)s: %(message)s")
            )
        logger.addHandler(file_handler)

    logger.setLevel(level)
    logger.propagate = True

    return logger
