import os
import sys
from datetime import datetime
from pathlib import Path
from PIL import Image, TiffImagePlugin
import piexif
from concurrent.futures import ThreadPoolExecutor, as_completed

def get_arbitrary_date():
    """Prompts the user for a date. Returns Unix epoch if blank, or the parsed target date."""
    while True:
        user_input = input("Enter the target date and time (YYYY-MM-DD HH:MM:SS) [Press Enter for 1970-01-01 00:00:00]: ").strip()
        
        if not user_input:
            return "1970:01:01 00:00:00"
            
        try:
            # Parse user's hyphenated input
            dt = datetime.strptime(user_input, "%Y-%m-%d %H:%M:%S")
            # Return colon-separated EXIF format
            return dt.strftime("%Y:%m:%d %H:%M:%S")
        except ValueError:
            print("Incorrect format. You must use YYYY-MM-DD HH:MM:SS or press Enter for 1970. Try again.", file=sys.stderr)

def update_file_timestamps(file_path, exif_date_str):
    """Updates EXIF metadata/native chunks based on format, and sets filesystem mtime/atime."""
    try:
        # Match the colon format returned by get_arbitrary_date
        dt = datetime.strptime(exif_date_str, "%Y:%m:%d %H:%M:%S")
        target_timestamp = dt.timestamp()
    except ValueError as e:
        return f"Failed to parse internal date string {exif_date_str}: {e}"
        
    file_suffix = file_path.suffix.lower()
    meta_updated = False

    # 1. Handle JPEG/WEBP via piexif
    if file_suffix in {'.jpg', '.jpeg', '.webp'}:
        try:
            try:
                exif_dict = piexif.load(str(file_path))
            except Exception:
                exif_dict = {"0th": {}, "Exif": {}, "GPS": {}, "Interop": {}, "1st": {}, "thumbnail": None}

            date_bytes = exif_date_str.encode('utf-8')
            exif_dict["0th"][piexif.ImageIFD.DateTime] = date_bytes
            exif_dict["Exif"][piexif.ExifIFD.DateTimeOriginal] = date_bytes
            exif_dict["Exif"][piexif.ExifIFD.DateTimeDigitized] = date_bytes

            exif_bytes = piexif.dump(exif_dict)
            piexif.insert(exif_bytes, str(file_path))
            meta_updated = True
        except Exception as e:
            return f"Warning: Failed to update EXIF for JPEG/WEBP {file_path.name}: {e}"

    # 2. Handle PNG via Pillow
    elif file_suffix == '.png':
        try:
            with Image.open(file_path) as img:
                exif_dict = {"0th": {}, "Exif": {}, "GPS": {}, "Interop": {}, "1st": {}, "thumbnail": None}
                date_bytes = exif_date_str.encode('utf-8')
                exif_dict["0th"][piexif.ImageIFD.DateTime] = date_bytes
                exif_dict["Exif"][piexif.ExifIFD.DateTimeOriginal] = date_bytes
                exif_dict["Exif"][piexif.ExifIFD.DateTimeDigitized] = date_bytes
                
                exif_bytes = piexif.dump(exif_dict)
                img.save(file_path, exif=exif_bytes)
                meta_updated = True
        except Exception as e:
            return f"Warning: Failed to inject EXIF chunk into PNG {file_path.name}: {e}"

    # 3. Handle TIFF via structured TiffImagePlugin tags
    elif file_suffix in {'.tiff', '.tif'}:
        try:
            with Image.open(file_path) as img:
                tiff_info = TiffImagePlugin.ImageFileDirectoryV2()
                tiff_info[306] = exif_date_str
                tiff_info[36867] = exif_date_str
                tiff_info[36868] = exif_date_str
                
                img.save(file_path, tiffinfo=tiff_info)
                meta_updated = True
        except Exception as e:
            return f"Warning: Failed to update TIFF tags for {file_path.name}: {e}"

    # 4. Apply filesystem updates (mtime and atime) to everything
    try:
        os.utime(str(file_path), (target_timestamp, target_timestamp))
        status = "META + FS" if meta_updated else "FS Only"
        return f"Updated ({status}): {file_path.name}"
    except Exception as e:
        return f"Failed to update filesystem time for {file_path.name}: {e}"

def main():
    try:
        current_dir = Path.cwd()
    except FileNotFoundError:
        pwd_env = os.environ.get('PWD')
        if pwd_env and os.path.exists(pwd_env):
            current_dir = Path(pwd_env)
        else:
            print("Error: The current directory does not exist. 'cd' into a valid directory and try again.", file=sys.stderr)
            sys.exit(1)

    valid_extensions = {
        '.jpg', '.jpeg', '.png', '.webp', '.tiff', '.tif', 
        '.bmp', '.gif', '.heic', '.heif', '.cr2', '.nef'
    }
    
    try:
        image_files = [
            f for f in current_dir.rglob('*') 
            if f.is_file() and f.suffix.lower() in valid_extensions
        ]
    except FileNotFoundError:
        print(f"Error: Directory path '{current_dir}' is unreachable.", file=sys.stderr)
        sys.exit(1)

    if not image_files:
        print("No matching images found in this directory or subdirectories.")
        return

    print(f"Found {len(image_files)} images recursively under {current_dir}")
    target_date = get_arbitrary_date()

    print(f"\nProcessing files concurrently using target date: {target_date}...")
    
    max_workers = min(32, (os.cpu_count() or 1) * 5)
    
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_img = {executor.submit(update_file_timestamps, img_path, target_date): img_path for img_path in image_files}
        
        for future in as_completed(future_to_img):
            result_message = future.result()
            if result_message:
                print(result_message)
        
    print("\nFinished processing all files.")

if __name__ == "__main__":
    main()