from __future__ import annotations

import logging
from pathlib import Path

import structlog
from rich.logging import RichHandler
from rich.progress import (
    BarColumn,
    Progress,
    SpinnerColumn,
    TaskProgressColumn,
    TextColumn,
    TimeElapsedColumn,
)


def setup_logging(level: str = "INFO", log_file: Path | None = None) -> None:
    """Configure structlog with rich console output and optional file logging."""
    log_level = getattr(logging, level.upper(), logging.INFO)

    shared_processors = [
        structlog.contextvars.merge_contextvars,
        structlog.stdlib.add_logger_name,
        structlog.stdlib.add_log_level,
        structlog.processors.TimeStamper(fmt="ISO"),
        structlog.processors.StackInfoRenderer(),
        structlog.dev.set_exc_info,
    ]

    structlog.configure(
        processors=shared_processors
        + [
            structlog.stdlib.ProcessorFormatter.wrap_for_formatter,
        ],
        logger_factory=structlog.stdlib.LoggerFactory(),
        wrapper_class=structlog.stdlib.BoundLogger,
        context_class=dict,
    )

    formatter = structlog.stdlib.ProcessorFormatter(
        foreign_pre_chain=shared_processors,
        processors=[
            structlog.stdlib.ProcessorFormatter.remove_processors_meta,
            structlog.dev.ConsoleRenderer(colors=True),
        ],
    )

    console_handler = RichHandler(
        rich_tracebacks=True,
        tracebacks_show_locals=False,
        markup=True,
        show_time=False,
        show_path=False,
    )
    console_handler.setFormatter(formatter)

    handlers: list[logging.Handler] = [console_handler]

    if log_file:
        log_file.parent.mkdir(parents=True, exist_ok=True)
        file_handler = logging.FileHandler(log_file)
        file_formatter = structlog.stdlib.ProcessorFormatter(
            foreign_pre_chain=shared_processors,
            processors=[
                structlog.stdlib.ProcessorFormatter.remove_processors_meta,
                structlog.processors.JSONRenderer(),
            ],
        )
        file_handler.setFormatter(file_formatter)
        handlers.append(file_handler)

    root_logger = logging.getLogger()
    root_logger.handlers = []
    root_logger.setLevel(log_level)
    for handler in handlers:
        root_logger.addHandler(handler)

    logging.getLogger("PIL").setLevel(logging.WARNING)
    logging.getLogger("paddle").setLevel(logging.WARNING)
    logging.getLogger("torch").setLevel(logging.WARNING)
    logging.getLogger("transformers").setLevel(logging.WARNING)
    logging.getLogger("watchdog").setLevel(logging.WARNING)


def get_logger(name: str) -> structlog.BoundLogger:
    return structlog.get_logger(name)


def create_progress() -> Progress:
    return Progress(
        SpinnerColumn(),
        TextColumn("[bold blue]{task.description}"),
        BarColumn(),
        TaskProgressColumn(),
        TimeElapsedColumn(),
        transient=True,
    )


class ProgressTracker:
    def __init__(self, total: int, description: str = "Processing"):
        self.total = total
        self.description = description
        self.progress: Progress | None = None
        self.task_id: int | None = None
        self.completed = 0
        self.failed = 0
        self.logger = get_logger(__name__)

    def __enter__(self) -> ProgressTracker:
        self.progress = create_progress()
        self.progress.start()
        self.task_id = self.progress.add_task(self.description, total=self.total)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        if self.progress:
            self.progress.stop()

    def update(self, advance: int = 1, success: bool = True) -> None:
        if self.progress and self.task_id is not None:
            self.progress.update(self.task_id, advance=advance)
        self.completed += 1
        if not success:
            self.failed += 1

    def log_result(self, path: Path, success: bool, error: str | None = None) -> None:
        if success:
            self.logger.info("processed", path=str(path))
        else:
            self.logger.error("failed", path=str(path), error=error)
