import os
import time

COOKIE_FILE = r"C:\Users\Administrator\Desktop\TikTokDownload-main\douyin_cookie.txt"

def get_douyin_cookie_auto(force_refresh=False, headless=True):
    """
    自动获取抖音访客Cookie（使用Edge本地驱动）
    """
    
    # 检查是否已有Cookie且不需要强制刷新
    if not force_refresh and os.path.exists(COOKIE_FILE):
        try:
            with open(COOKIE_FILE, 'r', encoding='utf-8') as f:
                cookie = f.read().strip()
                if cookie and len(cookie) > 100:
                    print("[自动Cookie] 使用已保存的Cookie")
                    return cookie
        except Exception as e:
            print(f"[自动Cookie] 读取Cookie失败: {e}")
    
    print("[自动Cookie] 开始自动获取...")
    
    try:
        # 尝试使用selenium with本地Edge
        from selenium import webdriver
        from selenium.webdriver.edge.options import Options as EdgeOptions
        from selenium.webdriver.edge.service import Service as EdgeService
        
        print("[自动Cookie] 启动Edge浏览器...")
        
        options = EdgeOptions()
        
        if headless:
            print("[自动Cookie] 使用无头模式...")
            options.add_argument("--headless")
            options.add_argument("--disable-gpu")
        
        options.add_argument("--no-sandbox")
        options.add_argument("--disable-dev-shm-usage")
        options.add_argument("--disable-blink-features=AutomationControlled")
        options.add_experimental_option('excludeSwitches', ['enable-automation'])
        options.add_experimental_option('useAutomationExtension', False)
        
        # 尝试使用系统Edge驱动
        try:
            # 方法1: 直接使用Edge（自动查找驱动）
            driver = webdriver.Edge(options=options)
            print("[自动Cookie] Edge浏览器启动成功（自动驱动）")
        except Exception as e1:
            print(f"[自动Cookie] 自动驱动失败: {e1}")
            
            # 方法2: 尝试常见驱动路径
            common_paths = [
                r"C:\Program Files (x86)\Microsoft\Edge\Application\msedgedriver.exe",
                r"C:\Program Files\Microsoft\Edge\Application\msedgedriver.exe",
                os.path.expandvars(r"%PROGRAMFILES%\Microsoft\Edge\Application\msedgedriver.exe"),
                os.path.expandvars(r"%PROGRAMFILES(X86)%\Microsoft\Edge\Application\msedgedriver.exe"),
            ]
            
            driver = None
            for driver_path in common_paths:
                if os.path.exists(driver_path):
                    print(f"[自动Cookie] 尝试使用驱动: {driver_path}")
                    try:
                        service = EdgeService(executable_path=driver_path)
                        driver = webdriver.Edge(service=service, options=options)
                        print("[自动Cookie] Edge浏览器启动成功（指定驱动）")
                        break
                    except Exception as e2:
                        print(f"[自动Cookie] 驱动 {driver_path} 失败: {e2}")
                        continue
            
            if driver is None:
                raise Exception("无法找到可用的Edge驱动")
        
        # 访问抖音
        print("[自动Cookie] 访问抖音网站...")
        driver.get("https://www.douyin.com")
        
        print("[自动Cookie] 等待Cookie生成（5秒）...")
        time.sleep(5)
        
        # 提取Cookie
        cookies = driver.get_cookies()
        
        if not cookies:
            print("[自动Cookie] 警告: 未获取到Cookie")
            driver.quit()
            return None
        
        # 转换为字符串
        cookie_str = "; ".join([f"{c['name']}={c['value']}" for c in cookies])
        
        print(f"[自动Cookie] 成功获取 {len(cookies)} 个Cookie项")
        print(f"[自动Cookie] Cookie项: {', '.join([c['name'] for c in cookies])}")
        print(f"[自动Cookie] Cookie长度: {len(cookie_str)}")
        
        # 保存
        try:
            with open(COOKIE_FILE, 'w', encoding='utf-8') as f:
                f.write(cookie_str)
            print(f"[自动Cookie] Cookie已保存")
        except Exception as e:
            print(f"[自动Cookie] 保存失败: {e}")
        
        driver.quit()
        return cookie_str
        
    except ImportError as e:
        print(f"[自动Cookie] 错误: selenium未安装 - {e}")
        print("[自动Cookie] 提示: 请运行 pip install selenium")
        return None
    except Exception as e:
        print(f"[自动Cookie] 错误: {e}")
        import traceback
        traceback.print_exc()
        return None

if __name__ == "__main__":
    print("="*60)
    print("测试：自动获取Cookie（Edge本地驱动）")
    print("="*60)
    
    cookie = get_douyin_cookie_auto(force_refresh=True, headless=False)
    
    if cookie:
        print(f"\n成功！Cookie长度: {len(cookie)}")
        print(f"Cookie前200字符:\n{cookie[:200]}...")
    else:
        print("\n失败：无法获取Cookie")