import json
import os
import redis
from kafka import KafkaProducer
from loguru import logger
from redis.cluster import ClusterNode, RedisCluster
from wake_v2.models import *
from wake_v2 import LoggerCollector
from wake_v2 import LoggerInterface
from wake_v2 import LoggerItem
from wake_v2 import LoggerStrategy
import pytz
from wake_v2.logger_conf import ip_address
from wake_v2.logger_conf import LOG_VERSION, LOG_ENABLE, LOGGER_PATH, LOG_TYPE, REDIS_CONF, WAKE_KAFKA_BOOTSTRAP, WAKE_TOPIC, WAKE_GROUP_ID, WAKE_TOPIC_PASSWORD, WAKE_TOPIC_USERNAME
import pathlib
from requests.exceptions import ProxyError
from requests.exceptions import ReadTimeout
import traceback

east_8 = pytz.timezone('Asia/Shanghai')
utc_tz = pytz.timezone('UTC')

LOG_CONF = {
    "level": "DEBUG",
    "dir": LOGGER_PATH,
    "rotation": "00:00",
    "retention": "1 week",
}
WORKER = {
    "id": ip_address if not os.environ.get("MY_POD_NAME") else os.environ.get("MY_POD_NAME"),
    "host": ip_address
}

KAFKA_CONF = {
    'bootstrap.servers': WAKE_KAFKA_BOOTSTRAP,  # Kafka服务器地址
    "batch.size": 16384,  # 控制批处理大小
    "linger.ms": 5,  # 控制批处理等待时间
    "compression.type": "snappy",  # 启用压缩，减少网络传输和存储需求
    "retries": 2,  # 重试次数
    "retry.backoff.ms": 100,  # 重试间隔
    'max_block_ms': 120000,
}

logger_format = (
    "{time:YYYY-MM-DD HH:mm:ss} | "  # 颜色>时间
    "{process.name}-{thread.name} | "  # 进程id-线程名
    "{module} {function}"  # 模块名.方法名
    ":{line} | "  # 行号
    "{level}: "  # 等级
    "{message}"
)  # 日志内容

# 是否为调试模式
log_debug: bool = LOG_CONF["level"] == "DEBUG"


# 定义一个函数，该函数将在线程中运行并发送消息
# def produce_messages(topic, message):
#     with producer_lock:  # 使用锁来同步对 producer 的访问
#         producer.send(topic, message)
#         producer.flush()


def configure_logger():
    """
    配置logger
    """
    log_dir = pathlib.Path(LOG_CONF["dir"])
    log_dir.mkdir(exist_ok=True)
    log_file = log_dir.joinpath("sending_log.log")

    # 移除控制台输出
    logger.remove(handler_id=None)
    # 写到文件
    logger.add(
        log_file,
        format=logger_format,
        level=LOG_CONF["level"],
        rotation=LOG_CONF["rotation"],
        retention=LOG_CONF["retention"],
        encoding="utf-8",
    )


configure_logger()


class QuickCrawlerLogItem(LoggerItem):
    def __init__(self,
                 sid: str = "",
                 taskid: str = "",
                 appid: str = "",
                 sub: str = "",
                 url: str = "",
                 typo: str = "",
                 domain: str = "",
                 action: str = "",
                 method: str = "",
                 status: str = "",
                 timestamps=None,
                 logs=None,
                 tags=None,
                 retry: int = 0,
                 proxy=None,
                 response=None
                 ):
        self.sid = sid
        self.taskid = taskid
        self.appid = appid
        self.sub = sub
        self.url = url
        self.domain = domain
        self.retry = retry
        if proxy is None:
            proxy = {}
        if response is None:
            response = {}
        self.response = response
        self.proxy = proxy
        self.worker = WORKER
        self.version = LOG_VERSION
        self.type = typo
        super().__init__(action, method, status, timestamps, logs, tags)

    @classmethod
    def from_dict(cls, data):
        return cls(
            sid=data.get("sid", ""),
            appid=data.get("appid", "dataapi"),
            url=data.get("url", ""),
            typo=data.get("type", ""),
            domain=data.get("domain", ""),
            action=data.get("action", ""),
            method=data.get("method", ""),
            status=data.get("status", ""),
            timestamps=data.get("timestamps", {}),
            logs=data.get("logs", []),
            tags=data.get("tags", {}),
            retries=data.get("retries", 0),
            resp=data.get("resp", {})
        )


class QuickLoggerInterface(LoggerInterface):
    def __init__(self):
        super().__init__()

    def info(self, message):
        _message = str(message)
        _message = _message.replace("\n", "\\n")
        self.item.log(_message)


def delivery_report(err, msg):
    if err is not None:
        logger.error("Message delivered error {}".format(err))
    else:
        logger.debug('Message delivered to {} [{}]'.format(msg.topic(), msg.partition()))


def init_redis(redis_conf: dict):
    hosts = redis_conf.get("hosts", ["127.0.0.1"])
    ports = redis_conf.get("ports", [6379])
    db = redis_conf.get("db", "db1")
    pwd = redis_conf.get("pwd", "")
    if redis_conf["cluster"]:
        startup_nodes = [ClusterNode(host, port) for port in ports for host in hosts]
        client = RedisCluster(
            startup_nodes=startup_nodes, password=pwd, decode_responses=True
        )
        return client
    else:
        client = redis.Redis(host=hosts[0], port=ports[0], db=db, password=pwd, decode_responses=True)
        return client


