#!/usr/bin/env python3
"""
Скачивает обложки аниме через Jikan API (jikan.moe) и встраивает их в аудиофайлы.
Использует SOCKS5 прокси для доступа к API и скачивания изображений.

Название аниме берётся из первых квадратных скобок в названии файла:
  TV-2 OP01 - Nobody Knows (Suga Shikao) [xxxHOLiC].mp3
                                           ^^^^^^^^
                                           извлекаем это

Jikan API - бесплатный API без аутентификации (парсит MyAnimeList)
https://jikan.moe/

Использование:
  python3 download_anime_covers_proxy.py run --dry-run /path/to/music
  python3 download_anime_covers_proxy.py run /path/to/music
  python3 download_anime_covers_proxy.py search "xxxHOLiC"  # Поиск обложки для теста
"""

import os
import sys
import argparse
import re
import time
import requests
from pathlib import Path
from typing import Optional, Tuple, Dict

try:
    from mutagen.flac import FLAC
    from mutagen.mp3 import MP3
    from mutagen.oggvorbis import OggVorbis
    from mutagen.oggopus import OggOpus
    from mutagen.mp4 import MP4
    from mutagen.id3 import APIC, ID3
except ImportError:
    print("Установите mutagen: pip install mutagen")
    sys.exit(1)

try:
    from urllib3.contrib.socks import SOCKSProxyManager
    SOCKS_AVAILABLE = True
except ImportError:
    SOCKS_AVAILABLE = False
    print("Warning: SOCKS proxy not available. Install: pip install requests[socks]")

AUDIO_EXTENSIONS = {".flac", ".mp3", ".ogg", ".opus", ".m4a"}

# ============================================================================
# КОНФИГУРАЦИЯ
# ============================================================================

# Jikan API (бесплатный, без ключа)
JIKAN_BASE_URL = "https://api.jikan.moe/v4"

# Прокси для доступа к API (SOCKS5)
PROXY_URL = "socks5://192.168.31.216:1080"

# Кэш обложек (чтобы не скачивать одну и ту же много раз)
COVER_CACHE: Dict[str, bytes] = {}

# Задержки для API (Jikan требует ~2 секунды между запросами)
JIKAN_DELAY = 2.0

# ============================================================================
# ФУНКЦИИ
# ============================================================================


def extract_anime_name(filepath: Path) -> Optional[str]:
    """
    Извлекает название аниме из первых квадратных скобок в названии файла.
    """
    stem = filepath.stem
    match = re.search(r'\[([^\]]+)\]', stem)
    if match:
        return match.group(1).strip()
    return None


class JikanAPI:
    """Класс для работы с Jikan API (jikan.moe) через прокси"""
    
    def __init__(self, proxy_url: str = None):
        self.base_url = JIKAN_BASE_URL
        self.session = self._create_session(proxy_url)
        self.last_request_time = 0
    
    def _create_session(self, proxy_url: str = None) -> requests.Session:
        """Создаёт сессию с прокси (если указан)"""
        session = requests.Session()
        
        # Jikan требует User-Agent
        session.headers.update({
            'User-Agent': 'Anime-Cover-Downloader/1.0 (shu@vespahomelab.ru)'
        })
        
        # Настраиваем прокси через SOCKS5
        if proxy_url and SOCKS_AVAILABLE:
            try:
                proxy_manager = SOCKSProxyManager(proxy_url)
                session.mount('http://', proxy_manager)
                session.mount('https://', proxy_manager)
                print(f"  ✓ Прокси настроен: {proxy_url}")
            except Exception as e:
                print(f"  ⚠ Warning: Не удалось настроить прокси: {e}. Будет использовано прямое соединение.")
        
        return session
    
    def _rate_limit(self):
        """Соблюдает лимиты API Jikan (~2 секунды между запросами)"""
        elapsed = time.time() - self.last_request_time
        if elapsed < JIKAN_DELAY:
            time.sleep(JIKAN_DELAY - elapsed)
        self.last_request_time = time.time()
    
    def search_anime(self, query: str, limit: int = 5) -> list:
        """Ищет аниме по названию"""
        self._rate_limit()
        
        url = f"{self.base_url}/anime"
        params = {
            'q': query,
            'limit': limit
        }
        
        try:
            response = self.session.get(url, params=params, timeout=30)
            response.raise_for_status()
            data = response.json()
            
            results = []
            for item in data.get('data', []):
                title = item.get('title', '')
                title_english = item.get('title_english', '')
                images = item.get('images', {})
                jpg = images.get('jpg', {})
                image_url = jpg.get('large', jpg.get('image_url', ''))
                
                if image_url:
                    results.append({
                        'title': title,
                        'title_english': title_english,
                        'image_url': image_url,
                        'source': 'Jikan'
                    })
            
            return results
            
        except Exception as e:
            print(f"  ✗ Ошибка поиска в Jikan: {e}")
            return []
    
    def get_anime_cover(self, anime_name: str) -> Optional[str]:
        """Получает URL обложки для аниме"""
        results = self.search_anime(anime_name)
        
        if results:
            return results[0]['image_url']
        
        return None


