"""
MoviePilot Cookie功能示例程序
从CookieCloud下载数据并注入浏览器访问网站
"""

import json
import time
import hashlib
import base64
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass
import urllib.request
import urllib.error
from Crypto.Cipher import AES


@dataclass
class CookieCloudConfig:
    """CookieCloud配置"""
    server: str
    username: str
    password: str
    timeout: int = 30


class CookieCloudDownloader:
    """从CookieCloud服务器下载cookie数据"""
    
    def __init__(self, config: CookieCloudConfig):
        self.config = config
    
    def _get_crypt_key(self) -> bytes:
        combined_string = f"{self.config.username}-{self.config.password}"
        return hashlib.md5(combined_string.encode('utf-8')).hexdigest()[:16].encode("utf-8")
    
    def _bytes_to_key(self, data: bytes, salt: bytes, output: int = 48) -> bytes:
        assert len(salt) == 8, len(salt)
        data += salt
        key = hashlib.md5(data).digest()
        final_key = key
        while len(final_key) < output:
            key = hashlib.md5(key + data).digest()
            final_key += key
        return final_key[:output]
    
    def _decrypt(self, encrypted: str, passphrase: bytes) -> bytes:
        encrypted_bytes = base64.b64decode(encrypted)
        assert encrypted_bytes.startswith(b"Salted__"), "Invalid encrypted data format"
        salt = encrypted_bytes[8:16]
        key_iv = self._bytes_to_key(passphrase, salt, 32 + 16)
        key = key_iv[:32]
        iv = key_iv[32:]
        aes = AES.new(key, AES.MODE_CBC, iv)
        decrypted_padded = aes.decrypt(encrypted_bytes[16:])
        padding_length = decrypted_padded[-1]
        if isinstance(padding_length, str):
            padding_length = ord(padding_length)
        return decrypted_padded[:-padding_length]
    
    def _get_url_domain(self, domain: str) -> str:
        if not domain:
            return ""
        domain = domain.lstrip('.')
        if ":" in domain:
            domain = domain.split(":")[0]
        parts = domain.split(".")
        if all(part.isdigit() for part in parts):
            return domain
        if len(parts) >= 2:
            return ".".join(parts[-2:])
        return domain
    
    def download_all(self) -> Tuple[Optional[Dict], str]:
        """下载所有cookie和local storage数据"""
        try:
            url = f"{self.config.server}/get/{self.config.username}"
            request = urllib.request.Request(
                url,
                headers={
                    'Content-Type': 'application/json',
                    'User-Agent': 'MoviePilot-Cookie-Client/1.0'
                },
                method='GET'
            )
            
            response = urllib.request.urlopen(request, timeout=self.config.timeout)
            
            if response.status != 200:
                return None, f"服务器返回错误状态码: {response.status}"
            
            result = json.loads(response.read().decode('utf-8'))
            
            if not result:
                return None, "服务器返回数据为空"
            
            encrypted = result.get("encrypted")
            if not encrypted:
                return None, "未获取到cookie密文"
            
            crypt_key = self._get_crypt_key()
            try:
                decrypted_data = self._decrypt(encrypted, crypt_key)
                result = json.loads(decrypted_data.decode("utf-8"))
            except Exception as e:
                return None, f"cookie解密失败: {str(e)}"
            
            if not result:
                return None, "cookie解密为空"
            
            cookie_data = result.get("cookie_data", {})
            local_storage_data = result.get("local_storage_data", {})
            
            total_cookies = sum(len(v) for v in cookie_data.values())
            total_ls = sum(len(v) for v in local_storage_data.values())
            print(f"  下载完成: {len(cookie_data)} 个域名的Cookie({total_cookies}个), {len(local_storage_data)} 个域名的Local Storage({total_ls}个)")
            
            return {"cookies": cookie_data, "local_storage": local_storage_data}, ""
            
        except urllib.error.HTTPError as e:
            if e.code == 401:
                return None, "认证失败，请检查用户名和密码"
            else:
                return None, f"HTTP错误: {e.code} {e.reason}"
        except urllib.error.URLError as e:
            return None, f"网络连接失败: {e.reason}"
        except Exception as e:
            return None, f"下载失败: {str(e)}"


