#!/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 struct
import sys
import time
from collections import Counter
from datetime import datetime, timedelta
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',
    '.mp4',
}

MAX_FILES = 10000  # ← меняй тут: сколько файлов обрабатывать за один запуск

# 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) ──────────────────

# Tracks the last EXIF field name used (for display in output)
_last_exif_field = ''

_EXIF_FIELD_NAMES = {36867: 'DateTimeOriginal', 36868: 'DateTimeDigitized', 306: 'DateTime', 0: 'MP4:mvhd'}


def _set_last_field(tag_id: int) -> None:
    global _last_exif_field
    _last_exif_field = _EXIF_FIELD_NAMES.get(tag_id, f'Tag#{tag_id}')


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', {})
    for tag_id in (piexif.ExifIFD.DateTimeOriginal, piexif.ExifIFD.DateTimeDigitized, piexif.ImageIFD.DateTime):
        raw = exif.get(tag_id) or image.get(tag_id)
        if raw:
            _set_last_field(tag_id)
            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', 'DateTimeDigitized'):
        tag_code = tag_map.get(tag_name)
        if tag_code and tag_code in raw:
            _set_last_field(tag_code)
            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', 'DateTimeDigitized'):
        tag_code = tag_map.get(tag_name)
        if tag_code and tag_code in raw:
            _set_last_field(tag_code)
            return raw[tag_code]
    return None


