import time
import json
import logging
import hashlib
from playwright.sync_api import Browser, BrowserContext, Page, sync_playwright, TimeoutError, Error
import redis as rd
import ddddocr
import re
import cv2
import base64
from io import BytesIO
from PIL import Image
import numpy as np
from func_timeout import func_set_timeout
from func_timeout.exceptions import FunctionTimedOut


# REDIS = {
#     # 'host': '127.0.0.1',
#     'host': '120.79.147.190',
#     'port': 6379,
#     'password': 'Vm5vQH4ydFXh',
#     'db': 10
# }

REDIS = {
    # 'host': '127.0.0.1',
    'host': 'wx.yswg.com.cn',
    'port': 6379,
    'password': 'yswg@2019',
    'db': 1
}


def singleton(cls, *args, **kw):
    """singleton mode.

    :param cls: classname
    :param args: args.
    :param kw: kwargs.
    :return:
    """

    instances = {}

    def _singleton():
        if cls not in instances:
            instances[cls] = cls(*args, **kw)
        return instances[cls]

    return _singleton


def md5(src: str, algorithm: str = "md5", digits: int = 32) -> str:
    """md5 algorithms.

    :param src: original string.
    :param algorithm: algorithm method.
    :param digits: 16 length or 32 length.
    :return: string.
    """

    algorithm = hashlib.new(algorithm)
    algorithm.update(src.encode('utf8'))
    if digits == 16:
        return algorithm.hexdigest()[8:24]
    else:
        return algorithm.hexdigest()

@singleton
class Redis(object):
    def __init__(self):
        self.host = REDIS['host']
        self.port = REDIS['port']
        self.db = REDIS['db']
        self.password = REDIS['password']
        self.pool = rd.ConnectionPool(
            host=self.host,
            port=self.port,
            db=self.db,
            password=self.password,
        )

    def get_instance(self):
        return rd.Redis(connection_pool=self.pool)


@func_set_timeout(10)
def sadd(key, value, use_md5=True):
    """add key-value to the sorted set.

    :param key: key
    :param value: value
    :param use_md5: weather use md5.
    :return: True for done, False for not.
    """

    r = Redis().get_instance()
    if use_md5:
        added = r.sadd(key, md5(value, digits=16))
    else:
        added = r.sadd(key, value)
    return added == 1


def spop(key, count=20) -> list:
    """spop

    :param count:
    :param key:
    :return:
    """

    r = Redis().get_instance()
    return r.spop(key, count)


def ladd(key, value, use_md5=True):
    """add key-value to the sorted set.

    :param key: key
    :param value: value
    :param use_md5: weather use md5.
    :return: True for done, False for not.
    """

    r = Redis().get_instance()
    if use_md5:
        added = r.lpush(key, md5(value, digits=16))
    else:
        added = r.lpush(key, value)
    return added == 1


@func_set_timeout(30)
def listpop(key) -> list:
    """lpop

    :param count:
    :param key:
    :return:
    """

    r = Redis().get_instance()
    return r.lpop(key)


def xadd(key, data):
    r = Redis().get_instance()
    return r.xadd(key, data)