class BrowserController:
    """使用Playwright控制浏览器"""
    
    def __init__(self, headless: bool = False, use_edge: bool = False):
        self.headless = headless
        self.use_edge = use_edge
        self.browser = None
        self.context = None
        self.page = None
        self.playwright = None
    
    def start(self) -> bool:
        try:
            from playwright.sync_api import sync_playwright
            
            self.playwright = sync_playwright().start()
            
            if self.use_edge:
                self.browser = self.playwright.chromium.launch(
                    channel="msedge",
                    headless=self.headless,
                    args=['--disable-blink-features=AutomationControlled', '--disable-dev-shm-usage', '--no-sandbox']
                )
            else:
                self.browser = self.playwright.chromium.launch(
                    headless=self.headless,
                    args=['--disable-blink-features=AutomationControlled', '--disable-dev-shm-usage', '--no-sandbox']
                )
            
            self.context = self.browser.new_context(
                user_agent='Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'
            )
            
            self.page = self.context.new_page()
            
            print(f"✓ 浏览器启动成功 ({'Edge' if self.use_edge else 'Chromium'})")
            return True
            
        except Exception as e:
            print(f"✗ 浏览器启动失败: {str(e)}")
            return False
    
    def inject_cookies(self, cookies: Dict):
        """注入cookie到浏览器"""
        if not self.context or not cookies:
            return
        
        try:
            playwright_cookies = []
            for domain_key, cookie_list in cookies.items():
                for cookie in cookie_list:
                    cookie_domain = cookie.get("domain", "").lstrip('.')
                    if not cookie_domain:
                        cookie_domain = domain_key.lstrip('.')
                    
                    pw_cookie = {
                        "name": cookie.get("name", ""),
                        "value": cookie.get("value", ""),
                        "domain": cookie_domain,
                        "path": cookie.get("path", "/"),
                    }
                    if cookie.get("secure"):
                        pw_cookie["secure"] = True
                    if cookie.get("httpOnly"):
                        pw_cookie["httpOnly"] = True
                    if cookie.get("expirationDate"):
                        pw_cookie["expires"] = int(cookie.get("expirationDate"))
                    playwright_cookies.append(pw_cookie)
            
            self.context.add_cookies(playwright_cookies)
            print(f"✓ Cookie注入成功 ({len(playwright_cookies)} 个)")
        except Exception as e:
            print(f"✗ Cookie注入失败: {str(e)}")
    
    def navigate(self, url: str, wait_time: int = 5, timeout: int = 60) -> bool:
        if not self.page:
            return False
        
        try:
            print(f"正在访问: {url}")
            
            self.page.goto(url, timeout=timeout * 1000)
            
            try:
                self.page.wait_for_load_state("networkidle", timeout=timeout * 1000)
            except Exception:
                print("  networkidle超时，尝试load状态...")
                try:
                    self.page.wait_for_load_state("load", timeout=30000)
                except Exception:
                    print("  load超时，尝试domcontentloaded状态...")
                    self.page.wait_for_load_state("domcontentloaded", timeout=10000)
            
            time.sleep(wait_time)
            
            print(f"✓ 页面加载完成")
            return True
            
        except Exception as e:
            print(f"✗ 页面访问失败: {str(e)}")
            return False
    
    def check_login_status(self, login_indicator: str) -> bool:
        if not self.page:
            return False
        
        try:
            html_content = self.page.content()
            
            if login_indicator:
                if login_indicator in html_content:
                    print(f"✓ 检测到登录指示器: {login_indicator}")
                    return True
                else:
                    print(f"✗ 未检测到登录指示器: {login_indicator}")
                    return False
            
            return False
            
        except Exception as e:
            print(f"✗ 登录状态检查失败: {str(e)}")
            return False
    
    def take_screenshot(self, filepath: str):
        if not self.page:
            return
        
        try:
            self.page.screenshot(path=filepath)
            print(f"✓ 截图已保存: {filepath}")
        except Exception as e:
            print(f"✗ 截图失败: {str(e)}")
    
    def close(self):
        try:
            if self.browser:
                self.browser.close()
            
            if self.playwright:
                self.playwright.stop()
            
            print("✓ 浏览器已关闭")
            
        except Exception as e:
            print(f"✗ 关闭浏览器失败: {str(e)}")


