#!/usr/bin/env python3
"""
开源情报目标采集器 - 核心模块
用于从公开源采集实体数据（人物、组织、军事装备等）
"""

import json
import time
import random
import os
from datetime import datetime
from typing import Dict, List, Optional, Any
from dataclasses import dataclass, asdict, field
from pathlib import Path

# 导入采集工具
try:
    from crawl4ai import Crawl4AI
    CRAWL4AI_AVAILABLE = True
except ImportError:
    CRAWL4AI_AVAILABLE = False
    print("⚠️ Crawl4AI 未安装")

try:
    import twscrape
    TWSCRAPE_AVAILABLE = True
except ImportError:
    TWSCRAPE_AVAILABLE = False
    print("⚠️ twscrape 未安装")


@dataclass
class DataPassport:
    """数据护照 - 记录数据采集的合规信息"""
    source: str                          # 数据来源
    source_url: str                      # 原始URL
    collection_time: str                 # 采集时间
    compliance_note: str                 # 合规说明
    data_quality: str                    # 数据质量评级
    fields: Dict[str, Any] = field(default_factory=dict)  # 各字段质量
    
    def to_dict(self) -> Dict:
        return asdict(self)


@dataclass
class CollectionConfig:
    """采集配置"""
    target_type: str           # 目标类型：人物/组织/装备等
    language: str = "en"       # 语言：en/zh/ja
    delay_min: int = 8         # 最小延迟（秒）
    delay_max: int = 12        # 最大延迟（秒）
    save_interval: int = 10    # 自动保存间隔
    retry_count: int = 3       # 重试次数
    resume: bool = True        # 启用续传
    output_dir: str = "./data" # 输出目录


class OSINTCollector:
    """开源情报目标采集器"""
    
    def __init__(self, config: CollectionConfig):
        self.config = config
        self.output_dir = Path(config.output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.results: List[Dict] = []
        self.log: List[str] = []
        self.current_index = 0
        
    def add_log(self, message: str):
        """添加日志"""
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        log_msg = f"[{timestamp}] {message}"
        self.log.append(log_msg)
        print(log_msg)
        
    def load_checkpoint(self, checkpoint_file: str) -> int:
        """加载检查点（续传）"""
        if not self.config.resume:
            return 0
        checkpoint_path = self.output_dir / checkpoint_file
        if checkpoint_path.exists():
            with open(checkpoint_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
                self.results = data.get('results', [])
                self.current_index = len(self.results)
                self.add_log(f"✅ 续传：从第 {self.current_index} 条继续")
                return self.current_index
        return 0
    
    def save_checkpoint(self, checkpoint_file: str):
        """保存检查点"""
        checkpoint_path = self.output_dir / checkpoint_file
        checkpoint_data = {
            'results': self.results,
            'index': self.current_index,
            'timestamp': datetime.now().isoformat()
        }
        with open(checkpoint_path, 'w', encoding='utf-8') as f:
            json.dump(checkpoint_data, f, ensure_ascii=False, indent=2)
            
    def random_delay(self):
        """随机延迟"""
        delay = random.randint(self.config.delay_min, self.config.delay_max)
        time.sleep(delay)
        
    def fetch_wikipedia(self, name: str, lang: str = "en") -> Optional[Dict]:
        """获取Wikipedia条目"""
        if not CRAWL4AI_AVAILABLE:
            self.add_log("⚠️ Crawl4AI 不可用")
            return None
            
        lang_map = {'en': 'en.wikipedia.org', 'zh': 'zh.wikipedia.org', 'ja': 'ja.wikipedia.org'}
        domain = lang_map.get(lang, 'en.wikipedia.org')
        url = f"https://{domain}/wiki/{name.replace(' ', '_')}"
        
        self.add_log(f"[FETCH] {url}")
        
        try:
            crawler = Crawl4AI()
            result = crawler.fetch(url)
            
            if result.success:
                return {
                    'url': url,
                    'html': result.html,
                    'markdown': result.markdown,
                    'text': result.text
                }
        except Exception as e:
            self.add_log(f"❌ 采集失败: {e}")
            
        return None
    
    def create_passport(self, source: str, source_url: str, quality: str = "B") -> DataPassport:
        """创建数据护照"""
        return DataPassport(
            source=source,
            source_url=source_url,
            collection_time=datetime.now().isoformat(),
            compliance_note="数据来源于公开网络资源，符合目标地区法律法规",
            data_quality=quality
        )
        
    def save_results(self, output_file: str):
        """保存结果"""
        output_path = self.output_dir / output_file
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(self.results, f, ensure_ascii=False, indent=2)
        self.add_log(f"💾 已保存: {output_path}")
        
    def save_log(self, log_file: str):
        """保存日志"""
        log_path = self.output_dir / log_file
        with open(log_path, 'w', encoding='utf-8') as f:
            f.write('\n'.join(self.log))
            
    def run(self, targets: List[Dict], processor_func=None):
        """
        运行采集任务
        
        Args:
            targets: 目标列表，每项包含 name, url 等
            processor_func: 处理函数，接收(target, fetch_result)，返回处理后的记录
        """
        start_index = self.load_checkpoint('checkpoint.json')
        
        for i, target in enumerate(targets[start_index:], start=start_index):
            self.current_index = i
            self.add_log(f"[{i+1}/{len(targets)}] 采集: {target.get('name', target)}")
            
            # 随机延迟
            self.random_delay()
            
            # 采集数据
            fetch_result = None
            if target.get('url'):
                fetch_result = self.fetch_wikipedia(
                    target['name'], 
                    target.get('lang', self.config.language)
                )
                
            # 处理数据
            record = target.copy()
            if processor_func and fetch_result:
                record = processor_func(target, fetch_result)
                
            self.results.append(record)
            
            # 自动保存
            if (i + 1) % self.config.save_interval == 0:
                self.save_checkpoint('checkpoint.json')
                self.add_log(f"📊 进度: {i+1}/{len(targets)}")
                
        # 最终保存
        self.save_checkpoint('checkpoint.json')
        self.save_log('collection.log')
        
        self.add_log(f"✅ 采集完成：{len(self.results)} 条记录")


def create_collector(target_type: str, **kwargs) -> OSINTCollector:
    """创建采集器实例"""
    config = CollectionConfig(
        target_type=target_type,
        **{k: v for k, v in kwargs.items() if k in ['language', 'delay_min', 'delay_max', 'save_interval', 'retry_count', 'resume', 'output_dir']}
    )
    return OSINTCollector(config)


if __name__ == "__main__":
    print("OSINT Target Collector Module")
    print("用法: from collector import create_collector, OSINTCollector")