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

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

Поддерживаемые форматы:
  JPEG, PNG, HEIC, TIFF, WebP, BMP, CR2 (Canon Raw), AVI

Использование:
    python organize.py [путь] [--no-color] [--dry-run] [--quiet] [--group-by month]
    python organize.py --help
"""

from __future__ import annotations

import argparse
import csv
import os
import shutil
import sys
import time
from collections import Counter
from datetime import datetime
from pathlib import Path
from typing import Any

# ── Optional imports ────────────────────────────────────────────────────────

HAS_TQDM = False
HAS_PIEXIF = False
HAS_RAWPY = False
HAS_PILLOW_HEIF = False

try:
    from tqdm import tqdm as _tqdm
    HAS_TQDM = True
except ImportError:
    pass

try:
    import piexif
    HAS_PIEXIF = True
except ImportError:
    pass

try:
    import rawpy
    HAS_RAWPY = True
except ImportError:
    pass

try:
    from pillow_heif import open_heif
    HAS_PILLOW_HEIF = True
except ImportError:
    pass

# ── Constants ───────────────────────────────────────────────────────────────

EXTENSIONS = {
    '.jpg', '.jpeg', '.png', '.heic', '.tiff', '.tif', '.webp', '.bmp',
    '.cr2', '.cr3', '.nef', '.arw', '.dng',
    '.avi',
}

# ANSI colour codes
_GREEN = '\033[32m'
_YELLOW = '\033[33m'
_RED = '\033[31m'
_CYAN = '\033[36m'
_BOLD = '\033[1m'
_RESET = '\033[0m'


def _c(text: str, colour: str, *, no_color: bool = False) -> str:
    """Colourise *text* unless *no_color* is True."""
    if no_color:
        return text
    return f'{colour}{text}{_RESET}'


# ── EXIF extraction (JPEG — piexif, 5-10x faster than PIL) ──────────────────

def _read_jpeg_exif_piexif(path: Path) -> str | None:
    """Read DateTimeOriginal from a JPEG using piexif (no PIL needed).

    Works on raw bytes — much faster than opening the image.
    """
    if not HAS_PIEXIF:
        return None
    try:
        exif_dict = piexif.load(str(path))
    except Exception:
        return None

    exif = exif_dict.get('Exif', {})
    image = exif_dict.get('Image', {})
    raw = (exif.get(piexif.ExifIFD.DateTimeOriginal) or
           image.get(piexif.ImageIFD.DateTime))
    if raw:
        try:
            return raw.decode('utf-8')
        except Exception:
            return None
    return None


def _read_jpeg_exif_pil(path: Path) -> str | None:
    """Fallback: PIL-based EXIF read for JPEG (slower but more compatible)."""
    try:
        from PIL import Image
        from PIL.ExifTags import TAGS
    except ImportError:
        return None

    try:
        img = Image.open(path)
        raw = img._getexif()
        img.close()
    except Exception:
        return None
    if not raw:
        return None

    tag_map = {v: k for k, v in TAGS.items()}
    for tag_name in ('DateTimeOriginal', 'DateTime'):
        tag_code = tag_map.get(tag_name)
        if tag_code and tag_code in raw:
            return raw[tag_code]
    return None


def _read_heic_exif(path: Path) -> str | None:
    """HEIC: needs pillow-heif + piexif."""
    if not HAS_PILLOW_HEIF or not HAS_PIEXIF:
        return None
    try:
        heif = open_heif(str(path))
        raw_bytes = heif.info.get('exif', b'')
        if not raw_bytes:
            return None
        exif_dict = piexif.load(raw_bytes)
        raw = exif_dict.get('Exif', {}).get(piexif.ExifIFD.DateTimeOriginal)
        if raw:
            return raw.decode('utf-8')
    except Exception:
        pass
    return None


def _read_other_exif_pil(path: Path) -> str | None:
    """Generic PIL-based EXIF read for TIFF, WebP, BMP."""
    try:
        from PIL import Image
        from PIL.ExifTags import TAGS
    except ImportError:
        return None

    try:
        img = Image.open(path)
        raw = img._getexif()
        img.close()
    except Exception:
        return None
    if not raw:
        return None

    tag_map = {v: k for k, v in TAGS.items()}
    for tag_name in ('DateTimeOriginal', 'DateTime'):
        tag_code = tag_map.get(tag_name)
        if tag_code and tag_code in raw:
            return raw[tag_code]
    return None


def _read_cr2_exif(path: Path) -> str | None:
    """CR2/CR3 (Canon Raw): try piexif first, then PIL embedded JPEG."""
    # Tier 1: piexif on CR2 (TIFF-based EXIF)
    if HAS_PIEXIF:
        try:
            exif_dict = piexif.load(str(path))
            raw_val = exif_dict.get('Exif', {}).get(piexif.ExifIFD.DateTimeOriginal)
            if raw_val:
                return raw_val.decode('utf-8')
        except Exception:
            pass

    # Tier 2: PIL can read EXIF from CR2's embedded JPEG preview
    try:
        from PIL import Image
        from PIL.ExifTags import TAGS
        img = Image.open(path)
        raw = img._getexif()
        img.close()
        if raw:
            tag_map = {v: k for k, v in TAGS.items()}
            for tag_name in ('DateTimeOriginal', 'DateTime'):
                tag_code = tag_map.get(tag_name)
                if tag_code and tag_code in raw:
                    return raw[tag_code]
    except Exception:
        pass
    return None


# ── Date helpers ────────────────────────────────────────────────────────────

def _parse_date(date_str: str) -> datetime | None:
    """Parse 'YYYY:MM:DD HH:MM:SS' to datetime."""
    try:
        return datetime.strptime(date_str.strip(), '%Y:%m:%d %H:%M:%S')
    except (ValueError, AttributeError):
        return None


def _get_created_date(path: Path) -> datetime | None:
    """Fallback: file system creation date (st_ctime on Windows)."""
    try:
        ts = path.stat().st_ctime
        return datetime.fromtimestamp(ts)
    except Exception:
        return None


# ── Core ─────────────────────────────────────────────────────────────────────

def organize(
    folder: Path,
    *,
    dry_run: bool = False,
    no_color: bool = False,
    group_by: str = 'day',
    quiet: bool = False,
) -> int:
    """Scan *folder* and sort images into year/month[/day] subfolders.

    Returns exit code (0 = ok, 1 = errors).
    """
    if not folder.is_dir():
        print(f'Error: folder not found — {folder}')
        return 1

    # ── Scan ────────────────────────────────────────────────────────────────
    files: list[Path] = [
        f for f in folder.iterdir()
        if f.is_file() and f.suffix.lower() in EXTENSIONS
    ]
    files.sort()

    if not files:
        print(_c('  No supported files found.', _YELLOW, no_color=no_color))
        return 0

    total = len(files)

    # ── Progress bar ────────────────────────────────────────────────────────
    use_tqdm = HAS_TQDM and not quiet and not dry_run
    iterator: Any = files
    if use_tqdm:
        iterator = _tqdm(files, desc='Processing', unit='file',
                         bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} '
                                     '[{elapsed}<{remaining}, {rate_fmt}]')

    # ── State ───────────────────────────────────────────────────────────────
    moved = 0
    skipped = 0
    no_date_count = 0
    exif_count = 0
    fs_count = 0
    year_counts: Counter[int] = Counter()
    move_log: list[dict[str, str]] = []

    start_time = time.time()

    for filepath in iterator:
        dt: datetime | None = None
        from_exif = False
        date_source_label: str | None = None

        ext = filepath.suffix.lower()

        # ── EXIF tier 1: JPEG — piexif (fastest) ────────────────────────
        if ext in ('.jpg', '.jpeg'):
            date_str = _read_jpeg_exif_piexif(filepath)
            if date_str:
                dt = _parse_date(date_str)
                if dt:
                    from_exif = True

        # ── EXIF tier 2: HEIC (pillow-heif + piexif) ────────────────────
        if dt is None and ext == '.heic':
            date_str = _read_heic_exif(filepath)
            if date_str:
                dt = _parse_date(date_str)
                if dt:
                    from_exif = True

        # ── EXIF tier 3: CR2/raw (piexif on JPEG preview) ──────────────
        if dt is None and ext in ('.cr2', '.cr3'):
            date_str = _read_cr2_exif(filepath)
            if date_str:
                dt = _parse_date(date_str)
                if dt:
                    from_exif = True

        # ── EXIF tier 4: JPEG fallback (PIL) ───────────────────────────
        if dt is None and ext in ('.jpg', '.jpeg'):
            date_str = _read_jpeg_exif_pil(filepath)
            if date_str:
                dt = _parse_date(date_str)
                if dt:
                    from_exif = True

        # ── EXIF tier 5: other formats (TIFF, WebP, BMP, PNG) ──────────
        if dt is None and ext not in ('.avi',):
            date_str = _read_other_exif_pil(filepath)
            if date_str:
                dt = _parse_date(date_str)
                if dt:
                    from_exif = True

        # ── Fallback: file creation date (always available) ────────────
        if dt is None:
            dt = _get_created_date(filepath)
            if dt:
                date_source_label = 'Created Date'

        # ── No date → skip ─────────────────────────────────────────────
        if dt is None:
            msg = (
                f'  [{_c("—", _RED, no_color=no_color)}] '
                f'{filepath.name} — no date, skipped'
            )
            if not quiet:
                print(msg)
            no_date_count += 1
            skipped += 1
            continue

        # ── Track stats ────────────────────────────────────────────────
        if from_exif:
            exif_count += 1
        else:
            fs_count += 1
        year_counts[dt.year] += 1

        # ── Build destination ──────────────────────────────────────────
        if group_by == 'month':
            dest = folder / str(dt.year) / f'{dt.month:02d}'
        else:
            dest = folder / str(dt.year) / f'{dt.month:02d}' / f'{dt.day:02d}'

        dest_file = dest / filepath.name

        # Handle filename collisions
        if dest_file.exists():
            stem = filepath.stem
            suffix = filepath.suffix
            n = 1
            while True:
                dest_file = dest / f'{stem}_{n}{suffix}'
                if not dest_file.exists():
                    break
                n += 1

        # ── Execute move or dry-run ────────────────────────────────────
        label = f'{dt.year}/{dt.month:02d}/{dt.day:02d}' if group_by == 'day' \
                else f'{dt.year}/{dt.month:02d}'

        if dry_run:
            src_tag = f' ({date_source_label})' if date_source_label else ''
            move_log.append({
                'source': str(filepath),
                'destination': str(dest_file),
                'date': str(dt),
                'date_source': 'EXIF' if from_exif else (date_source_label or '?'),
            })
            if not quiet and not use_tqdm:
                marker = _c('D', _CYAN, no_color=no_color)
                icon = _c('✓', _GREEN, no_color=no_color) if from_exif \
                    else _c('✓', _YELLOW, no_color=no_color)
                src_tag = f' ({date_source_label})' if date_source_label else ''
                msg = f'  [{marker}] {filepath.name} → {label}{src_tag}'
                print(msg)
            moved += 1
            continue

        # Real move
        dest.mkdir(parents=True, exist_ok=True)
        try:
            shutil.move(str(filepath), str(dest_file))
        except OSError as e:
            msg = (
                f'  [{_c("✗", _RED, no_color=no_color)}] '
                f'{filepath.name} — move failed: {e}'
            )
            if not quiet:
                print(msg)
            skipped += 1
            continue

        move_log.append({
            'source': str(filepath),
            'destination': str(dest_file),
            'date': str(dt),
            'date_source': 'EXIF' if from_exif else (date_source_label or '?'),
        })

        if not quiet and not use_tqdm:
            src_tag = f' ({date_source_label})' if date_source_label else ''
            colour = _GREEN if from_exif else _YELLOW
            msg = (
                f'  [{_c("✓", colour, no_color=no_color)}] '
                f'{filepath.name} → {label}{src_tag}'
            )
            print(msg)

        moved += 1

        if not dry_run and not quiet and use_tqdm and isinstance(iterator, _tqdm):
            iterator.set_postfix(moved=moved, skipped=skipped)

    # ── Summary ─────────────────────────────────────────────────────────
    elapsed = time.time() - start_time
    print()
    print(_c('=== Year Summary ===', _BOLD, no_color=no_color))
    if year_counts:
        for year in sorted(year_counts):
            print(f'  {year}: {year_counts[year]} files')
    else:
        print('  (no entries with a determinable year)')
    print()

    mode_label = 'DRY-RUN (no files moved)' if dry_run else ''
    if mode_label:
        print(_c(mode_label, _CYAN, no_color=no_color))

    print(f'  EXIF date:              {exif_count}')
    print(f'  Created Date:           {fs_count}')
    print(f'  Skipped (no date):      {skipped}')
    print(f'  Total processed:        {moved + skipped}')
    print(f'  Time:                   {elapsed:.1f}s')
    print()

    if dry_run:
        print(_c(f'Dry-run: {moved} would be moved, {skipped} skipped (no date).',
                  _CYAN, no_color=no_color))
    else:
        print(f'Done! Moved: {moved}, skipped: {skipped} (no date: {no_date_count})')

    # ── Write move log ──────────────────────────────────────────────────
    if move_log and not dry_run:
        log_path = folder / '.raskidator-log.csv'
        undo_path = folder / 'raskidator-undo.bat'
        _write_logs(move_log, log_path, undo_path)
        print()
        print(_c(f'  Log:  {log_path}', _CYAN, no_color=no_color))
        print(_c(f'  Undo: {undo_path} (run to restore files)',
                 _CYAN, no_color=no_color))

    return 0


# ── Logging helpers ─────────────────────────────────────────────────────────

def _write_logs(entries: list[dict[str, str]],
                log_path: Path,
                undo_path: Path) -> None:
    """Write CSV move log + undo.bat."""
    fieldnames = ['source', 'destination', 'date', 'date_source']
    with log_path.open('w', newline='', encoding='utf-8') as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(entries)

    # Undo .bat — restores files in reverse order of columns
    with undo_path.open('w', encoding='utf-8') as f:
        f.write('@echo off\r\n')
        f.write('title Raskidator-3000 — Undo move\r\n')
        f.write('echo Restoring files...\r\n')
        f.write('echo.\r\n')
        # Reverse: move destination → source
        for entry in reversed(entries):
            src = entry['source']
            dst = entry['destination']
            f.write(f'move "{dst}" "{src}" >nul 2>&1\r\n')
            if not f.write:  # impossible but keep mypy happy
                pass
        f.write('echo.\r\n')
        f.write(f'echo Restored {len(entries)} files.\r\n')
        f.write('pause\r\n')


# ── CLI entry point ─────────────────────────────────────────────────────────

def main() -> None:
    parser = argparse.ArgumentParser(
        prog='organize.py',
        description='Raskidator-3000 — sort images by EXIF date.',
        add_help=True,
    )
    parser.add_argument(
        'path',
        nargs='?',
        default=None,
        help='Folder to scan (default: current directory)',
    )
    parser.add_argument(
        '--no-color',
        action='store_true',
        default=False,
        help='Disable ANSI color output (for pipes)',
    )
    parser.add_argument(
        '--dry-run',
        action='store_true',
        default=False,
        help='Show what would be moved without actually moving',
    )
    parser.add_argument(
        '--quiet', '-q',
        action='store_true',
        default=False,
        help='Suppress per-file output (shows summary only)',
    )
    parser.add_argument(
        '--group-by',
        choices=('day', 'month'),
        default='day',
        help='Output folder granularity: year/month/day (default) or year/month',
    )

    args = parser.parse_args()

    folder = Path(args.path) if args.path else Path.cwd()
    print(_c('Raskidator-3000', _BOLD, no_color=args.no_color))
    print(f'  Folder: {folder.resolve()}')
    if args.dry_run:
        print(_c('  Mode:   DRY-RUN', _CYAN, no_color=args.no_color))
    print()

    rc = organize(
        folder,
        dry_run=args.dry_run,
        no_color=args.no_color,
        group_by=args.group_by,
        quiet=args.quiet,
    )
    sys.exit(rc)


if __name__ == '__main__':
    main()