class CookieDemo:
    """Cookie功能示例"""
    
    def __init__(self, cookiecloud_config: CookieCloudConfig):
        self.downloader = CookieCloudDownloader(cookiecloud_config)
        self.browser = None
    
    def run_sites(self, sites: list, headless: bool = False, use_edge: bool = False):
        """下载一次数据，依次访问多个网站"""
        # 下载数据
        print("=" * 60)
        print("从CookieCloud下载数据")
        print("=" * 60)
        
        data, error = self.downloader.download_all()
        
        if error:
            print(f"✗ 数据下载失败: {error}")
            return
        
        # 启动浏览器
        print(f"\n启动浏览器")
        print("-" * 60)
        
        self.browser = BrowserController(headless=headless, use_edge=use_edge)
        
        if not self.browser.start():
            return
        
        # 注入所有cookie
        if data and data.get("cookies"):
            print(f"\n注入所有cookie到浏览器")
            print("-" * 60)
            self.browser.inject_cookies(data["cookies"])
        
        # 依次访问每个网站
        for idx, site in enumerate(sites, 1):
            print(f"\n\n{'=' * 60}")
            print(f"访问网站 {idx}/{len(sites)}: {site['name']}")
            print(f"URL: {site['url']}")
            print("=" * 60)
            
            if not self.browser.navigate(site['url']):
                print(f"✗ 页面访问失败")
                continue
            
            # 验证登录状态
            from urllib.parse import urlparse
            parsed = urlparse(site['url'])
            target_domain = parsed.netloc
            
            login_indicator = site.get('login_indicator')
            is_logged_in = self.browser.check_login_status(login_indicator)
            
            if is_logged_in:
                print("✓ Cookie验证成功：已登录状态")
            else:
                print("⚠ Cookie验证失败：未检测到登录状态")
            
            # 截图
            screenshot_path = f"{target_domain.replace('.', '_').replace(':', '_')}_screenshot.png"
            self.browser.take_screenshot(screenshot_path)
            
            if idx < len(sites):
                print("\n等待3秒后继续下一个网站...")
                time.sleep(3)
        
        # 关闭浏览器
        print(f"\n关闭浏览器")
        print("-" * 60)
        self.browser.close()
        
        print("\n" + "=" * 60)
        print("所有网站访问完成")
        print("=" * 60)


def main():
    # 配置CookieCloud
    config = CookieCloudConfig(
        server="http://192.168.1.100:3000/cookiecloud",
        username="qyZpTrgiP5mVwiRZwfBhjz",
        password="iGn1B4FnWftp3oj4Ko3jxA",
        timeout=30
    )
    
    # 创建示例程序
    demo = CookieDemo(config)
    
    # 测试网站列表
    test_sites = [
        {
            "url": "https://lmkbi.95155.com/bi-system/#/login",
            "name": "BI系统",
            "login_indicator": "彭峰"
        },
        {
            "url": "http://192.168.1.100:3000/#/subscribe/movie",
            "name": "MoviePilot订阅",
            "login_indicator": "搜索"
        },
        {
            "url": "http://192.168.1.102:8418/bwadmin/QueryCarInfo2",
            "name": "QueryCarInfo2",
            "login_indicator": "工单管理"
        }
    ]
    
    # 下载一次数据，依次访问所有网站
    demo.run_sites(
        sites=test_sites,
        headless=False,
        use_edge=True
    )


if __name__ == "__main__":
    main()