def _read_cr2_exif(path: Path) -> str | None:
    """CR2/CR3 (Canon Raw): try piexif, then PIL, then raw bytes scan."""
    # ── Tier 1: piexif (TIFF-based EXIF) ──────────────────────
    if HAS_PIEXIF:
        try:
            exif_dict = piexif.load(str(path))
            # All possible date tag locations in EXIF
            for ifd_name, tag_id in (
                ('Exif', piexif.ExifIFD.DateTimeOriginal),    # 36867
                ('Exif', piexif.ExifIFD.DateTimeDigitized),   # 36868
                ('Image', piexif.ImageIFD.DateTime),           # 306
            ):
                raw_val = exif_dict.get(ifd_name, {}).get(tag_id)
                if raw_val:
                    _set_last_field(tag_id)
                    return raw_val.decode('utf-8')
        except Exception:
            pass

    # ── Tier 2: PIL (embedded JPEG preview) ───────────────────
    try:
        from PIL import Image
        from PIL.ExifTags import TAGS
        img = Image.open(path)
        raw = img._getexif()
        img.close()
        if raw:
            for tag_code in (36867, 36868, 306):  # DateTimeOriginal, Digitized, DateTime
                val = raw.get(tag_code)
                if val:
                    _set_last_field(tag_code)
                    return val
    except Exception:
        pass

    # ── Tier 3: raw bytes scan (last resort) ──────────────────
    try:
        data = path.read_bytes()
        # CR2 stores dates as ASCII 'YYYY:MM:DD HH:MM:SS'
        import re
        m = re.search(rb'(\d{4}:\d{2}:\d{2} \d{2}:\d{2}:\d{2})', data)
        if m:
            return m.group(1).decode('ascii')
    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 _read_mp4_date(path: Path) -> str | None:
    """Read creation date from MP4 (moov > mvhd atom).

    MP4 stores creation time in the Movie Header Box (mvhd)
    as seconds since January 1, 1904 (Mac epoch).  No external
    dependencies needed — pure binary parsing.
    """
    try:
        data = path.read_bytes()
    except Exception:
        return None

    pos = 0
    size = len(data)
    _MAC_EPOCH = datetime(1904, 1, 1)

    try:
        while pos < size - 16:
            box_size = struct.unpack_from('>I', data, pos)[0]
            box_type = data[pos + 4:pos + 8]
            if box_size == 0:
                break
            if box_type == b'moov':
                # Found moov box — scan for mvhd inside
                end = pos + box_size
                inner = pos + 8
                while inner < end - 16:
                    inner_sz = struct.unpack_from('>I', data, inner)[0]
                    inner_ty = data[inner + 4:inner + 8]
                    if inner_sz == 0:
                        break
                    if inner_ty == b'mvhd':
                        ver = data[inner + 8]
                        if ver == 0:
                            ctime = struct.unpack_from('>I', data, inner + 12)[0]
                        else:
                            ctime = struct.unpack_from('>Q', data, inner + 16)[0]
                        dt = _MAC_EPOCH + timedelta(seconds=ctime)
                        return dt.strftime('%Y:%m:%d %H:%M:%S')
                    inner += inner_sz
                break
            pos += box_size
    except Exception:
        pass
    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,
    max_files: int | None = None,
) -> 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 max_files is not None and len(files) > max_files:
        print(_c(f'  First {max_files} of {len(files)} files '
                 f'(edit MAX_FILES in organize.py to change)',
                 _CYAN, no_color=no_color))
        files = files[:max_files]

    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:
        _last_exif_field = ''
        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', '.mp4'):
            date_str = _read_other_exif_pil(filepath)
            if date_str:
                dt = _parse_date(date_str)
                if dt:
                    from_exif = True

        # ── EXIF tier 6: MP4 (moov > mvhd binary parse) ──────────────
        if dt is None and ext == '.mp4':
            date_str = _read_mp4_date(filepath)
            if date_str:
                dt = _parse_date(date_str)
                if dt:
                    _set_last_field(0)  # mark as MP4
                    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
            date_source_label = _last_exif_field or 'EXIF'
        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('chcp 65001 >nul\r\n')
        f.write('title Raskidator-3000 — Undo move\r\n')
        f.write('setlocal enabledelayedexpansion\r\n')
        f.write('echo Restoring files...\r\n')
        f.write('echo.\r\n')
        f.write('set FAILED=0\r\n')
        # Reverse: move destination → source
        for entry in reversed(entries):
            src = entry['source']
            dst = entry['destination']
            # Normalise to backslashes for Windows
            src = src.replace('/', '\\')
            dst = dst.replace('/', '\\')
            fname = dst.rsplit('\\', 1)[-1] if '\\' in dst else dst

            src_drive = src[:2] if len(src) >= 2 and src[1] == ':' else ''
            dst_drive = dst[:2] if len(dst) >= 2 and dst[1] == ':' else ''

            if src_drive and dst_drive and src_drive.upper() != dst_drive.upper():
                # Cross-drive — robocopy /move
                src_dir = src.rsplit('\\', 1)[0] if '\\' in src else src
                dst_dir = dst.rsplit('\\', 1)[0] if '\\' in dst else dst
                f.write(f'robocopy "{dst_dir}" "{src_dir}" "{fname}" /move /ndl /njh /njs /nfl >nul\r\n')
                f.write(f'if errorlevel 8 echo FAILED (cross-drive: {src_drive} → {dst_drive})\r\n')
                f.write(f'if errorlevel 8 set /a FAILED+=1\r\n')
            else:
                f.write(f'move "{dst}" "{src}"\r\n')
                f.write(f'if errorlevel 1 echo FAILED: {fname}\r\n')
                f.write(f'if errorlevel 1 set /a FAILED+=1\r\n')
        f.write('echo.\r\n')
        f.write(f'echo Restored {len(entries)} files. Failures: !FAILED!\r\n')
        f.write('if !FAILED! gtr 0 (\r\n')
        f.write('  echo.\r\n')
        f.write('  echo Some files could not be restored.^ See above for details.\r\n')
        f.write(')\r\n')
        f.write('echo.\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',
    )
    parser.add_argument(
        '--max-files',
        type=int,
        default=None,
        help=f'Max files to process per run (default: {MAX_FILES}, but use MAX_FILES constant at top of organize.py)',
    )

    args = parser.parse_args()
    # If --max-files not set, use MAX_FILES constant
    max_files = args.max_files if args.max_files is not None else MAX_FILES

    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))
    if max_files:
        print(f'  Max files: {max_files}')
    print()

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


if __name__ == '__main__':
    main()