#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
图像存储管理器
统一管理小红书笔记的图片和视频下载和存储
"""

import os
import json
import hashlib
import time
import requests
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
import logging
logger = logging.getLogger(__name__)
from urllib.parse import urlparse

@dataclass
class MediaInfo:
    """媒体文件信息"""
    media_id: str
    media_type: str  # 'image' or 'video'
    original_url: str
    local_path: str
    file_size: int
    download_time: datetime
    note_id: str
    note_title: str
    user_id: str
    user_name: str
    
    def to_dict(self) -> Dict:
        return {
            'media_id': self.media_id,
            'media_type': self.media_type,
            'original_url': self.original_url,
            'local_path': self.local_path,
            'file_size': self.file_size,
            'download_time': self.download_time.isoformat(),
            'note_id': self.note_id,
            'note_title': self.note_title,
            'user_id': self.user_id,
            'user_name': self.user_name
        }

class ImageStorageManager:
    """图像存储管理器"""
    
    def __init__(self, base_storage_path: str = "data/media"):
        self.base_storage_path = Path(base_storage_path)
        self.media_index_file = self.base_storage_path / "media_index.json"
        self.media_index: Dict[str, MediaInfo] = {}
        self.chunk_size = 1024 * 1024  # 1MB chunks
        
        # 创建存储目录
        self.image_path = self.base_storage_path / "images"
        self.video_path = self.base_storage_path / "videos"
        self.thumbnail_path = self.base_storage_path / "thumbnails"
        
        for path in [self.image_path, self.video_path, self.thumbnail_path]:
            path.mkdir(parents=True, exist_ok=True)
        
        self.load_media_index()
    
    def load_media_index(self) -> None:
        """加载媒体索引文件"""
        try:
            if self.media_index_file.exists():
                with open(self.media_index_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    
                for media_id, media_data in data.items():
                    self.media_index[media_id] = MediaInfo(
                        media_id=media_data['media_id'],
                        media_type=media_data['media_type'],
                        original_url=media_data['original_url'],
                        local_path=media_data['local_path'],
                        file_size=media_data['file_size'],
                        download_time=datetime.fromisoformat(media_data['download_time']),
                        note_id=media_data['note_id'],
                        note_title=media_data['note_title'],
                        user_id=media_data['user_id'],
                        user_name=media_data['user_name']
                    )
                logger.info(f"加载了 {len(self.media_index)} 个媒体文件索引")
        except Exception as e:
            logger.error(f"加载媒体索引失败: {str(e)}")
    
    def save_media_index(self) -> None:
        """保存媒体索引文件"""
        try:
            data = {media_id: media_info.to_dict() 
                   for media_id, media_info in self.media_index.items()}
            
            with open(self.media_index_file, 'w', encoding='utf-8') as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
            logger.debug("保存媒体索引成功")
        except Exception as e:
            logger.error(f"保存媒体索引失败: {str(e)}")
    
    def _generate_media_id(self, url: str, note_id: str) -> str:
        """生成媒体文件ID"""
        content = f"{url}_{note_id}"
        return hashlib.md5(content.encode()).hexdigest()
    
    def _get_file_extension(self, url: str, media_type: str) -> str:
        """获取文件扩展名"""
        parsed_url = urlparse(url)
        path = parsed_url.path
        
        if path and '.' in path:
            return os.path.splitext(path)[1].lower()
        
        # 默认扩展名
        return '.jpg' if media_type == 'image' else '.mp4'
    
    def _normalize_filename(self, filename: str) -> str:
        """规范化文件名"""
        # 移除不安全字符
        unsafe_chars = '<>:"/\\|?*'
        for char in unsafe_chars:
            filename = filename.replace(char, '_')
        
        # 限制长度
        if len(filename) > 100:
            filename = filename[:100]
        
        return filename
    
    def _download_media(self, url: str, local_path: str, max_retries: int = 3) -> Tuple[bool, str, int]:
        """下载媒体文件"""
        for attempt in range(max_retries):
            try:
                headers = {
                    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
                }
                
                response = requests.get(url, headers=headers, stream=True, timeout=30)
                response.raise_for_status()
                
                file_size = 0
                with open(local_path, 'wb') as f:
                    for chunk in response.iter_content(chunk_size=self.chunk_size):
                        if chunk:
                            f.write(chunk)
                            file_size += len(chunk)
                
                logger.info(f"下载成功: {local_path} ({file_size} bytes)")
                return True, "下载成功", file_size
                
            except requests.exceptions.RequestException as e:
                logger.warning(f"下载失败 (尝试 {attempt + 1}/{max_retries}): {str(e)}")
                if attempt < max_retries - 1:
                    time.sleep(2 ** attempt)  # 指数退避
                else:
                    return False, f"下载失败: {str(e)}", 0
            except Exception as e:
                logger.error(f"下载出错: {str(e)}")
                return False, f"下载出错: {str(e)}", 0
        
        return False, "下载失败", 0
    
    def download_image(self, url: str, note_info: Dict) -> Optional[MediaInfo]:
        """下载图片"""
        media_id = self._generate_media_id(url, note_info['note_id'])
        
        # 检查是否已存在
        if media_id in self.media_index:
            existing_media = self.media_index[media_id]
            if os.path.exists(existing_media.local_path):
                logger.info(f"图片已存在: {existing_media.local_path}")
                return existing_media
        
        # 生成本地路径
        extension = self._get_file_extension(url, 'image')
        user_name = self._normalize_filename(note_info.get('nickname', 'unknown'))
        note_title = self._normalize_filename(note_info.get('title', 'untitled'))
        
        filename = f"{user_name}_{note_info['note_id']}_{media_id}{extension}"
        local_path = self.image_path / filename
        
        # 下载文件
        success, message, file_size = self._download_media(url, str(local_path))
        
        if success:
            media_info = MediaInfo(
                media_id=media_id,
                media_type='image',
                original_url=url,
                local_path=str(local_path),
                file_size=file_size,
                download_time=datetime.now(),
                note_id=note_info['note_id'],
                note_title=note_info.get('title', ''),
                user_id=note_info.get('user_id', ''),
                user_name=note_info.get('nickname', '')
            )
            
            self.media_index[media_id] = media_info
            self.save_media_index()
            return media_info
        else:
            logger.error(f"图片下载失败: {message}")
            return None
    
    def download_video(self, url: str, note_info: Dict, thumbnail_url: Optional[str] = None) -> Optional[MediaInfo]:
        """下载视频"""
        media_id = self._generate_media_id(url, note_info['note_id'])
        
        # 检查是否已存在
        if media_id in self.media_index:
            existing_media = self.media_index[media_id]
            if os.path.exists(existing_media.local_path):
                logger.info(f"视频已存在: {existing_media.local_path}")
                return existing_media
        
        # 生成本地路径
        extension = self._get_file_extension(url, 'video')
        user_name = self._normalize_filename(note_info.get('nickname', 'unknown'))
        note_title = self._normalize_filename(note_info.get('title', 'untitled'))
        
        filename = f"{user_name}_{note_info['note_id']}_{media_id}{extension}"
        local_path = self.video_path / filename
        
        # 下载视频文件
        success, message, file_size = self._download_media(url, str(local_path))
        
        if success:
            # 如果有缩略图，也下载缩略图
            if thumbnail_url:
                thumbnail_filename = f"{user_name}_{note_info['note_id']}_{media_id}_thumb.jpg"
                thumbnail_path = self.thumbnail_path / thumbnail_filename
                self._download_media(thumbnail_url, str(thumbnail_path))
            
            media_info = MediaInfo(
                media_id=media_id,
                media_type='video',
                original_url=url,
                local_path=str(local_path),
                file_size=file_size,
                download_time=datetime.now(),
                note_id=note_info['note_id'],
                note_title=note_info.get('title', ''),
                user_id=note_info.get('user_id', ''),
                user_name=note_info.get('nickname', '')
            )
            
            self.media_index[media_id] = media_info
            self.save_media_index()
            return media_info
        else:
            logger.error(f"视频下载失败: {message}")
            return None
    
    def download_note_media(self, note_info: Dict, download_images: bool = True, download_videos: bool = True) -> List[MediaInfo]:
        """下载笔记的所有媒体文件"""
        downloaded_media = []
        
        # 下载图片
        if download_images and note_info.get('image_list'):
            for image_url in note_info['image_list']:
                media_info = self.download_image(image_url, note_info)
                if media_info:
                    downloaded_media.append(media_info)
        
        # 下载视频
        if download_videos and note_info.get('video_addr'):
            media_info = self.download_video(
                note_info['video_addr'], 
                note_info, 
                note_info.get('video_cover')
            )
            if media_info:
                downloaded_media.append(media_info)
        
        return downloaded_media
    
    def get_media_by_note(self, note_id: str) -> List[MediaInfo]:
        """根据笔记ID获取媒体文件"""
        return [media for media in self.media_index.values() if media.note_id == note_id]
    
    def get_media_by_user(self, user_id: str) -> List[MediaInfo]:
        """根据用户ID获取媒体文件"""
        return [media for media in self.media_index.values() if media.user_id == user_id]
    
    def get_storage_statistics(self) -> Dict:
        """获取存储统计信息"""
        total_files = len(self.media_index)
        total_size = sum(media.file_size for media in self.media_index.values())
        
        image_count = sum(1 for media in self.media_index.values() if media.media_type == 'image')
        video_count = sum(1 for media in self.media_index.values() if media.media_type == 'video')
        
        image_size = sum(media.file_size for media in self.media_index.values() if media.media_type == 'image')
        video_size = sum(media.file_size for media in self.media_index.values() if media.media_type == 'video')
        
        return {
            'total_files': total_files,
            'total_size': total_size,
            'total_size_mb': round(total_size / (1024 * 1024), 2),
            'image_count': image_count,
            'video_count': video_count,
            'image_size': image_size,
            'video_size': video_size,
            'image_size_mb': round(image_size / (1024 * 1024), 2),
            'video_size_mb': round(video_size / (1024 * 1024), 2),
            'storage_path': str(self.base_storage_path)
        }
    
    def cleanup_orphaned_files(self) -> int:
        """清理孤立的文件（没有索引记录的文件）"""
        orphaned_count = 0
        
        # 检查图片目录
        if self.image_path.exists():
            for file_path in self.image_path.iterdir():
                if file_path.is_file():
                    file_found = False
                    for media_info in self.media_index.values():
                        if Path(media_info.local_path) == file_path:
                            file_found = True
                            break
                    
                    if not file_found:
                        try:
                            file_path.unlink()
                            orphaned_count += 1
                            logger.info(f"删除孤立文件: {file_path}")
                        except Exception as e:
                            logger.error(f"删除孤立文件失败: {str(e)}")
        
        # 检查视频目录
        if self.video_path.exists():
            for file_path in self.video_path.iterdir():
                if file_path.is_file():
                    file_found = False
                    for media_info in self.media_index.values():
                        if Path(media_info.local_path) == file_path:
                            file_found = True
                            break
                    
                    if not file_found:
                        try:
                            file_path.unlink()
                            orphaned_count += 1
                            logger.info(f"删除孤立文件: {file_path}")
                        except Exception as e:
                            logger.error(f"删除孤立文件失败: {str(e)}")
        
        logger.info(f"清理了 {orphaned_count} 个孤立文件")
        return orphaned_count
    
    def remove_media(self, media_id: str) -> bool:
        """删除媒体文件"""
        if media_id not in self.media_index:
            return False
        
        media_info = self.media_index[media_id]
        
        # 删除文件
        try:
            if os.path.exists(media_info.local_path):
                os.remove(media_info.local_path)
                logger.info(f"删除媒体文件: {media_info.local_path}")
        except Exception as e:
            logger.error(f"删除媒体文件失败: {str(e)}")
            return False
        
        # 从索引中删除
        del self.media_index[media_id]
        self.save_media_index()
        
        return True 