class KafkaLoggerCollector(LoggerCollector):
    def __init__(self):
        super().__init__()
        # producer = KafkaProducer(bootstrap_servers=KAFKA_BOOTSTRAP,
        #                          compression_type='snappy')

        producer = KafkaProducer(
            bootstrap_servers=json.loads(WAKE_KAFKA_BOOTSTRAP),
            security_protocol="SASL_PLAINTEXT",
            sasl_mechanism="SCRAM-SHA-256",
            sasl_plain_username=WAKE_TOPIC_USERNAME,
            sasl_plain_password=WAKE_TOPIC_PASSWORD,

            acks="all",

            value_serializer=lambda v: v.encode("utf-8") if isinstance(v, str) else v,
        )

        self.producer = producer

    def flush(self, item: LoggerItem):
        logger.info(LOG_ENABLE)
        if LOG_ENABLE:
            logger.info("flushing log")
            try:
                # producer.produce(LOG_CONSUMER_KEY, json.dumps(item.to_dict(), ensure_ascii=False),
                #                  callback=delivery_report)
                self.producer.send(WAKE_TOPIC,
                                   value=json.dumps(item.to_dict(), ensure_ascii=False))
                self.producer.flush(timeout=3)
            except Exception as dil_error:
                traceback.format_exc()
                logger.error(f"{dil_error}{traceback.format_exc()}")
                logger.error(json.dumps(item.to_dict()))
            finally:
                logger.info("flushed log")
        else:
            logger.info(json.dumps(item.to_dict(), ensure_ascii=False))


class RedisLoggerCollector(LoggerCollector):
    def __init__(self):
        super().__init__()
        self.redis_log_client = init_redis(REDIS_CONF)

    def flush(self, item: LoggerItem):
        if LOG_ENABLE:
            logger.info("flushing log")
            self.redis_log_client.rpush(WAKE_TOPIC, json.dumps(item.to_dict(), ensure_ascii=False))
            logger.info("flushed log")
        else:
            logger.info(json.dumps(item.to_dict(), ensure_ascii=False))


class LoggerQuickCrawlStrategy(LoggerStrategy):
    def __init__(self, interface: LoggerInterface, collector: LoggerCollector, clazz):
        super().__init__(interface, collector, clazz)
        self.action = None
        self.command = None
        self.domain = None
        self.type = None
        self.resp = None

    def set_base_info(self, action: CrawlerAction, command: CrawlerCommand):
        self.action = action.value
        self.command = command.value

    def set_domain(self, domain):
        self.domain = domain

    def set_type(self, typo: CrawlerType):
        self.type = typo.value

    def set_resp(self, resp):
        self.resp = Resp.from_dict(resp)

    def set_tag(self, key, val):
        flush_item = self.interface.get_item()
        if flush_item:
            try:
                flush_item.tag(key, val)
            except Exception as e:
                pass

    def set_tag_failed_block(self):
        flush_item = self.interface.get_item()
        if flush_item:
            try:
                flush_item.tag("failed.reason", "block")
            except Exception as e:
                flush_item.tag("failed.reason", "others")

    def set_tag_failed_reason(self, e: Exception):
        flush_item = self.interface.get_item()
        if flush_item:
            try:
                str_info = str(traceback.format_exc())
                if isinstance(e, ProxyError):
                    flush_item.tag("failed.reason", "proxy")
                elif isinstance(e, ReadTimeout):
                    flush_item.tag("failed.reason", "timeout")
                elif (str_info.__contains__("ERR_PROXY_CONNECTION_FAILED")
                      or str_info.__contains__("PROXY")
                      or str_info.__contains__("proxy")):
                    flush_item.tag("failed.reason", "proxy")
                else:
                    flush_item.tag("failed.reason", "others")
            except Exception as e:
                flush_item.tag("failed.reason", "others")

    def reset(self, item: dict):
        try:
            item = self.clazz.from_dict(item)
            item.action = self.action
            self.interface.reset(item)
        except Exception as e:
            logger.error(e)

    '''
    1. 取出需要投递的信息
    2. 投递消息
    3. 清除消息
    '''

    def flush(self, status: CrawlerStatus = CrawlerStatus.SUCCESS):
        try:
            flush_item = self.interface.get_item()
            if flush_item is not None:
                flush_item.set_status(status.value)
                flush_item.command = self.command
                flush_item.domain = self.domain
                flush_item.type = self.type
                flush_item.resp = self.resp
                self.collector.flush(flush_item)
            self.interface.clean()
        except Exception as e:
            logger.error("push error")

    def info(self, message):
        self.interface.info(message)

    def times(self, times):
        self.interface.times(times)


def build_log():
    if LOG_TYPE == 'REDIS':
        collector = RedisLoggerCollector()
    else:
        collector: LoggerCollector = KafkaLoggerCollector()
    log = LoggerQuickCrawlStrategy(QuickLoggerInterface(), collector, QuickCrawlerLogItem)
    return log