def download_image(url: str, proxy_url: str = None) -> Optional[bytes]:
    """Скачивает изображение по URL через прокси"""
    try:
        session = requests.Session()
        
        if proxy_url and SOCKS_AVAILABLE:
            try:
                proxy_manager = SOCKSProxyManager(proxy_url)
                session.mount('http://', proxy_manager)
                session.mount('https://', proxy_manager)
            except:
                pass
        
        response = session.get(url, timeout=30)
        response.raise_for_status()
        return response.content
    except Exception as e:
        print(f"  ✗ Ошибка скачивания изображения: {e}")
        return None


def get_cover_image(anime_name: str, jikan_api: JikanAPI, proxy_url: str = None) -> Optional[bytes]:
    """
    Получает обложку для аниме через Jikan API.
    Использует кэш для повторных запросов.
    """
    # Проверяем кэш
    cache_key = anime_name.lower()
    if cache_key in COVER_CACHE:
        print(f"  ✓ Найдено в кэше: {anime_name}")
        return COVER_CACHE[cache_key]
    
    print(f"  Поиск обложки для: {anime_name}")
    
    # Ищем через Jikan
    image_url = jikan_api.get_anime_cover(anime_name)
    if image_url:
        print(f"  ✓ Найдено: {image_url}")
        image_data = download_image(image_url, proxy_url)
        if image_data:
            COVER_CACHE[cache_key] = image_data
            return image_data
    
    print(f"  ✗ Обложка не найдена")
    return None


def embed_cover(filepath: Path, image_data: bytes, dry_run: bool = False) -> bool:
    """Встраивает обложку в аудиофайл."""
    if dry_run:
        return True
    
    ext = filepath.suffix.lower()
    
    try:
        if ext == ".mp3":
            audio = MP3(str(filepath), ID3=ID3)
            if audio.tags is None:
                audio.add_tags()
            
            # Удаляем старую обложку если есть
            audio.tags.delall('APIC')
            
            # Добавляем новую
            audio.tags.add(APIC(
                encoding=3,
                mime='image/jpeg',
                type=3,  # Front cover
                desc='Cover',
                data=image_data
            ))
            audio.save()
            
        elif ext == ".flac":
            audio = FLAC(str(filepath))
            audio.clear_pictures()
            audio.add_picture(image_data)
            audio.save()
            
        elif ext == ".ogg":
            audio = OggVorbis(str(filepath))
            from mutagen.flac import Picture
            pic = Picture()
            pic.type = 3
            pic.mime = 'image/jpeg'
            pic.data = image_data
            audio.add_picture(pic)
            audio.save()
            
        elif ext == ".opus":
            audio = OggOpus(str(filepath))
            from mutagen.flac import Picture
            pic = Picture()
            pic.type = 3
            pic.mime = 'image/jpeg'
            pic.data = image_data
            audio.add_picture(pic)
            audio.save()
            
        elif ext == ".m4a":
            audio = MP4(str(filepath))
            audio['covr'] = [image_data]
            audio.save()
            
        else:
            return False
        
        return True
        
    except Exception as e:
        print(f"    ✗ Ошибка встраивания обложки: {e}")
        return False


