from DrissionPage._base.chromium import Chromium
from DrissionPage._configs.chromium_options import ChromiumOptions
from DrissionPage._pages.mix_tab import MixTab
from applog import logger
import time
import json
import os

class Lmk_BI(object):
    tab: MixTab

    def __init__(self):
        self.url = 'https://lmkbi.95155.com/bi-system/#/login'
        self.url_info = 'lmkbi.95155.com'
        self.browser_path = r'C:\Program Files (x86)\Microsoft\Edge\Application\msedge.exe'
        self.cookies_file = 'lmk_cookies.json'  # 保存cookies的文件
        self.co = ChromiumOptions().set_browser_path(self.browser_path).set_local_port(9555)
        # self.co.headless(False)  # 设为False以便查看操作过程
        # self.co.ignore_certificate_errors(True)
        
        self.browser = Chromium(self.co)
        self.tab = self.browser.latest_tab
        self.tab.set.window.max()
        self.tab.set.cookies.clear

    def quit(self):
        self.browser.quit()
        logger.info("浏览器已退出")

    def load_cookies_from_file(self):
        """从文件加载cookies"""
        try:
            if os.path.exists(self.cookies_file):
                with open(self.cookies_file, 'r', encoding='utf-8') as f:
                    cookies = json.load(f)
                logger.info(f"从文件 {self.cookies_file} 加载cookies成功")
                return cookies
            else:
                logger.warning(f"cookies文件 {self.cookies_file} 不存在")
                return None
        except Exception as e:
            logger.error(f"加载cookies文件失败: {str(e)}")
            return None

    def save_cookies_to_file(self, cookies):
        """保存cookies到文件"""
        try:
            with open(self.cookies_file, 'w', encoding='utf-8') as f:
                json.dump(cookies, f, ensure_ascii=False, indent=2)
            logger.info(f"cookies已保存到文件 {self.cookies_file}")
        except Exception as e:
            logger.error(f"保存cookies到文件失败: {str(e)}")

    def init_cookies(self, cookies):
        """设置传入的cookies到浏览器"""
        logger.info('使用传入的cookies登录')
        
        # 处理字符串类型的cookies
        if isinstance(cookies, str):
            try:
                parsed_cookies = json.loads(cookies)
                logger.info(f"成功将字符串cookies解析为JSON对象，共 {len(parsed_cookies)} 条")
                return self.init_cookies(parsed_cookies)
            except json.JSONDecodeError as e:
                logger.error(f'字符串cookies解析失败: {str(e)}')
                return False
        
        # 处理字典类型的cookies
        elif isinstance(cookies, dict):
            try:
                self.browser.set.cookies(cookies)
                logger.info('成功设置字典类型的cookies')
                return True
            except Exception as e:
                logger.error(f'设置字典类型的cookies失败: {str(e)}')
                return False
        
        # 处理列表类型的cookies
        elif isinstance(cookies, list):
            try:
                for cookie in cookies:
                    self.browser.set.cookies(cookie)
                logger.info('成功设置列表类型的cookies')
                return True
            except Exception as e:
                logger.error(f'设置列表类型的cookies失败: {str(e)}')
                return False
        
        # 处理未知类型
        else:
            logger.error(f'传入的cookies格式不正确，类型为: {type(cookies)}，应为列表、字典或JSON字符串')
            return False

    def save_cookies(self):
        """获取并返回当前浏览器tab的cookies"""
        try:
            self.tab.refresh()
            cookies = self.tab.cookies()
            self.save_cookies_to_file(cookies)
            return cookies
        except Exception as e:
            logger.error(f"获取cookies失败: {str(e)}")
            return None

    def is_login(self):
        """判断是否已登录"""
        logger.info("判断是否登录")
        try:
            # 刷新页面以确保获取最新状态
            self.tab.refresh()
            time.sleep(3)
            verify_code_button = self.tab.eles('@@tag()=button@@text()=获取验证码')
            if verify_code_button:
                logger.info("发现'获取验证码'按钮，视为未登录")
                return False
            else:
                logger.info("未发现'获取验证码'按钮，视为已登录")
                return True            
        except Exception as e:
            logger.error(f"判断登录状态时出错: {str(e)}")
            return False

    def open_web_with_cookies(self):
        """使用保存的cookies打开网站"""
        try:
            # 加载cookies
            cookies = self.load_cookies_from_file()
            self.init_cookies(cookies)           
            # 打开网站
            logger.info(f"加载 cookies 成功")
        except Exception as e:
            logger.error(f"加载 cookies 失败: {str(e)}")
        self.tab.get(self.url)
        time.sleep(5)  # 等待页面加载
        return self.is_login()

    def first_time_login(self):
        """首次登录，需要手动操作保存cookies"""
        try:
            logger.info("请手动登录网站...")
            self.tab.get(self.url)
            
            # 等待用户手动登录
            logger.info("请在浏览器中完成登录操作，然后按Enter键继续...")
            input("登录完成后按Enter键继续...")
            
            # 检查是否登录成功
            if self.is_login():
                # 保存cookies
                self.save_cookies()
                logger.info("首次登录成功，cookies已保存")
                return True
            else:
                logger.warning("登录失败，请重试")
                return False
        except Exception as e:
            logger.error(f"首次登录时出错: {str(e)}")
            return False

    def run(self):
        """运行主程序"""
        try:
            # 尝试使用cookies登录
            if not self.open_web_with_cookies():
                # 如果cookies登录失败，提示用户手动登录
                logger.info("cookies登录失败，准备手动登录")
                self.first_time_login()
            
            # 保持浏览器打开一段时间
            logger.info("操作完成，浏览器将保持打开状态")
            time.sleep(30)
        except Exception as e:
            logger.error(f"运行时出错: {str(e)}")
        finally:
            # 关闭浏览器
            self.save_cookies()
            self.quit()

if __name__ == "__main__":
    # 创建实例并运行
    lmk_bi = Lmk_BI()
    lmk_bi.run()