#!/usr/bin/env python3
"""
图片/头像采集器
从Wikipedia、官网、社交媒体采集目标头像和图片
"""

import os
import re
import json
import hashlib
import time
import random
from pathlib import Path
from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass
from urllib.parse import urljoin, urlparse
from datetime import datetime

# 尝试导入图片处理库
try:
    from PIL import Image
    PIL_AVAILABLE = True
except ImportError:
    PIL_AVAILABLE = False

try:
    import requests
    REQUESTS_AVAILABLE = True
except ImportError:
    REQUESTS_AVAILABLE = False

try:
    from crawl4ai import Crawl4AI
    CRAWL4AI_AVAILABLE = True
except ImportError:
    CRAWL4AI_AVAILABLE = False


@dataclass
class ImageConfig:
    """图片采集配置"""
    output_dir: str = "./images"           # 输出目录
    max_size_mb: int = 5                   # 最大文件大小(MB)
    min_size_kb: int = 2                   # 最小文件大小(KB)
    timeout: int = 30                      # 下载超时(秒)
    delay_min: int = 3                     # 最小延迟(秒)
    delay_max: int = 6                     # 最大延迟(秒)
    user_agent: str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
    allowed_formats: Tuple[str, ...] = ("jpg", "jpeg", "png", "webp", "gif")
    save_thumbnails: bool = True           # 保存缩略图
    thumbnail_size: Tuple[int, int] = (200, 200)  # 缩略图尺寸


