#!/usr/bin/env python3
"""
Raskidator-3000 — сортировка фотографий по EXIF дате.

Сканирует указанную папку (по умолчанию текущую) на изображения,
читает EXIF дату съёмки (DateTimeOriginal) и перемещает файлы
в подпапки Год/Месяц/День.

Файлы без EXIF даты остаются на месте.

Поддерживаемые форматы: .jpg, .jpeg, .png, .heic, .tiff, .tif, .webp, .bmp

Использование:
    python organize.py [путь_к_папке]
    python organize.py --no-color [путь_к_папке]
    python organize.py [путь_к_папке] --no-color
"""

import argparse
import os
import sys
import shutil
from collections import Counter
from datetime import datetime
from pathlib import Path

EXTENSIONS = {'.jpg', '.jpeg', '.png', '.heic', '.tiff', '.tif', '.webp', '.bmp'}

# ANSI color codes
_GREEN = '\033[32m'
_YELLOW = '\033[33m'
_RED = '\033[31m'
_RESET = '\033[0m'


def _colorize(text: str, color: str, no_color: bool) -> str:
    """Wrap *text* in ANSI *color*, unless *no_color* is True."""
    if no_color:
        return text
    return f'{color}{text}{_RESET}'


def get_file_created_date(filepath: Path) -> datetime | None:
    """Возвращает дату создания файла из файловой системы (Windows: st_ctime)."""
    try:
        ts = filepath.stat().st_ctime
        return datetime.fromtimestamp(ts)
    except Exception:
        return None


def get_exif_date(filepath: Path) -> str | None:
    """Извлекает дату съёмки из EXIF (DateTimeOriginal)."""
    try:
        # HEIC требует pillow-heif
        if filepath.suffix.lower() == '.heic':
            try:
                from pillow_heif import open_heif
                heif_file = open_heif(str(filepath))
                exif_data = heif_file.info.get('exif', b'')
                if not exif_data:
                    return None
                import piexif
                exif_dict = piexif.load(exif_data)
                if piexif.ExifIFD.DateTimeOriginal in exif_dict.get('Exif', {}):
                    return exif_dict['Exif'][piexif.ExifIFD.DateTimeOriginal].decode()
            except Exception:
                return None
        else:
            from PIL import Image
            from PIL.ExifTags import TAGS

            img = Image.open(filepath)
            exif_data = img._getexif()
            if not exif_data:
                return None

            tag_map = {v: k for k, v in TAGS.items()}
            datetime_original_tag = tag_map.get('DateTimeOriginal')
            if datetime_original_tag and datetime_original_tag in exif_data:
                return exif_data[datetime_original_tag]

            # Fallback: DateTime (если DateTimeOriginal нет)
            datetime_tag = tag_map.get('DateTime')
            if datetime_tag and datetime_tag in exif_data:
                return exif_data[datetime_tag]

    except Exception:
        return None

    return None


def parse_date(date_str: str) -> datetime | None:
    """Парсит EXIF дату в формате 'YYYY:MM:DD HH:MM:SS'."""
    try:
        return datetime.strptime(date_str.strip(), '%Y:%m:%d %H:%M:%S')
    except (ValueError, AttributeError):
        return None


def organize(folder: Path, no_color: bool = False) -> None:
    """Сканирует папку и сортирует изображения."""
    if not folder.exists() or not folder.is_dir():
        print(f'Ошибка: папка не существует — {folder}')
        sys.exit(1)

    files = [f for f in folder.iterdir() if f.is_file() and f.suffix.lower() in EXTENSIONS]
    files.sort()

    moved = 0
    skipped = 0
    no_date = 0
    exif_count = 0
    fs_count = 0
    year_counts: Counter[int] = Counter()

    for filepath in files:
        dt = None
        date_source = None
        from_exif = False

        # 1. EXIF DateTimeOriginal
        date_str = get_exif_date(filepath)
        if date_str:
            dt = parse_date(date_str)
            if dt:
                from_exif = True

        # 2. Fallback: дата создания файла (Windows: st_ctime)
        if dt is None:
            dt = get_file_created_date(filepath)
            if dt:
                date_source = 'дата создания файла'

        if dt is None:
            skipped += 1
            no_date += 1
            msg = f'  [—] {filepath.name} — нет даты (EXIF + CreatedDate), оставлен на месте'
            print(_colorize(msg, _RED, no_color))
            continue

        # Track source and year for final statistics
        if from_exif:
            exif_count += 1
        else:
            fs_count += 1
        year_counts[dt.year] += 1

        dest_dir = folder / str(dt.year) / f'{dt.month:02d}' / f'{dt.day:02d}'
        dest_dir.mkdir(parents=True, exist_ok=True)

        dest_path = dest_dir / filepath.name

        # Если файл с таким именем уже есть — добавляем суффикс
        if dest_path.exists():
            stem = filepath.stem
            suffix = filepath.suffix
            counter = 1
            while True:
                dest_path = dest_dir / f'{stem}_{counter}{suffix}'
                if not dest_path.exists():
                    break
                counter += 1

        shutil.move(str(filepath), str(dest_path))
        moved += 1
        src = f' ({date_source})' if date_source else ''
        msg = f'  [✓] {filepath.name} → {dt.year}/{dt.month:02d}/{dt.day:02d}/{src}'

        if from_exif:
            # Green — valid EXIF date
            print(_colorize(msg, _GREEN, no_color))
        else:
            # Yellow — fallback to file system creation date
            print(_colorize(msg, _YELLOW, no_color))

    print()
    print('=== Year Summary ===')
    if year_counts:
        for year in sorted(year_counts):
            print(f'  {year}: {year_counts[year]} entries')
    else:
        print('  (no entries with a determinable year)')
    print()
    print(f'  Total with EXIF date:     {exif_count}')
    print(f'  Total with FS date:       {fs_count}')
    print(f'  Total skipped:            {skipped}')
    print(f'  Total processed:          {moved + skipped}')
    print()
    print(f'Готово! Перемещено: {moved}, пропущено: {skipped} (без даты: {no_date})')


def main() -> None:
    parser = argparse.ArgumentParser(
        prog='organize.py',
        description='Raskidator-3000 — сортировка фотографий по EXIF дате.',
        add_help=True,
    )
    parser.add_argument(
        'path',
        nargs='?',
        default=None,
        help='Путь к папке с фотографиями (по умолчанию текущая папка)',
    )
    parser.add_argument(
        '--no-color',
        action='store_true',
        default=False,
        help='Отключить цветной вывод ANSI (для совместимости с пайпами)',
    )

    args = parser.parse_args()

    folder = Path(args.path) if args.path else Path.cwd()
    no_color = args.no_color

    print(f'Raskidator-3000 — сортировка фотографий')
    print(f'Папка: {folder.resolve()}')
    print()
    organize(folder, no_color=no_color)


if __name__ == '__main__':
    main()