class ChinataxSpider(object):
    seed_key = 'finance:sp_invoice_queue'
    save_key = 'finance:sp_invoice_result'

    def __init__(self):
        logging.basicConfig(format='%(asctime)s %(name)s %(levelname)s %(message)s',
                            level=logging.INFO)
        self.browser: Browser = None
        self.context: BrowserContext = None
        self.page: Page = None
        self.padding_error = 1
        self.seeds = [{'@class': 'com.alibaba.fastjson.JSONObject', 'fpdm': None, 'kprq': '20231229', 'u_key': '1742416930985218050', 'kjje': '3861.62', 'fphm': '23952000000063440070', 'jym': None}]

    def base64_to_image(self, base64_str):
        """
        将base64字符串转换成图片
        :param base64_str:
        :param color:
        :return:
        """
        base64_data = re.sub('^data:image/.+;base64,', '', base64_str)
        byte_data = base64.b64decode(base64_data)
        image_data = BytesIO(byte_data)
        img = Image.open(image_data)
        img = img.convert("RGB")
        return img

    def get_img_base64(self):
        # yzminfo = ""
        color = "black"
        # img_base64 = ""
        # count = 1
        yzminfo = self.page.query_selector("#yzminfo").as_element().text_content()
        self.page.wait_for_timeout(1000)
        img_base64 = self.page.query_selector("#yzm_img").get_attribute("src")
        self.page.wait_for_timeout(1000)

        if "蓝色" in yzminfo:
            color = 'blue'
        if "红色" in yzminfo:
            color = 'red'
        if "黄色" in yzminfo:
            color = 'yellow'

        return color, img_base64

    def get_img(self, count):
        # count = 1
        if self.page.query_selector("#yzm_img").get_attribute("src") == "images/code.png":
            self.page.query_selector("#yzm_img").click()
            self.page.wait_for_timeout(3000)
            self.page.wait_for_load_state()
        elif count != 0:
            self.page.query_selector("#yzm_img").click()
            self.page.wait_for_timeout(3000)
            self.page.wait_for_load_state()
        color, img_base64 = self.get_img_base64()
        # print(f"第{count}次是{color}色")
        # while color != 'black':
        #     self.page.query_selector("#yzm_img").click()
        #     self.page.wait_for_timeout(3000)
        #     self.page.wait_for_load_state()
        #     count = count + 1
        #     color, img_base64 = self.get_img_base64()
        #     color = color
        #     img_base64 = img_base64
        #     print(f"第{count}次是{color}色")
            # self.base64_to_image(img_base64).save(f"test第{count}次.png")
        # image = self.base64_to_image(img_base64)
        # image.save("test.png")
        return color, img_base64

    def ddddocr_imge_get_code(self, color, img_base64):
        ocr = ddddocr.DdddOcr()
        img = self.base64_to_image(img_base64)
        if color == 'black':
            return ocr.classification(img)
        # 转换为cv2可以使用的对象
        image = np.asarray(img)
        one_color_img = self.save_appointed_color(image, color)
        img = Image.fromarray(cv2.cvtColor(one_color_img, cv2.COLOR_BGR2RGB))
        #  todo 预处理
        return ocr.classification(img)

    def save_appointed_color(self, image, color):
        hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        if color == "red":
            lower_red = np.array([0, 100, 100])
            upper_red = np.array([10, 255, 255])
            mask = cv2.inRange(hsv_image, lower_red, upper_red)
            print("保留 red")
        elif color == "blue":
            # 定义要保留的颜色的范围
            lower_blue = np.array([100, 50, 50])
            upper_blue = np.array([140, 255, 255])
            mask = cv2.inRange(hsv_image, lower_blue, upper_blue)
            print("保留 blue")
        elif color == "yellow":
            # 定义要保留的颜色的范围
            lower_yellow = np.array([26, 43, 46])
            upper_yellow = np.array([34, 255, 255])
            mask = cv2.inRange(hsv_image, lower_yellow, upper_yellow)
            print("保留 yellow")
        else:
            # 定义要保留的颜色的范围（这里以红色为例）
            lower_red = np.array([0, 50, 50])
            upper_red = np.array([10, 255, 255])
            # HSV颜色空间中红色的另一个范围
            lower_red2 = np.array([170, 50, 50])
            upper_red2 = np.array([180, 255, 255])
            mask1 = cv2.inRange(hsv_image, lower_red, upper_red)
            mask2 = cv2.inRange(hsv_image, lower_red2, upper_red2)
            mask = mask1 + mask2
        white = np.full_like(image, (255, 255, 255), dtype=np.uint8)
        result = cv2.bitwise_and(white, white, mask=mask)
        return result

    def get_seed(self):
        while True:
            try:
                seeds = listpop(self.seed_key)
                break
            except FunctionTimedOut as e:
                logging.info(f"get seed timeout {e}")
                continue
        # seeds = self.seeds.pop()
        if seeds:
            seed = json.loads(seeds)
            # seed = seeds
            if seed.get("fphm") and seed.get("kprq") and (seed.get("kjje") or seed.get("jym")):
                return seed
            else:
                error_msg = "seed error"
                data = {
                    "u_key": seed.get("u_key"),
                    "dom": error_msg,
                }
                xadd(self.save_key, data)
                logging.info(f"seeds 缺少字段 {seed}")
                return None
        else:
            return None

    def recaptcha(self, count):
        color, img_base64 = self.get_img(count)

        code = self.ddddocr_imge_get_code(color, img_base64)
        if not code:
            logging.info("验证码识别失败")
            return False
        logging.info(f"验证码识别 {code}")
        self.page.locator("#yzm").fill(code)
        self.page.wait_for_timeout(1000)
        # 点击空白页失去焦点
        self.page.locator("#pageshow").click()
        self.page.wait_for_timeout(1000)
        if "display: none" not in self.page.query_selector("#checkfp").get_attribute("style"):
            self.page.locator("#checkfp").click()
            self.page.wait_for_timeout(1000)
            # 判断验证码是否通过
            if not self.page.query_selector("#dialog-body"):
                if error_msg := self.page.query_selector("#popup_message"):
                    if "超过该张发票当日查验次数" in error_msg.text_content():
                        self.page.locator("#popup_ok").click()
                        self.page.wait_for_timeout(1000)
                        return "count error"
                    if "验证码请求次数过于频繁" in error_msg.text_content():
                        self.page.locator("#popup_ok").click()
                        self.page.wait_for_timeout(1000)
                        return "recaptcha count error"
                self.page.locator("#popup_ok").click()
                self.page.wait_for_timeout(1000)
                # 刷新验证码
                # self.page.locator("#yzm_img").click()
                # self.page.wait_for_timeout(1000)
                logging.info("验证码处理错误")
                return False
            else:
                return True
        else:
            return "seed error"

    def get_item(self, seed):
        logging.info(f"fpdm --> {seed.get('fpdm', '')}")
        logging.info(f"seed --> {seed}")
        # fpdm --> None
        if seed.get('fpdm'):
            self.page.locator("#fpdm").fill(seed.get("fpdm", ""))
            self.page.wait_for_timeout(5000)
        self.page.locator("#fphm").fill(seed.get("fphm"))
        self.page.wait_for_timeout(1000)
        self.page.locator("#kprq").fill(seed.get("kprq"))
        self.page.wait_for_timeout(3000)
        if "开具金额" in self.page.query_selector(
                "span[id='context']").text_content() or "价税合计" in self.page.query_selector(
                "span[id='context']").text_content():
            kjje = seed.get("kjje")
        else:

            kjje = seed.get("jym", "")[-6::]
        self.page.locator("#kjje").fill(kjje)

        self.page.locator("#pageshow").click()
        self.page.wait_for_timeout(2000)
        if error_msg := self.page.query_selector("#fpdmjy").text_content().strip():
            data = {
                "u_key": seed.get("u_key"),
                "dom": error_msg,
            }
            xadd(self.save_key, data)
            logging.info(f"{error_msg}")
            self.page.close()
            self.context.close()
        elif error_msg := self.page.query_selector("#fphmjy").text_content().strip():
            data = {
                "u_key": seed.get("u_key"),
                "dom": error_msg,
            }
            xadd(self.save_key, data)
            logging.info(f"{error_msg}")
            self.page.close()
            self.context.close()
        elif error_msg := self.page.query_selector("#kprqjy").text_content().strip():
            data = {
                "u_key": seed.get("u_key"),
                "dom": error_msg,
            }
            xadd(self.save_key, data)
            logging.info(f"{error_msg}")
            self.page.close()
            self.context.close()
        elif error_msg := self.page.query_selector("#kjjejy").text_content().strip():
            data = {
                "u_key": seed.get("u_key"),
                "dom": error_msg,
            }
            xadd(self.save_key, data)
            logging.info(f"{error_msg}")
            self.page.close()
            self.context.close()
        else:
            error = 0
            for i in range(8):
                # 判断验证码是否通过
                if error_msg := self.recaptcha(i):
                    if error_msg in ["count error", "seed error"]:
                        data = {
                            "u_key": seed.get("u_key"),
                            "dom": error_msg,
                        }
                        xadd(self.save_key, data)
                        logging.info(f"{error_msg} {seed}")
                        error = 1
                        self.page.close()
                        self.context.close()
                        break
                    if error_msg == "recaptcha count error":
                        logging.info(f"recaptcha count error {seed}")
                        break
                    self.page.frame_locator("#dialog-body")
                    self.page.wait_for_timeout(2000)
                    text_contents = self.page.frame("dialog-body").content()
                    data = {
                        "u_key": seed.get("u_key"),
                        "dom": text_contents,
                    }
                    xadd(self.save_key, data)
                    logging.info(f"seed spider ok --> {seed}")
                    error = 1
                    self.page.close()
                    self.context.close()
                    break
                else:
                    logging.info("验证码处理失败重试")

            if error == 0:
                logging.info("失败8次处理,将任务重新推送到redis")
                ladd(self.seed_key, json.dumps(seed), use_md5=False)
                self.page.close()
                self.context.close()

    def crawl(self, url, seed):
        # 需要打开的网站
        self.page.goto(url)
        self.page.wait_for_timeout(1000)
        self.get_item(seed)

    def change_user(self):
        user_agent = 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_14_5) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/96.0.4664.45 Safari/537.36'
        self.context = self.browser.new_context(
            locale='en-GB',
            bypass_csp=True,
            user_agent=user_agent,
            ignore_https_errors=True,
            no_viewport=True,
        )
        # self.page = self.browser.new_page()
        self.page = self.context.new_page()

        # 模拟浏览器参数
        self.page.locator("body").click()
        js = """
        Object.defineProperties(navigator, {webdriver:{get:()=>undefined}});
        """
        self.page.add_init_script(js)
        self.page.set_default_timeout(300000)
        # self.page.add_init_script(
        #     "const newProto = navigator.__proto__; delete newProto.webdriver; navigator.__proto__ = newProto;")

    def run(self):
        while True:
            seed = self.get_seed()
            try:
                if seed:
                    self.change_user()
                    url = "https://inv-veri.chinatax.gov.cn/index.html"
                    self.crawl(url, seed)
                else:
                    time.sleep(30)
                    logging.info('no task sleep 30s')
            except Error as e:
                logging.info(f"--> playwright error ")
                self.page.close()
                self.context.close()
                ladd(self.seed_key, json.dumps(seed), use_md5=False)
                continue
            except FunctionTimedOut as e:
                logging.info(f"--> FunctionTimedOut ")
                self.page.close()
                self.context.close()
                ladd(self.seed_key, json.dumps(seed), use_md5=False)
                continue
            except Exception as e:
                self.page.close()
                self.context.close()
                if f"{e}" == "Incorrect padding":
                    ladd(self.seed_key, json.dumps(seed), use_md5=False)
                    logging.info(f"--> Incorrect padding error {e}")
                    self.padding_error += 1
                    if self.padding_error >= 5:
                        time.sleep(1200)
                        logging.info(f"等待1200分钟")
                        self.padding_error = 1
                    continue
                else:
                    data = {
                        "u_key": seed.get("u_key"),
                        "dom": f"code error",
                    }
                    logging.info(f"Incorrect padding error {e}")
                    xadd(self.save_key, data)
                    continue

    def send_msg(self, account, title, content):
        import requests
        url = 'http://47.112.96.71:8082/selection/sendMessage'
        data = {
            'account': account,
            'title': title,
            'content': content
        }
        print(data)
        requests.post(url=url, data=data, timeout=15)

    def main(self):
        headless = False
        # headless = True
        logging.info(f"{headless}")
        with sync_playwright() as _playwright:
            # self.browser = _playwright.chromium.launch_persistent_context(
            #     headless=False,
            #     executable_path="C:\Program Files\Google\Chrome\Application\chrome.exe",
            #     user_data_dir=r"D:\chrome_data\test01",
            #     ignore_https_errors=True,
            #     no_viewport=True,
            #     bypass_csp=True
            # )
            self.browser = _playwright.chromium.launch(
                headless=False,
                executable_path="C:\Program Files\Google\Chrome\Application\chrome.exe",
            )
            try:
                self.run()
            except Exception as e:
                print(e)
                account = 'hezhe'
                title = f'税务爬取脚本 error'
                content = f"税务爬取脚本 error 时间：{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}"
                self.send_msg(account, title, content)


if __name__ == '__main__':
    chinatax = ChinataxSpider()
    chinatax.main()
    # seed = {
    #     "u_key": "xxx",
    #     "fpdm": "044002311111",
    #     "fphm": "59210491",
    #     "kprq": "20231214",
    #     "kjje": "303791"
    # }
    # ladd('finance:sp_invoice_queue', json.dumps(seed), use_md5=False)