class ImageCollector:
    """图片/头像采集器"""
    
    # 常见头像容器选择器
    AVATAR_SELECTORS = {
        'infobox': [
            '.infobox .portable-infobox .pi-image img',
            '.infobox img[alt*="portrait"]',
            '.infobox img[alt*="photo"]',
            '.infobox img:first-of-type',
            'table.infobox img',
        ],
        'sidebar': [
            '.sidebar img',
            '.portable-infobox img',
        ],
        'general': [
            '.mw-parser-output img[alt*="portrait"]',
            '.mw-parser-output img[alt*="photo"]',
            '.avatar',
            '.user-avatar img',
            '.profile-image img',
        ]
    }
    
    # 常见图片文件名模式
    AVATAR_FILENAME_PATTERNS = [
        r'avatar',
        r'portrait',
        r'photo',
        r'headshot',
        r'profile',
        r'profile_image',
        r'profile-pic',
    ]
    
    def __init__(self, config: ImageConfig = None):
        self.config = config or ImageConfig()
        self.output_dir = Path(self.config.output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
        # 初始化HTTP会话
        self.session = None
        if REQUESTS_AVAILABLE:
            self.session = requests.Session()
            self.session.headers.update({'User-Agent': self.config.user_agent})
            
    def _random_delay(self):
        """随机延迟"""
        delay = random.uniform(self.config.delay_min, self.config.delay_max)
        time.sleep(delay)
        
    def _is_valid_image_format(self, url: str) -> bool:
        """检查是否为有效图片格式"""
        parsed = urlparse(url)
        path = parsed.path.lower()
        return any(path.endswith(f'.{fmt}') for fmt in self.config.allowed_formats)
        
    def _download_image(self, url: str, filename: str) -> Optional[Dict]:
        """下载图片"""
        if not REQUESTS_AVAILABLE:
            return None
            
        try:
            self._random_delay()
            
            response = self.session.get(url, timeout=self.config.timeout, stream=True)
            response.raise_for_status()
            
            # 检查内容类型
            content_type = response.headers.get('Content-Type', '')
            if 'image' not in content_type and not self._is_valid_image_format(url):
                return None
                
            # 检查文件大小
            content_length = response.headers.get('Content-Length')
            if content_length:
                size_mb = int(content_length) / (1024 * 1024)
                if size_mb > self.config.max_size_mb:
                    return None
                    
            # 保存图片
            filepath = self.output_dir / filename
            with open(filepath, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)
                    
            # 验证图片
            if PIL_AVAILABLE:
                try:
                    img = Image.open(filepath)
                    width, height = img.size
                    file_size_kb = filepath.stat().st_size / 1024
                    
                    if file_size_kb < self.config.min_size_kb:
                        filepath.unlink()
                        return None
                        
                    # 生成缩略图
                    thumbnail_path = None
                    if self.config.save_thumbnails:
                        thumb_name = f"thumb_{filename}"
                        thumb_path = self.output_dir / thumb_name
                        img.thumbnail(self.config.thumbnail_size)
                        img.save(thumb_path)
                        thumbnail_path = str(thumb_path)
                        
                    return {
                        'filename': filename,
                        'filepath': str(filepath),
                        'thumbnail_path': thumbnail_path,
                        'width': width,
                        'height': height,
                        'file_size_kb': round(file_size_kb, 2),
                        'format': img.format,
                        'download_time': datetime.now().isoformat()
                    }
                except Exception:
                    return None
            else:
                return {
                    'filename': filename,
                    'filepath': str(filepath),
                    'download_time': datetime.now().isoformat()
                }
                
        except Exception as e:
            return None
            
    def _sanitize_filename(self, name: str) -> str:
        """清理文件名"""
        name = re.sub(r'[^\w\s-]', '', name)
        name = re.sub(r'[-\s]+', '_', name)
        return name[:50]
        
    def _generate_hash(self, content: bytes) -> str:
        """生成内容哈希"""
        return hashlib.md5(content).hexdigest()[:12]
        
    def extract_images_from_html(self, html: str, base_url: str = "") -> List[Dict]:
        """从HTML中提取图片信息"""
        images = []
        
        # 提取所有img标签
        img_pattern = r'<img[^>]+>'
        for img_tag in re.findall(img_pattern, html, re.IGNORECASE):
            # 获取src
            src_match = re.search(r'src=["\']([^"\']+)["\']', img_tag)
            if not src_match:
                continue
                
            src = src_match.group(1)
            if src.startswith('//'):
                src = 'https:' + src
            elif src.startswith('/'):
                src = urljoin(base_url, src)
                
            # 过滤无效URL
            if not src.startswith('http'):
                continue
            if 'icon' in src.lower() or 'logo' in src.lower() and 'avatar' not in src.lower():
                continue
                
            # 获取alt属性
            alt_match = re.search(r'alt=["\']([^"\']*)["\']', img_tag)
            alt = alt_match.group(1) if alt_match else ""
            
            # 获取width/height
            width_match = re.search(r'width=["\'](\d+)["\']', img_tag)
            height_match = re.search(r'height=["\'](\d+)["\']', img_tag)
            
            images.append({
                'url': src,
                'alt': alt,
                'width': int(width_match.group(1)) if width_match else 0,
                'height': int(height_match.group(1)) if height_match else 0,
            })
            
        return images
        
    def extract_images_from_markdown(self, markdown: str, base_url: str = "") -> List[Dict]:
        """从Markdown中提取图片信息"""
        images = []
        
        # 提取![alt](url)格式
        pattern = r'!\[([^\]]*)\]\(([^)]+)\)'
        for alt, url in re.findall(pattern, markdown):
            if url.startswith('//'):
                url = 'https:' + url
            elif url.startswith('/'):
                url = urljoin(base_url, url)
                
            if url.startswith('http'):
                images.append({
                    'url': url,
                    'alt': alt,
                    'width': 0,
                    'height': 0,
                })
                
        return images
        
    def find_best_avatar(self, images: List[Dict], target_name: str = "") -> Optional[Dict]:
        """找到最佳头像"""
        if not images:
            return None
            
        # 优先选择规则:
        # 1. 包含portrait/photo关键词
        # 2. 有一定的尺寸(宽度>100)
        # 3. 匹配目标名称
        
        candidates = []
        for img in images:
            score = 0
            url = img.get('url', '').lower()
            alt = img.get('alt', '').lower()
            target = target_name.lower()
            
            # 关键词匹配
            for kw in self.AVATAR_FILENAME_PATTERNS:
                if re.search(kw, url) or re.search(kw, alt):
                    score += 10
                    
            # 尺寸评分
            width = img.get('width', 0)
            if width > 200:
                score += 5
            elif width > 100:
                score += 3
                
            # 名称匹配
            if target and (target in url or target in alt):
                score += 5
                
            candidates.append((score, img))
            
        if candidates:
            candidates.sort(key=lambda x: -x[0])
            return candidates[0][1]
            
        # 如果没有匹配，返回最大的图片
        valid = [img for img in images if img.get('width', 0) > 100]
        if valid:
            return max(valid, key=lambda x: x.get('width', 0))
            
        return images[0] if images else None
        
    def collect_from_wikipedia(self, target_name: str, url: str) -> Optional[Dict]:
        """从Wikipedia采集头像"""
        if not CRAWL4AI_AVAILABLE:
            return None
            
        try:
            from crawl4ai import Crawl4AI
            crawler = Crawl4AI()
            result = crawler.fetch(url)
            
            if not result.success:
                return None
                
            # 提取图片
            html = result.html or ""
            base_url = url
            images = self.extract_images_from_html(html, base_url)
            images.extend(self.extract_images_from_markdown(result.markdown or "", base_url))
            
            # 找到最佳头像
            avatar = self.find_best_avatar(images, target_name)
            
            if not avatar:
                return None
                
            # 下载头像
            safe_name = self._sanitize_filename(target_name)
            filename = f"{safe_name}_avatar{Path(urlparse(avatar['url']).suffix).name}"
            if not filename.endswith(('.jpg', '.jpeg', '.png', '.webp', '.gif')):
                filename += '.jpg'
                
            download_result = self._download_image(avatar['url'], filename)
            
            if download_result:
                return {
                    'target_name': target_name,
                    'source_url': url,
                    'image_url': avatar['url'],
                    'alt_text': avatar.get('alt', ''),
                    'collection_time': datetime.now().isoformat(),
                    'data_passport': {
                        'source': 'Wikipedia',
                        'source_url': url,
                        'collection_time': datetime.now().isoformat(),
                        'type': 'avatar',
                        'compliance_note': '数据来源于公开网络资源'
                    },
                    **download_result
                }
                
        except Exception as e:
            pass
            
        return None
        
    def collect_from_url(self, target_name: str, url: str) -> Optional[Dict]:
        """从指定URL采集图片"""
        if not CRAWL4AI_AVAILABLE:
            return None
            
        try:
            from crawl4ai import Crawl4AI
            crawler = Crawl4AI()
            result = crawler.fetch(url)
            
            if not result.success:
                return None
                
            base_url = url
            images = self.extract_images_from_html(result.html or "", base_url)
            images.extend(self.extract_images_from_markdown(result.markdown or "", base_url))
            
            avatar = self.find_best_avatar(images, target_name)
            
            if not avatar:
                return None
                
            safe_name = self._sanitize_filename(target_name)
            filename = f"{safe_name}_{self._generate_hash(avatar['url'].encode())[:6]}"
            ext = Path(urlparse(avatar['url']).suffix)
            filename += ext if ext else '.jpg'
            
            download_result = self._download_image(avatar['url'], filename)
            
            if download_result:
                return {
                    'target_name': target_name,
                    'source_url': url,
                    'image_url': avatar['url'],
                    'alt_text': avatar.get('alt', ''),
                    'collection_time': datetime.now().isoformat(),
                    'data_passport': {
                        'source': 'Web',
                        'source_url': url,
                        'collection_time': datetime.now().isoformat(),
                        'type': 'image',
                        'compliance_note': '数据来源于公开网络资源'
                    },
                    **download_result
                }
                
        except Exception as e:
            pass
            
        return None


def create_image_collector(**kwargs) -> ImageCollector:
    """创建图片采集器"""
    config = ImageConfig(**kwargs)
    return ImageCollector(config)


def collect_target_image(target: Dict, config: ImageConfig = None) -> Optional[Dict]:
    """
    采集目标图片/头像
    
    Args:
        target: 目标信息，包含 name, url, wiki_url 等
        config: 采集配置
        
    Returns:
        采集结果或None
    """
    collector = ImageCollector(config)
    
    name = target.get('name') or target.get('name_cn') or target.get('name_en')
    if not name:
        return None
        
    # 优先尝试Wikipedia
    wiki_url = target.get('url') or target.get('wiki_url') or target.get('wiki链接')
    if wiki_url and 'wikipedia' in wiki_url.lower():
        return collector.collect_from_wikipedia(name, wiki_url)
        
    # 其他URL
    if wiki_url:
        return collector.collect_from_url(name, wiki_url)
        
    return None


def collect_batch(targets: List[Dict], config: ImageConfig = None) -> List[Dict]:
    """批量采集目标图片"""
    collector = ImageCollector(config)
    results = []
    
    for target in targets:
        name = target.get('name') or target.get('name_cn') or target.get('name_en')
        url = target.get('url') or target.get('wiki_url') or target.get('wiki链接')
        
        if not name or not url:
            continue
            
        print(f"📷 采集图片: {name}")
        
        result = None
        if 'wikipedia' in url.lower():
            result = collector.collect_from_wikipedia(name, url)
        else:
            result = collector.collect_from_url(name, url)
            
        if result:
            results.append(result)
            print(f"   ✅ 保存: {result.get('filename')}")
        else:
            print(f"   ❌ 未找到图片")
            
    return results


if __name__ == "__main__":
    print("Image Collector Module")
    print("用法: from image_collector import ImageCollector, collect_target_image")