import json
import logging
from logging.handlers import RotatingFileHandler
from pathlib import Path


class Ansi:
    RESET = "\033[0m"
    BOLD = "\033[1m"

    RED = "\033[31m"
    GREEN = "\033[32m"
    YELLOW = "\033[33m"
    CYAN = "\033[36m"
    MAGENTA = "\033[35m"


def setup_logger(log_file: str = "logs/cloudenv_steps.log") -> logging.Logger:
    logger = logging.getLogger("cloudenv.step")
    logger.setLevel(logging.DEBUG)

    if logger.handlers:
        return logger

    Path(log_file).parent.mkdir(parents=True, exist_ok=True)

    console = logging.StreamHandler()

    class ColorFormatter(logging.Formatter):
        def __init__(self):
            super().__init__("%(message)s")

        def format(self, record):
            if not hasattr(record, "episode_id"):
                return super().format(record)

            status = record.__dict__.get("status", "UNKNOWN")

            if status == "SUCCESS":
                status_col = f"{Ansi.GREEN}SUCCESS{Ansi.RESET}"
            elif status == "ACCESS_DENIED":
                status_col = f"{Ansi.YELLOW}ACCESS_DENIED{Ansi.RESET}"
            else:
                status_col = f"{Ansi.RED}ERROR{Ansi.RESET}"

            return (
                f"{Ansi.BOLD}{Ansi.CYAN}[STEP]{Ansi.RESET} "
                f"ep={Ansi.YELLOW}{record.episode_id}{Ansi.RESET} "
                f"sid={Ansi.YELLOW}{record.task_id}{Ansi.RESET} "
                f"step={Ansi.MAGENTA}{record.step}{Ansi.RESET} "
                f"{record.service}.{record.operation} "
                f"=> {status_col}"
            )

    console.setFormatter(ColorFormatter())
    logger.addHandler(console)

    file_handler = RotatingFileHandler(
        log_file,
        # 5MB
        maxBytes=5 * 1024 * 1024,
        backupCount=5,
    )

    class JsonFormatter(logging.Formatter):
        def __init__(self):
            super().__init__("%(message)s")

        def format(self, record):
            payload = {
                "episode_id": record.__dict__.get("episode_id"),
                "task_id": record.__dict__.get("task_id"),
                "step": record.__dict__.get("step"),
                "service": record.__dict__.get("service"),
                "operation": record.__dict__.get("operation"),
                "status": record.__dict__.get("status"),
                "input": record.__dict__.get("input"),
                "output": record.__dict__.get("output"),
            }
            return json.dumps(payload)

    file_handler.setFormatter(JsonFormatter())
    logger.addHandler(file_handler)

    return logger


logger = setup_logger()