def process_directory(root_path: str, dry_run: bool = False):
    """Обработать директорию."""
    root = Path(root_path)
    
    if not root.is_dir():
        print(f"Директория не найдена: {root_path}")
        sys.exit(1)
    
    # Инициализируем Jikan API с прокси
    jikan_api = JikanAPI(PROXY_URL)
    
    prefix_str = "[DRY] " if dry_run else ""
    processed = 0
    skipped = 0
    errors = 0
    no_cover = 0
    
    audio_files = sorted([
        f for f in root.rglob("*")
        if f.is_file() and f.suffix.lower() in AUDIO_EXTENSIONS
    ])
    
    if not audio_files:
        print("Аудиофайлы не найдены")
        return
    
    print(f"{prefix_str}Найдено файлов: {len(audio_files)}\n")
    
    # Группируем файлы по названию аниме для кэширования обложек
    anime_covers: Dict[str, Optional[bytes]] = {}
    
    for idx, filepath in enumerate(audio_files, 1):
        anime_name = extract_anime_name(filepath)
        
        if not anime_name:
            skipped += 1
            continue
        
        print(f"[{idx}/{len(audio_files)}] {filepath.name}")
        
        # Получаем обложку (с кэшированием)
        if anime_name not in anime_covers:
            cover_data = get_cover_image(anime_name, jikan_api, PROXY_URL)
            anime_covers[anime_name] = cover_data
        else:
            cover_data = anime_covers[anime_name]
            print(f"  ✓ Используем кэш для: {anime_name}")
        
        if cover_data is None:
            no_cover += 1
            continue
        
        success = embed_cover(filepath, cover_data, dry_run)
        
        if success:
            processed += 1
            if dry_run:
                print(f"    ✓ Будет встроена обложка")
            else:
                print(f"    ✓ Обложка встроена")
        else:
            errors += 1
    
    print(f"\n{'=' * 60}")
    print(f"  {'Будет обработано' if dry_run else 'Обработано'}: {processed}")
    print(f"  Пропущено (нет названия аниме): {skipped}")
    print(f"  Не найдено обложек: {no_cover}")
    print(f"  Ошибок: {errors}")
    print(f"{'=' * 60}")


def search_cover(anime_name: str):
    """Тестовый поиск обложки для аниме"""
    jikan_api = JikanAPI(PROXY_URL)
    
    print(f"Поиск обложки для: {anime_name}\n")
    
    print("Jikan (MyAnimeList):")
    results = jikan_api.search_anime(anime_name, limit=5)
    for i, result in enumerate(results, 1):
        print(f"  {i}. {result['title']}")
        if result['title_english']:
            print(f"     English: {result['title_english']}")
        print(f"     {result['image_url']}")
    
    if not results:
        print("  Не найдено")


def main():
    parser = argparse.ArgumentParser(description="Скачивание и встраивание обложек аниме (с прокси)")
    sub = parser.add_subparsers(dest="command")
    
    p_run = sub.add_parser("run", help="Обработать директорию")
    p_run.add_argument("directory", help="Директория с музыкой")
    p_run.add_argument("--dry-run", action="store_true", help="Показать что будет сделано без изменений")
    
    p_search = sub.add_parser("search", help="Поиск обложки для теста")
    p_search.add_argument("anime_name", help="Название аниме для поиска")
    
    args = parser.parse_args()
    
    if args.command == "search":
        search_cover(args.anime_name)
    elif args.command == "run":
        process_directory(args.directory, args.dry_run)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()