# -*- coding: utf-8 -*-

from __future__ import annotations

import logging
import threading

from pygelf import GelfTcpHandler, GelfTlsHandler, GelfUdpHandler

from parsedmarc import (
    parsed_aggregate_reports_to_csv_rows,
    parsed_failure_reports_to_csv_rows,
    parsed_smtp_tls_reports_to_csv_rows,
)
from parsedmarc.types import AggregateReport, FailureReport, SMTPTLSReport

log_context_data = threading.local()


class ContextFilter(logging.Filter):
    def filter(self, record):
        record.parsedmarc = log_context_data.parsedmarc
        return True


class GelfClient(object):
    """A client for the Graylog Extended Log Format"""

    def __init__(self, host, port, mode):
        """
        Initializes the GelfClient
        Args:
            host (str): The GELF host
            port (int): The GELF port
            mode (str): The GELF transport mode
        """
        self.host = host
        self.port = port
        self.logger = logging.getLogger("parsedmarc_gelf")
        self.logger.setLevel(logging.INFO)
        self.logger.addFilter(ContextFilter())
        self.gelf_mode = {
            "udp": GelfUdpHandler,
            "tcp": GelfTcpHandler,
            "tls": GelfTlsHandler,
        }
        self.handler = self.gelf_mode[mode](
            host=self.host, port=self.port, include_extra_fields=True
        )
        self.logger.addHandler(self.handler)

    def save_aggregate_report_to_gelf(
        self, aggregate_reports: AggregateReport | list[AggregateReport]
    ):
        rows = parsed_aggregate_reports_to_csv_rows(aggregate_reports)
        for row in rows:
            log_context_data.parsedmarc = row
            self.logger.info("parsedmarc aggregate report")

        log_context_data.parsedmarc = None

    def save_failure_report_to_gelf(
        self, failure_reports: FailureReport | list[FailureReport]
    ):
        rows = parsed_failure_reports_to_csv_rows(failure_reports)
        for row in rows:
            log_context_data.parsedmarc = row
            self.logger.info("parsedmarc failure report")

    def save_smtp_tls_report_to_gelf(
        self, smtp_tls_reports: SMTPTLSReport | list[SMTPTLSReport]
    ):
        rows = parsed_smtp_tls_reports_to_csv_rows(smtp_tls_reports)
        for row in rows:
            log_context_data.parsedmarc = row
            self.logger.info("parsedmarc smtptls report")

    def close(self):
        """Remove and close the GELF handler, releasing its connection."""
        self.logger.removeHandler(self.handler)
        self.handler.close()

    # Backward-compatible alias
    save_forensic_report_to_gelf = save_failure_report_to_gelf
