"""
Cookie和Session管理模块
"""
import json
from typing import Optional, List, Dict, Any
from pathlib import Path
from playwright.async_api import BrowserContext, Cookie
from loguru import logger
from datetime import datetime


class CookieManager:
    """
    Cookie管理器
    支持保存、加载、导入、导出Cookie
    """
    
    def __init__(self, context: BrowserContext):
        self.context = context
    
    async def get_all(self) -> List[Cookie]:
        """获取所有Cookie"""
        return await self.context.cookies()
    
    async def get(self, domain: str) -> List[Cookie]:
        """
        获取指定域名的Cookie
        
        Args:
            domain: 域名
        """
        return await self.context.cookies([domain])
    
    async def set(self, cookies: List[Dict[str, Any]]):
        """
        设置Cookie
        
        Args:
            cookies: Cookie列表
        """
        await self.context.add_cookies(cookies)
        logger.info(f"设置Cookie: {len(cookies)} 个")
    
    async def set_one(
        self,
        name: str,
        value: str,
        domain: str,
        path: str = "/",
        expires: Optional[int] = None,
        http_only: bool = False,
        secure: bool = False,
    ):
        """
        设置单个Cookie
        
        Args:
            name: Cookie名称
            value: Cookie值
            domain: 域名
            path: 路径
            expires: 过期时间（Unix时间戳）
            http_only: 是否HttpOnly
            secure: 是否Secure
        """
        cookie = {
            "name": name,
            "value": value,
            "domain": domain,
            "path": path,
            "httpOnly": http_only,
            "secure": secure,
        }
        if expires:
            cookie["expires"] = expires
        
        await self.context.add_cookies([cookie])
        logger.info(f"设置Cookie: {name}={value[:20]}...")
    
    async def clear(self, domain: Optional[str] = None):
        """
        清除Cookie
        
        Args:
            domain: 指定域名，None表示清除所有
        """
        if domain:
            await self.context.clear_cookies([domain])
            logger.info(f"清除Cookie: {domain}")
        else:
            await self.context.clear_cookies()
            logger.info("清除所有Cookie")
    
    async def save(self, filepath: str, domain: Optional[str] = None):
        """
        保存Cookie到文件
        
        Args:
            filepath: 文件路径
            domain: 指定域名，None表示保存所有
        """
        cookies = await self.get(domain) if domain else await self.get_all()
        
        data = {
            "cookies": cookies,
            "saved_at": datetime.now().isoformat(),
            "domain": domain,
        }
        
        Path(filepath).parent.mkdir(parents=True, exist_ok=True)
        with open(filepath, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        
        logger.info(f"Cookie已保存: {filepath}, 共 {len(cookies)} 个")
    
    async def load(self, filepath: str) -> int:
        """
        从文件加载Cookie
        
        Args:
            filepath: 文件路径
            
        Returns:
            int: 加载的Cookie数量
        """
        with open(filepath, "r", encoding="utf-8") as f:
            data = json.load(f)
        
        cookies = data.get("cookies", [])
        await self.set(cookies)
        
        logger.info(f"Cookie已加载: {filepath}, 共 {len(cookies)} 个")
        return len(cookies)
    
    async def export_netscape(self, filepath: str, domain: Optional[str] = None):
        """
        导出为Netscape格式（wget/curl兼容）
        
        Args:
            filepath: 文件路径
            domain: 指定域名
        """
        cookies = await self.get(domain) if domain else await self.get_all()
        
        lines = ["# Netscape HTTP Cookie File\n"]
        for cookie in cookies:
            line = f".{cookie['domain']}\tTRUE\t{cookie['path']}\t"
            line += "TRUE\t" if cookie.get("secure") else "FALSE\t"
            line += f"{cookie.get('expires', 0)}\t{cookie['name']}\t{cookie['value']}\n"
            lines.append(line)
        
        with open(filepath, "w", encoding="utf-8") as f:
            f.writelines(lines)
        
        logger.info(f"Cookie已导出(Netscape格式): {filepath}")
    
    async def import_netscape(self, filepath: str) -> int:
        """
        从Netscape格式导入
        
        Args:
            filepath: 文件路径
            
        Returns:
            int: 导入的Cookie数量
        """
        cookies = []
        
        with open(filepath, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith("#"):
                    continue
                
                parts = line.split("\t")
                if len(parts) >= 7:
                    cookies.append({
                        "domain": parts[0].lstrip("."),
                        "path": parts[2],
                        "secure": parts[3] == "TRUE",
                        "expires": int(parts[4]),
                        "name": parts[5],
                        "value": parts[6],
                    })
        
        await self.set(cookies)
        logger.info(f"Cookie已导入(Netscape格式): {filepath}, 共 {len(cookies)} 个")
        return len(cookies)
    
    async def print_all(self):
        """打印所有Cookie"""
        cookies = await self.get_all()
        print(f"\n=== Cookies ({len(cookies)} 个) ===")
        for cookie in cookies:
            print(f"  {cookie['name']}: {cookie['value'][:30]}... ({cookie['domain']})")


class SessionManager:
    """
    Session管理器
    支持保存和恢复整个浏览器会话状态
    """
    
    def __init__(self, context: BrowserContext):
        self.context = context
        self.cookie_manager = CookieManager(context)
    
    async def save(self, filepath: str):
        """
        保存会话状态
        
        Args:
            filepath: 文件路径
        """
        cookies = await self.context.cookies()
        storage = await self.context.storage_state()
        
        data = {
            "cookies": cookies,
            "storage_state": storage,
            "saved_at": datetime.now().isoformat(),
        }
        
        Path(filepath).parent.mkdir(parents=True, exist_ok=True)
        with open(filepath, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        
        logger.info(f"会话已保存: {filepath}")
    
    async def load(self, filepath: str):
        """
        加载会话状态
        
        Args:
            filepath: 文件路径
        """
        with open(filepath, "r", encoding="utf-8") as f:
            data = json.load(f)
        
        if "cookies" in data:
            await self.context.add_cookies(data["cookies"])
        
        # 注意：storage_state需要在创建context时设置
        # 这里只能设置cookies
        
        logger.info(f"会话已加载: {filepath}")
    
    async def login_state_save(self, name: str, storage_dir: str = "./sessions"):
        """
        保存登录状态
        
        Args:
            name: 状态名称（如网站名）
            storage_dir: 存储目录
        """
        filepath = f"{storage_dir}/{name}_session.json"
        await self.save(filepath)
        return filepath
    
    async def login_state_load(self, name: str, storage_dir: str = "./sessions") -> bool:
        """
        加载登录状态
        
        Args:
            name: 状态名称
            storage_dir: 存储目录
            
        Returns:
            bool: 是否加载成功
        """
        filepath = f"{storage_dir}/{name}_session.json"
        
        if not Path(filepath).exists():
            logger.warning(f"登录状态文件不存在: {filepath}")
            return False
        
        await self.load(filepath)
        return True
